Add unit tests

4 files changed, 213 insertions, 28 deletions

diff --git a/src/lib.rs b/src/lib.rs
index a1aa18a..33f2d81 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,5 +1,6 @@
 mod messages;
 mod machine;
+mod noise;
 
 #[cfg(test)]
 mod tests {
diff --git a/src/machine.rs b/src/machine.rs
index a9a3904..a8d4d07 100644
--- a/src/machine.rs
+++ b/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(())
+    }
+
+    /// 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);
     }
 
-    /* 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>), ()> {
-        Err(())
+    /// 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(())
+        }
     }
 }
diff --git a/src/messages.rs b/src/messages.rs
index c9c4239..1ec50bf 100644
--- a/src/messages.rs
+++ b/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[..]));
+    }
+}
diff --git a/src/noise.rs b/src/noise.rs
new file mode 100644
index 0000000..189ada0
--- /dev/null
+++ b/src/noise.rs
@@ -0,0 +1,5 @@
+use crate::machine::Peer;
+
+fn create_initiation(st : &mut Peer) {
+
+}