Add unit tests

This commit is contained in:
Mathias Hall-Andersen
2019-07-11 21:53:44 +02:00
parent 9154c997fa
commit 0c05104e8b
4 changed files with 214 additions and 29 deletions

View File

@@ -1,5 +1,6 @@
mod messages;
mod machine;
mod noise;
#[cfg(test)]
mod tests {

View File

@@ -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(())
}
}
}

View File

@@ -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
View File

@@ -0,0 +1,5 @@
use crate::machine::Peer;
fn create_initiation(st : &mut Peer) {
}