Begin processing of initation

6 files changed, 223 insertions, 47 deletions

diff --git a/src/device.rs b/src/device.rs
index 928863a..85adc69 100644
--- a/src/device.rs
+++ b/src/device.rs
@@ -8,6 +8,7 @@ use x25519_dalek::PublicKey;
 use x25519_dalek::StaticSecret;
 
 use crate::noise;
+use crate::messages;
 use crate::types::*;
 use crate::peer::Peer;
 
@@ -109,9 +110,9 @@ impl Device {
     /// # Arguments
     ///
     /// * `pk` - Public key of peer to initiate handshake for
-    pub fn begin(&self, pk : PublicKey) -> Result<Vec<u8>, HandshakeError> {
+    pub fn begin(&self, pk : &PublicKey) -> Result<Vec<u8>, HandshakeError> {
         match self.pkmap.get(pk.as_bytes()) {
-            None => Err(HandshakeError::new()),
+            None => Err(HandshakeError::UnknownPublicKey),
             Some(&idx) => {
                 let peer = &self.peers[idx];
                 let id = self.allocate(idx);
@@ -120,15 +121,27 @@ impl Device {
         }
     }
 
+    pub fn lookup(&self, pk : &PublicKey) -> Result<&Peer, HandshakeError> {
+        match self.pkmap.get(pk.as_bytes()) {
+            Some(&idx) => Ok(&self.peers[idx]),
+            _ => Err(HandshakeError::UnknownPublicKey)
+        }
+    }
+
     /// Process a handshake message.
     ///
     /// # Arguments
     ///
     /// * `msg` - Byte slice containing the message (untrusted input)
     pub fn process(&self, msg : &[u8]) -> Result<Output, HandshakeError> {
-        // inspect type field
         match msg.get(0) {
-            _ => Err(HandshakeError::new())
+            Some(&messages::TYPE_INITIATION) => {
+                noise::process_initiation(self, msg)
+            },
+            Some(&messages::TYPE_RESPONSE) => {
+                Err(HandshakeError::InvalidMessageFormat)
+            },
+            _ => Err(HandshakeError::InvalidMessageFormat)
         }
     }
 }
@@ -147,3 +160,38 @@ impl Device {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn handshake() {
+        // generate new keypairs
+
+        let mut rng = OsRng::new().unwrap();
+
+        let sk1 = StaticSecret::new(&mut rng);
+        let pk1 = PublicKey::from(&sk1);
+
+        let sk2 = StaticSecret::new(&mut rng);
+        let pk2 = PublicKey::from(&sk2);
+
+        // intialize devices on both ends
+
+        let mut dev1 = Device::new(sk1);
+        let mut dev2 = Device::new(sk2);
+
+        dev1.add(pk2).unwrap();
+        dev2.add(pk1).unwrap();
+
+        // create initiation
+
+        let msg1 = dev1.begin(&pk2).unwrap();
+
+        // process initiation and create response
+
+        let out1 = dev2.process(&msg1).unwrap();
+
+    }
+}
diff --git a/src/messages.rs b/src/messages.rs
index 340f5e0..da26e48 100644
--- a/src/messages.rs
+++ b/src/messages.rs
@@ -1,12 +1,16 @@
 use std::fmt;
 use std::mem;
+use std::convert::TryFrom;
+
+use crate::types::*;
 
 const SIZE_TAG : usize = 16;
 const SIZE_X25519_POINT : usize = 32;
 const SIZE_TIMESTAMP : usize = 12;
 
-pub const TYPE_INITIATION : u32 = 1;
-pub const TYPE_RESPONSE : u32 = 2;
+pub const TYPE_INITIATION : u8 = 1;
+pub const TYPE_RESPONSE : u8 = 2;
+
 
 /* Wireguard handshake (noise) initiation message
  * initator -> responder
@@ -23,22 +27,42 @@ pub struct Initiation {
     pub f_timestamp_tag  : [u8; SIZE_TAG],
 }
 
-impl From<&[u8]> for Initiation {
-    fn from(b: &[u8]) -> Self {
+impl TryFrom<&[u8]> for Initiation {
+
+    type Error = HandshakeError;
+
+    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
+
+        // check length of slice matches message
+
+        if value.len() != mem::size_of::<Self>() {
+            return Err(HandshakeError::InvalidMessageFormat);
+        }
+
         // create owned copy
+
         let mut owned = [0u8; mem::size_of::<Self>()];
         let mut msg : Self;
-        owned.copy_from_slice(b);
+        owned.copy_from_slice(value);
 
         // cast to Initiation
+
         unsafe {
             msg = mem::transmute::<[u8; mem::size_of::<Self>()], Self>(owned);
         };
 
         // correct endianness
+
         msg.f_type = msg.f_type.to_le();
         msg.f_sender = msg.f_sender.to_le();
-        msg
+
+        // check type and reserved fields
+
+        if msg.f_type != (TYPE_INITIATION as u32) {
+            return Err(HandshakeError::InvalidMessageFormat);
+        }
+
+        Ok(msg)
     }
 }
 
@@ -71,7 +95,7 @@ impl fmt::Debug for Initiation {
 impl Default for Initiation {
     fn default() -> Self {
         Self {
-            f_type          : TYPE_INITIATION,
+            f_type          : TYPE_INITIATION as u32,
             f_sender        : 0,
             f_ephemeral     : [0u8; SIZE_X25519_POINT],
             f_static        : [0u8; SIZE_X25519_POINT],
@@ -112,23 +136,43 @@ pub struct Response {
     pub f_empty_tag : [u8; SIZE_TAG],
 }
 
-impl From<&[u8]> for Response {
-    fn from(b: &[u8]) -> Self {
+impl TryFrom<&[u8]> for Response {
+
+    type Error = HandshakeError;
+
+    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
+
+        // check length of slice matches message
+
+        if value.len() != mem::size_of::<Self>() {
+            return Err(HandshakeError::InvalidMessageFormat);
+        }
+
         // create owned copy
+
         let mut owned = [0u8; mem::size_of::<Self>()];
         let mut msg : Self;
-        owned.copy_from_slice(b);
+        owned.copy_from_slice(value);
 
         // cast to MessageResponse
+
         unsafe {
             msg = mem::transmute::<[u8; mem::size_of::<Self>()], Self>(owned);
         };
 
         // correct endianness
+
         msg.f_type = msg.f_type.to_le();
         msg.f_sender = msg.f_sender.to_le();
         msg.f_receiver = msg.f_receiver.to_le();
-        msg
+
+        // check type and reserved fields
+
+        if msg.f_type != (TYPE_RESPONSE as u32) {
+            return Err(HandshakeError::InvalidMessageFormat);
+        }
+
+        Ok(msg)
     }
 }
 
@@ -177,7 +221,7 @@ mod tests {
     #[test]
     fn message_response_identity() {
         let msg = Response {
-            f_type : TYPE_RESPONSE,
+            f_type : TYPE_RESPONSE as u32,
             f_sender : 146252,
             f_receiver : 554442,
             f_ephemeral : [
@@ -195,13 +239,14 @@ mod tests {
         };
 
         let buf : Vec<u8> = msg.into();
-        assert_eq!(msg, Response::from(&buf[..]));
+        let msg_p : Response = Response::try_from(&buf[..]).unwrap();
+        assert_eq!(msg, msg_p);
     }
 
     #[test]
     fn message_initiate_identity() {
         let msg = Initiation {
-            f_type : TYPE_RESPONSE,
+            f_type : TYPE_INITIATION as u32,
             f_sender : 575757,
             f_ephemeral : [
                 // ephemeral public key
@@ -218,7 +263,7 @@ mod tests {
                 0xeb, 0x6a, 0xec, 0xc3, 0x3c, 0xda, 0x47, 0xe1
             ],
             f_static_tag : [
-                // tag
+                // poly1305 tag
                 0x45, 0xac, 0x8d, 0x43, 0xea, 0x1b, 0x2f, 0x02,
                 0x45, 0x5d, 0x86, 0x37, 0xee, 0x83, 0x6b, 0x42
             ],
@@ -228,13 +273,13 @@ mod tests {
                 0x78, 0x28, 0x57, 0x42
             ],
             f_timestamp_tag : [
-                // tag
+                // poly1305 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, Initiation::from(&buf[..]));
+        assert_eq!(msg, Initiation::try_from(&buf[..]).unwrap());
     }
 }
diff --git a/src/noise.rs b/src/noise.rs
index 7cc5f6c..9de03cd 100644
--- a/src/noise.rs
+++ b/src/noise.rs
@@ -1,3 +1,5 @@
+use std::convert::TryFrom;
+
 // DH
 use x25519_dalek::PublicKey;
 use x25519_dalek::StaticSecret;
@@ -15,7 +17,7 @@ use rand::rngs::OsRng;
 use crate::types::*;
 use crate::peer::{State, Peer};
 use crate::device::Device;
-use crate::messages;
+use crate::messages::{Initiation, Response};
 use crate::timestamp;
 
 type HMACBlake2s = Hmac<Blake2s>;
@@ -129,6 +131,19 @@ macro_rules! SEAL {
     }
 }
 
+macro_rules! OPEN {
+    ($key:expr, $aead:expr, $pt:expr, $ct:expr, $tag:expr) => {
+        {
+            let mut aead = ChaCha20Poly1305::new($key, &ZERO_NONCE, $aead);
+            if !aead.decrypt($ct, $pt, $tag) {
+                Err(HandshakeError::DecryptionFailure)
+            } else {
+                Ok(())
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -156,7 +171,7 @@ pub fn create_initiation(
     id : u32
 ) -> Result<Vec<u8>, HandshakeError> {
     let mut rng = OsRng::new().unwrap();
-    let mut msg : messages::Initiation = Default::default();
+    let mut msg : Initiation = Default::default();
 
     // initialize state
 
@@ -185,10 +200,7 @@ pub fn create_initiation(
 
     // (C, k) := Kdf2(C, DH(E_priv, S_pub))
 
-    let (ck, key) = KDF2!(
-        &ck,
-        sk.diffie_hellman(&peer.pk).as_bytes()
-    );
+    let (ck, key) = KDF2!(&ck, sk.diffie_hellman(&peer.pk).as_bytes());
 
     // msg.static := Aead(k, 0, S_pub, H)
 
@@ -206,10 +218,7 @@ pub fn create_initiation(
 
     // (C, k) := Kdf2(C, DH(S_priv, S_pub))
 
-    let (ck, key) = KDF2!(
-        &ck,
-        peer.ss.as_bytes() // precomputed static-static
-    );
+    let (ck, key) = KDF2!(&ck, peer.ss.as_bytes());
 
     // msg.timestamp := Aead(k, 0, Timestamp(), H)
 
@@ -236,9 +245,80 @@ pub fn create_initiation(
 
     // return message as vector
 
-    Ok(messages::Initiation::into(msg))
+    Ok(Initiation::into(msg))
 }
 
-pub fn process_initiation(device : &Device, peer : &Peer) -> Result<Output, ()> {
-    Err(())
+pub fn process_initiation(device : &Device, msg : &[u8]) -> Result<Output, HandshakeError> {
+
+    // parse message
+
+    let msg = Initiation::try_from(msg)?;
+
+    // initialize state
+
+    let ck = INITIAL_CK;
+    let hs = INITIAL_HS;
+    let hs = HASH!(&hs, device.pk.as_bytes());
+
+    // C := Kdf(C, E_pub)
+
+    let ck = KDF1!(&ck, &msg.f_ephemeral);
+
+    // H := HASH(H, msg.ephemeral)
+
+    let hs = HASH!(&hs, &msg.f_ephemeral);
+
+    // (C, k) := Kdf2(C, DH(E_priv, S_pub))
+
+    let eph = PublicKey::from(msg.f_ephemeral);
+    let (ck, key) = KDF2!(
+        &ck,
+        device.sk.diffie_hellman(&eph).as_bytes()
+    );
+
+    // msg.static := Aead(k, 0, S_pub, H)
+
+    let mut pk = [0u8; 32];
+
+    OPEN!(
+        &key,
+        &hs,              // ad
+        &mut pk,          // pt
+        &msg.f_static,    // ct
+        &msg.f_static_tag // tag
+    )?;
+
+    let peer = device.lookup(&PublicKey::from(pk))?;
+
+    // H := Hash(H || msg.static)
+
+    let hs = HASH!(&hs, &msg.f_static, &msg.f_static_tag);
+
+    // (C, k) := Kdf2(C, DH(S_priv, S_pub))
+
+    let (ck, key) = KDF2!(&ck, peer.ss.as_bytes());
+
+    // msg.timestamp := Aead(k, 0, Timestamp(), H)
+
+    let mut ts = timestamp::zero();
+
+    OPEN!(
+        &key,
+        &hs,                 // ad
+        &mut ts,             // pt
+        &msg.f_timestamp,    // ct
+        &msg.f_timestamp_tag // tag
+    )?;
+
+    // update state of peer
+
+    peer.set_state_timestamp(
+        State::InitiationSent{
+            hs : hs,
+            ck : ck
+        },
+        &ts
+    )?;
+
+    Ok(Output(None, None))
 }
diff --git a/src/peer.rs b/src/peer.rs
index 3c81f32..6c4bb3c 100644
--- a/src/peer.rs
+++ b/src/peer.rs
@@ -76,7 +76,7 @@ impl Peer {
         &self,
         state_new : State,
         timestamp_new : &timestamp::TAI64N
-    ) -> bool {
+    ) -> Result<(), HandshakeError> {
         let mut state = self.state.lock().unwrap();
         let mut timestamp = self.timestamp.lock().unwrap();
         match *timestamp {
@@ -84,15 +84,15 @@ impl Peer {
                 // no prior timestamp know
                 *state = state_new;
                 *timestamp = Some(*timestamp_new);
-                true
+                Ok(())
             },
             Some(timestamp_old) => if timestamp::compare(&timestamp_old, &timestamp_new) {
                 // new timestamp is strictly greater
                 *state = state_new;
                 *timestamp = Some(*timestamp_new);
-                true
+                Ok(())
             } else {
-                false
+                Err(HandshakeError::OldTimestamp)
             }
         }
     }
diff --git a/src/timestamp.rs b/src/timestamp.rs
index 6ac20b8..52ab154 100644
--- a/src/timestamp.rs
+++ b/src/timestamp.rs
@@ -1,7 +1,11 @@
 pub type TAI64N = [u8; 12];
 
+pub fn zero() -> TAI64N {
+    [0u8; 12]
+}
+
 pub fn now() -> TAI64N {
-    [0u8; 12] // TODO
+    [0u8; 12] // TODO, return current timestamp
 }
 
 pub fn compare(old : &TAI64N, new : &TAI64N) -> bool {
diff --git a/src/types.rs b/src/types.rs
index 81c1dc7..464ad81 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -36,12 +36,11 @@ impl Error for ConfigError {
 // handshake error
 
 #[derive(Debug)]
-pub struct HandshakeError {}
-
-impl HandshakeError {
-    pub fn new() -> Self {
-        HandshakeError{}
-    }
+pub enum HandshakeError {
+    DecryptionFailure,
+    UnknownPublicKey,
+    InvalidMessageFormat,
+    OldTimestamp
 }
 
 impl fmt::Display for HandshakeError {
@@ -74,8 +73,8 @@ pub struct KeyPair {
 }
 
 pub struct Output (
-    Option<KeyPair>, // resulting key-pair of successful handshake
-    Option<Vec<u8>>  // message to send
+    pub Option<KeyPair>, // resulting key-pair of successful handshake
+    pub Option<Vec<u8>>  // message to send
 );
 
 // per-peer state machine