Add unit tests
This commit is contained in:
@@ -1,5 +1,6 @@
|
||||
mod messages;
|
||||
mod machine;
|
||||
mod noise;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
111
src/machine.rs
111
src/machine.rs
@@ -2,6 +2,18 @@ use x25519_dalek::PublicKey;
|
||||
use x25519_dalek::StaticSecret;
|
||||
use x25519_dalek::SharedSecret;
|
||||
|
||||
use std::sync::Mutex;
|
||||
use std::sync::mpsc::channel;
|
||||
use std::collections::HashMap;
|
||||
|
||||
pub struct Peer {
|
||||
m : StateMutable,
|
||||
sk : StaticSecret,
|
||||
pk : PublicKey,
|
||||
ss : SharedSecret,
|
||||
psk : [u8; 32]
|
||||
}
|
||||
|
||||
/* Mutable part of handshake state */
|
||||
enum StateMutable {
|
||||
Reset,
|
||||
@@ -12,43 +24,92 @@ enum StateMutable {
|
||||
|
||||
/* Immutable part of the handshake state */
|
||||
struct StateFixed {
|
||||
sk : StaticSecret,
|
||||
pk : PublicKey,
|
||||
ss : SharedSecret
|
||||
}
|
||||
|
||||
struct State {
|
||||
m : StateMutable,
|
||||
f : StateFixed,
|
||||
struct StateMachine {
|
||||
peers : Vec<Mutex<Peer>>, // peer index -> state
|
||||
pkmap : HashMap<[u8; 32], usize>, // public key -> peer index
|
||||
ids : Mutex<HashMap<u32, usize>> // receive ids -> peer index
|
||||
}
|
||||
|
||||
struct Key {
|
||||
key : [u8; 32],
|
||||
id : u32
|
||||
}
|
||||
|
||||
struct KeyPair {
|
||||
send : [u8; 32],
|
||||
recv : [u8; 32]
|
||||
confimed : bool, // has the key-pair been confirmed
|
||||
send : Key, // key for outbound messages
|
||||
recv : Key // key for inbound messages
|
||||
}
|
||||
|
||||
impl State {
|
||||
/* Initialize a new handshake state machine
|
||||
*/
|
||||
fn new(sk : StaticSecret, pk : PublicKey) -> State {
|
||||
let ss = sk.diffie_hellman(&pk);
|
||||
State {
|
||||
m : StateMutable::Reset,
|
||||
f : StateFixed{sk, pk, ss}
|
||||
struct Output (
|
||||
Option<KeyPair>, // resulting key-pair of successful handshake
|
||||
Option<u32> // id to be released
|
||||
);
|
||||
|
||||
impl StateMachine {
|
||||
/// Initialize a new handshake state machine
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `sk` - x25519 scalar representing the local private key
|
||||
pub fn new(sk : StaticSecret) -> StateMachine {
|
||||
StateMachine {
|
||||
peers : vec![],
|
||||
pkmap : HashMap::new(),
|
||||
ids : Mutex::new(HashMap::new())
|
||||
}
|
||||
}
|
||||
|
||||
/* Begin a new handshake, returns the initial handshake message
|
||||
*/
|
||||
fn begin(&self) -> Vec<u8> {
|
||||
vec![]
|
||||
/// Add a new public key to the state machine
|
||||
/// To remove public keys, you must create a new machine instance
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `pk` - The public key to add
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// The call might fail if the public key corresponds to the secret key of the machine
|
||||
pub fn add(&mut self, pk : PublicKey) -> Result<(), ()> {
|
||||
// let ss = sk.diffie_hellman(&pk);
|
||||
Err(())
|
||||
}
|
||||
|
||||
/* Process a handshake message.
|
||||
*
|
||||
* Result is either a new state (and optionally a new key pair) or an error
|
||||
*/
|
||||
fn process(&self, msg : &[u8]) -> Result<(State, Option<KeyPair>), ()> {
|
||||
/// Release an id back to the pool
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `id` - The (sender) id to release
|
||||
pub fn release(&self, id : u32) {
|
||||
self.ids.lock().unwrap().remove(&id);
|
||||
}
|
||||
|
||||
/// Begin a new handshake
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `pk` - Public key of peer to initiate handshake for
|
||||
pub fn begin(&self, pk : PublicKey) -> Result<Output, ()> {
|
||||
match self.pkmap.get(pk.as_bytes()) {
|
||||
None => Err(()),
|
||||
Some(&idx) => {
|
||||
let mut peer = self.peers.get(idx).unwrap().lock().unwrap();
|
||||
Err(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Process a handshake message.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `msg` - Byte slice containing the message (untrusted input)
|
||||
fn process(&self, msg : &[u8]) -> Result<Output, ()> {
|
||||
// inspect type field
|
||||
match msg.get(0) {
|
||||
_ => Err(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
122
src/messages.rs
122
src/messages.rs
@@ -1,3 +1,4 @@
|
||||
use std::fmt;
|
||||
use std::mem;
|
||||
|
||||
const SIZE_TAG : usize = 16;
|
||||
@@ -16,7 +17,8 @@ pub const TYPE_RESPONSE : u8 = 2;
|
||||
#[repr(C)]
|
||||
#[derive(Copy, Clone)]
|
||||
struct MessageInitiate {
|
||||
f_type : u32,
|
||||
f_type : u8,
|
||||
f_reserved : [u8; 3],
|
||||
f_sender : u32,
|
||||
f_ephemeral : [u8; SIZE_X25519_POINT],
|
||||
f_static : [u8; SIZE_X25519_POINT + SIZE_TAG],
|
||||
@@ -59,13 +61,39 @@ impl Into<Vec<u8>> for MessageInitiate {
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for MessageInitiate {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f,
|
||||
"MessageInitiate {{ type = {} }}",
|
||||
self.f_type
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
impl PartialEq for MessageInitiate {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.f_type == other.f_type &&
|
||||
self.f_reserved == other.f_reserved &&
|
||||
self.f_sender == other.f_sender &&
|
||||
self.f_ephemeral[..] == other.f_ephemeral[..] &&
|
||||
self.f_static[..] == other.f_static[..] &&
|
||||
self.f_timestamp[..] == other.f_timestamp
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
impl Eq for MessageInitiate {}
|
||||
|
||||
|
||||
/* Wireguard handshake responder message
|
||||
* responder -> initator
|
||||
*/
|
||||
#[repr(C)]
|
||||
#[derive(Copy, Clone)]
|
||||
struct MessageResponse {
|
||||
f_type : u32,
|
||||
f_type : u8,
|
||||
f_reserved : [u8; 3],
|
||||
f_sender : u32,
|
||||
f_receiver : u32,
|
||||
f_ephemeral : [u8; SIZE_X25519_POINT],
|
||||
@@ -109,3 +137,93 @@ impl Into<Vec<u8>> for MessageResponse {
|
||||
array.to_vec()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for MessageResponse {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f,
|
||||
"MessageResponse {{ type = {} }}",
|
||||
self.f_type
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
impl PartialEq for MessageResponse {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.f_type == other.f_type &&
|
||||
self.f_reserved == other.f_reserved &&
|
||||
self.f_sender == other.f_sender &&
|
||||
self.f_receiver == other.f_receiver &&
|
||||
self.f_ephemeral == other.f_ephemeral &&
|
||||
self.f_empty == other.f_empty
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn message_response_identity() {
|
||||
let msg = MessageResponse {
|
||||
f_type : TYPE_RESPONSE,
|
||||
f_reserved : [0u8; 3],
|
||||
f_sender : 146252,
|
||||
f_receiver : 554442,
|
||||
f_ephemeral : [
|
||||
// ephemeral public key
|
||||
0xc1, 0x66, 0x0a, 0x0c, 0xdc, 0x0f, 0x6c, 0x51,
|
||||
0x0f, 0xc2, 0xcc, 0x51, 0x52, 0x0c, 0xde, 0x1e,
|
||||
0xf7, 0xf1, 0xca, 0x90, 0x86, 0x72, 0xad, 0x67,
|
||||
0xea, 0x89, 0x45, 0x44, 0x13, 0x56, 0x52, 0x1f
|
||||
],
|
||||
f_empty : [
|
||||
// tag
|
||||
0x60, 0x0e, 0x1e, 0x95, 0x41, 0x6b, 0x52, 0x05,
|
||||
0xa2, 0x09, 0xe1, 0xbf, 0x40, 0x05, 0x2f, 0xde
|
||||
]
|
||||
};
|
||||
|
||||
let buf : Vec<u8> = msg.into();
|
||||
assert_eq!(msg, MessageResponse::from(&buf[..]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn message_initiate_identity() {
|
||||
let msg = MessageInitiate {
|
||||
f_type : TYPE_RESPONSE,
|
||||
f_reserved : [0u8; 3],
|
||||
f_sender : 575757,
|
||||
f_ephemeral : [
|
||||
// ephemeral public key
|
||||
0xc1, 0x66, 0x0a, 0x0c, 0xdc, 0x0f, 0x6c, 0x51,
|
||||
0x0f, 0xc2, 0xcc, 0x51, 0x52, 0x0c, 0xde, 0x1e,
|
||||
0xf7, 0xf1, 0xca, 0x90, 0x86, 0x72, 0xad, 0x67,
|
||||
0xea, 0x89, 0x45, 0x44, 0x13, 0x56, 0x52, 0x1f
|
||||
],
|
||||
f_static : [
|
||||
// encrypted static public key
|
||||
0xdc, 0x33, 0x90, 0x15, 0x8f, 0x82, 0x3e, 0x06,
|
||||
0x44, 0xa0, 0xde, 0x4c, 0x15, 0x6c, 0x5d, 0xa4,
|
||||
0x65, 0x99, 0xf6, 0x6c, 0xa1, 0x14, 0x77, 0xf9,
|
||||
0xeb, 0x6a, 0xec, 0xc3, 0x3c, 0xda, 0x47, 0xe1,
|
||||
|
||||
// tag
|
||||
0x45, 0xac, 0x8d, 0x43, 0xea, 0x1b, 0x2f, 0x02,
|
||||
0x45, 0x5d, 0x86, 0x37, 0xee, 0x83, 0x6b, 0x42
|
||||
],
|
||||
f_timestamp : [
|
||||
// timestamp
|
||||
0x4f, 0x1c, 0x60, 0xec, 0x0e, 0xf6, 0x36, 0xf0,
|
||||
0x78, 0x28, 0x57, 0x42,
|
||||
|
||||
// tag
|
||||
0x60, 0x0e, 0x1e, 0x95, 0x41, 0x6b, 0x52, 0x05,
|
||||
0xa2, 0x09, 0xe1, 0xbf, 0x40, 0x05, 0x2f, 0xde
|
||||
]
|
||||
};
|
||||
|
||||
let buf : Vec<u8> = msg.into();
|
||||
assert_eq!(msg, MessageInitiate::from(&buf[..]));
|
||||
}
|
||||
}
|
||||
|
||||
5
src/noise.rs
Normal file
5
src/noise.rs
Normal file
@@ -0,0 +1,5 @@
|
||||
use crate::machine::Peer;
|
||||
|
||||
fn create_initiation(st : &mut Peer) {
|
||||
|
||||
}
|
||||
Reference in New Issue
Block a user