Finish handshake exchange

4 files changed, 239 insertions, 107 deletions

diff --git a/src/device.rs b/src/device.rs
index 0635a5e..1acbbd4 100644
--- a/src/device.rs
+++ b/src/device.rs
@@ -1,4 +1,4 @@
-use std::sync::Mutex;
+use spin::RwLock;
 use std::collections::HashMap;
 
 use rand::prelude::*;
@@ -13,11 +13,11 @@ use crate::types::*;
 use crate::peer::Peer;
 
 pub struct Device {
-    pub sk : StaticSecret,              // static secret key
-    pub pk : PublicKey,                 // static public key
-    peers  : Vec<Peer>,                 // peer index  -> state
-    pkmap  : HashMap<[u8; 32], usize>,  // public key  -> peer index
-    ids    : Mutex<HashMap<u32, usize>> // receive ids -> peer index
+    pub sk : StaticSecret,               // static secret key
+    pub pk : PublicKey,                  // static public key
+    peers  : Vec<Peer>,                  // peer index  -> state
+    pk_map : HashMap<[u8; 32], usize>,   // public key  -> peer index
+    id_map : RwLock<HashMap<u32, usize>> // receive ids -> peer index
 }
 
 /* A mutable reference to the state machine needs to be held,
@@ -31,11 +31,11 @@ impl Device {
     /// * `sk` - x25519 scalar representing the local private key
     pub fn new(sk : StaticSecret) -> Device {
         Device {
-            pk    : PublicKey::from(&sk),
-            sk    : sk,
-            peers : vec![],
-            pkmap : HashMap::new(),
-            ids   : Mutex::new(HashMap::new())
+            pk     : PublicKey::from(&sk),
+            sk     : sk,
+            peers  : vec![],
+            pk_map : HashMap::new(),
+            id_map : RwLock::new(HashMap::new())
         }
     }
 
@@ -48,7 +48,7 @@ impl Device {
     pub fn add(&mut self, pk : PublicKey) -> Result<(), ConfigError> {
         // check that the pk is not added twice
 
-        if let Some(_) = self.pkmap.get(pk.as_bytes()) {
+        if let Some(_) = self.pk_map.get(pk.as_bytes()) {
             return Err(ConfigError::new("Duplicate public key"));
         };
 
@@ -61,7 +61,7 @@ impl Device {
         // map : pk -> new index
 
         let idx = self.peers.len();
-        self.pkmap.insert(*pk.as_bytes(), idx);
+        self.pk_map.insert(*pk.as_bytes(), idx);
 
         // map : new index -> peer
 
@@ -83,7 +83,7 @@ impl Device {
     ///
     /// The call might fail if the public key is not found
     pub fn psk(&mut self, pk : PublicKey, psk : Option<Psk>) -> Result<(), ConfigError> {
-        match self.pkmap.get(pk.as_bytes()) {
+        match self.pk_map.get(pk.as_bytes()) {
             Some(&idx) => {
                 let peer = &mut self.peers[idx];
                 peer.psk = match psk {
@@ -102,7 +102,7 @@ impl Device {
     ///
     /// * `id` - The (sender) id to release
     pub fn release(&self, id : u32) {
-        self.ids.lock().unwrap().remove(&id);
+        self.id_map.write().remove(&id);
     }
 
     /// Begin a new handshake
@@ -111,7 +111,7 @@ impl Device {
     ///
     /// * `pk` - Public key of peer to initiate handshake for
     pub fn begin(&self, pk : &PublicKey) -> Result<Vec<u8>, HandshakeError> {
-        match self.pkmap.get(pk.as_bytes()) {
+        match self.pk_map.get(pk.as_bytes()) {
             None => Err(HandshakeError::UnknownPublicKey),
             Some(&idx) => {
                 let peer = &self.peers[idx];
@@ -130,41 +130,42 @@ impl Device {
         match msg.get(0) {
             Some(&messages::TYPE_INITIATION) => {
                 // consume the initiation
-                let (peer, receiver, hs, ck) = noise::consume_initiation(self, msg)?;
+                let (peer, st) = noise::consume_initiation(self, msg)?;
 
-                // allocate index for response
+                // allocate new index for response
                 let sender = self.allocate(peer.idx);
 
                 // create response
-                noise::create_response(self, peer, sender, receiver, hs, ck).map_err(|e| {
+                noise::create_response(self, peer, sender, st).map_err(|e| {
                     self.release(sender);
                     e
                 })
             },
-            Some(&messages::TYPE_RESPONSE) => {
-                Err(HandshakeError::InvalidMessageFormat)
-            },
+            Some(&messages::TYPE_RESPONSE) => noise::consume_response(self, msg),
             _ => Err(HandshakeError::InvalidMessageFormat)
         }
     }
 
-    pub fn lookup(&self, pk : &PublicKey) -> Result<&Peer, HandshakeError> {
-        match self.pkmap.get(pk.as_bytes()) {
+    pub(crate) fn lookup_pk(&self, pk : &PublicKey) -> Result<&Peer, HandshakeError> {
+        match self.pk_map.get(pk.as_bytes()) {
             Some(&idx) => Ok(&self.peers[idx]),
             _ => Err(HandshakeError::UnknownPublicKey)
         }
     }
-}
 
-impl Device {
-    // allocate a new index (id), for peer with idx
+    pub(crate) fn lookup_id(&self, id : u32) -> Result<&Peer, HandshakeError> {
+        match self.id_map.read().get(&id) {
+            Some(&idx) => Ok(&self.peers[idx]),
+            _ => Err(HandshakeError::UnknownReceiverId)
+        }
+    }
+
     fn allocate(&self, idx : usize) -> u32 {
         let mut rng = OsRng::new().unwrap();
-        let mut table = self.ids.lock().unwrap();
         loop {
             let id = rng.gen();
-            if !table.contains_key(&id) {
-                table.insert(id, idx);
+            if !self.id_map.read().contains_key(&id) {
+                self.id_map.write().insert(id, idx);
                 return id;
             }
         }
diff --git a/src/noise.rs b/src/noise.rs
index 9f127e8..f5d1dc1 100644
--- a/src/noise.rs
+++ b/src/noise.rs
@@ -23,9 +23,13 @@ use crate::device::Device;
 use crate::messages::{Initiation, Response};
 use crate::timestamp;
 
+// HMAC hasher (generic construction)
+
 type HMACBlake2s = Hmac<Blake2s>;
 
-/* Internal functions for processing and creating noise messages */
+// convenient alias to pass state temporarily into device.rs and back
+
+type TemporaryState = (u32, PublicKey, GenericArray<u8, U32>, GenericArray<u8, U32>);
 
 const SIZE_CK : usize = 32;
 const SIZE_HS : usize = 32;
@@ -121,6 +125,18 @@ macro_rules! KDF2 {
     }
 }
 
+macro_rules! KDF3 {
+    ($ck:expr, $input:expr) => {
+        {
+            let t0 = HMAC!($ck, $input);
+            let t1 = HMAC!(&t0, &[0x1]);
+            let t2 = HMAC!(&t0, &t1, &[0x2]);
+            let t3 = HMAC!(&t0, &t2, &[0x3]);
+            (t1, t2, t3)
+        }
+    }
+}
+
 macro_rules! SEAL {
     ($key:expr, $aead:expr, $pt:expr, $ct:expr, $tag:expr) => {
         {
@@ -171,8 +187,9 @@ mod tests {
 pub fn create_initiation(
     device : &Device,
     peer : &Peer,
-    id : u32
+    sender : u32
 ) -> Result<Vec<u8>, HandshakeError> {
+
     let mut rng = OsRng::new().unwrap();
     let mut msg : Initiation = Default::default();
 
@@ -182,20 +199,20 @@ pub fn create_initiation(
     let hs = INITIAL_HS;
     let hs = HASH!(&hs, peer.pk.as_bytes());
 
-    msg.f_sender = id;
+    msg.f_sender = sender;
 
     // (E_priv, E_pub) := DH-Generate()
 
-    let sk = StaticSecret::new(&mut rng);
-    let pk = PublicKey::from(&sk);
+    let eph_sk = StaticSecret::new(&mut rng);
+    let eph_pk = PublicKey::from(&eph_sk);
 
     // C := Kdf(C, E_pub)
 
-    let ck = KDF1!(&ck, pk.as_bytes());
+    let ck = KDF1!(&ck, eph_pk.as_bytes());
 
     // msg.ephemeral := E_pub
 
-    msg.f_ephemeral = *pk.as_bytes();
+    msg.f_ephemeral = *eph_pk.as_bytes();
 
     // H := HASH(H, msg.ephemeral)
 
@@ -203,7 +220,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, eph_sk.diffie_hellman(&peer.pk).as_bytes());
 
     // msg.static := Aead(k, 0, S_pub, H)
 
@@ -239,12 +256,7 @@ pub fn create_initiation(
 
     // update state of peer
 
-    peer.set_state(
-        State::InitiationSent{
-            hs : hs,
-            ck : ck
-        }
-    );
+    peer.set_state(State::InitiationSent{hs, ck, eph_sk, sender});
 
     // return message as vector
 
@@ -254,7 +266,7 @@ pub fn create_initiation(
 pub fn consume_initiation<'a>(
     device : &'a Device,
     msg : &[u8]
-) -> Result<(&'a Peer, u32, GenericArray<u8, U32>, GenericArray<u8, U32>), HandshakeError> {
+) -> Result<(&'a Peer, TemporaryState), HandshakeError> {
 
     // parse message
 
@@ -276,10 +288,10 @@ pub fn consume_initiation<'a>(
 
     // (C, k) := Kdf2(C, DH(E_priv, S_pub))
 
-    let eph = PublicKey::from(msg.f_ephemeral);
+    let eph_r_pk = PublicKey::from(msg.f_ephemeral);
     let (ck, key) = KDF2!(
         &ck,
-        device.sk.diffie_hellman(&eph).as_bytes()
+        device.sk.diffie_hellman(&eph_r_pk).as_bytes()
     );
 
     // msg.static := Aead(k, 0, S_pub, H)
@@ -294,7 +306,7 @@ pub fn consume_initiation<'a>(
         &msg.f_static_tag // tag
     )?;
 
-    let peer = device.lookup(&PublicKey::from(pk))?;
+    let peer = device.lookup_pk(&PublicKey::from(pk))?;
 
     // H := Hash(H || msg.static)
 
@@ -318,25 +330,134 @@ pub fn consume_initiation<'a>(
 
     // check and update timestamp
 
-    peer.check_timestamp(&ts)?;
+    peer.check_timestamp(device, &ts)?;
 
     // return state (to create response)
 
-    Ok((peer, msg.f_sender, hs, ck))
+    Ok((peer, (msg.f_sender, eph_r_pk, hs, ck)))
 }
 
 pub fn create_response(
     device   : &Device,
     peer     : &Peer,
-    sender   : u32,
-    receiver : u32,
-    hs       : GenericArray<u8, U32>,
-    ck       : GenericArray<u8, U32>
+    sender   : u32,           // sending identifier
+    state    : TemporaryState // state from "consume_initiation"
 ) -> Result<Output, HandshakeError> {
 
+    let mut rng = OsRng::new().unwrap();
     let mut msg : Response = Default::default();
 
+    let (receiver, eph_r_pk, hs, ck) = state;
+
+    // (E_priv, E_pub) := DH-Generate()
+
+    let eph_sk = StaticSecret::new(&mut rng);
+    let eph_pk = PublicKey::from(&eph_sk);
+
+    // C := Kdf1(C, E_pub)
+
+    let ck = KDF1!(&ck, eph_pk.as_bytes());
+
+    // msg.ephemeral := E_pub
+
+    msg.f_ephemeral = *eph_pk.as_bytes();
+
+    // H := Hash(H || msg.ephemeral)
+
+    let hs = HASH!(&hs, &msg.f_ephemeral);
+
+    // C := Kdf1(C, DH(E_priv, E_pub))
+
+    let ck = KDF1!(&ck, eph_sk.diffie_hellman(&eph_r_pk).as_bytes());
+
+    // C := Kdf1(C, DH(E_priv, S_pub))
+
+    let ck = KDF1!(&ck, eph_sk.diffie_hellman(&peer.pk).as_bytes());
+
+    // (C, tau, k) := Kdf3(C, Q)
+
+    let (ck, tau, key) = KDF3!(&ck, &peer.psk);
+
+    // H := Hash(H || tau)
+
+    let hs = HASH!(&hs, tau);
+
+    // msg.empty := Aead(k, 0, [], H)
+
+    SEAL!(
+        &key,
+        &hs,                 // ad
+        &[],                 // pt
+        &mut [],             // ct
+        &mut msg.f_empty_tag // tag
+    );
+
+    // H := Hash(H || msg.empty)
+
+    // let hs = HASH!(&hs, &msg.f_empty_tag); // not strictly needed
+
+    // derive key-pair
+
+    let (key_recv, key_send) = KDF2!(&ck, &[]);
+
+    Ok(Output(None, None))
+}
+
+pub fn consume_response(
+    device : &Device,
+    msg : &[u8]
+) -> Result<Output, HandshakeError> {
+
     // parse message
 
+    let msg = Response::try_from(msg)?;
+
+    // retrieve peer and associated state
+
+    let peer = device.lookup_id(msg.f_receiver)?;
+    let (hs, ck, sender, eph_sk) = match peer.get_state() {
+        State::Reset => Err(HandshakeError::InvalidState),
+        State::InitiationSent{hs, ck, sender, eph_sk} => Ok((hs, ck, sender, eph_sk))
+    }?;
+
+    // C := Kdf1(C, E_pub)
+
+    let ck = KDF1!(&ck, &msg.f_ephemeral);
+
+    // H := Hash(H || msg.ephemeral)
+
+    let hs = HASH!(&hs, &msg.f_ephemeral);
+
+    // C := Kdf1(C, DH(E_priv, E_pub))
+
+    let eph_r_pk = PublicKey::from(msg.f_ephemeral);
+    let ck = KDF1!(&ck, eph_sk.diffie_hellman(&eph_r_pk).as_bytes());
+
+    // C := Kdf1(C, DH(E_priv, S_pub))
+
+    let ck = KDF1!(&ck, eph_sk.diffie_hellman(&peer.pk).as_bytes());
+
+    // (C, tau, k) := Kdf3(C, Q)
+
+    let (ck, tau, key) = KDF3!(&ck, &peer.psk);
+
+    // H := Hash(H || tau)
+
+    let hs = HASH!(&hs, tau);
+
+    // msg.empty := Aead(k, 0, [], H)
+
+    OPEN!(
+        &key,
+        &hs,             // ad
+        &mut [],         // pt
+        &[],             // ct
+        &msg.f_empty_tag // tag
+    );
+
+    // derive key-pair
+
+    let (key_send, key_recv) = KDF2!(&ck, &[]);
+
     Ok(Output(None, None))
 }
diff --git a/src/peer.rs b/src/peer.rs
index e656d56..7678e85 100644
--- a/src/peer.rs
+++ b/src/peer.rs
@@ -4,10 +4,12 @@ use generic_array::typenum::U32;
 use generic_array::GenericArray;
 
 use x25519_dalek::PublicKey;
+use x25519_dalek::StaticSecret;
 use x25519_dalek::SharedSecret;
 
 use crate::types::*;
 use crate::timestamp;
+use crate::device::Device;
 
 /* Represents the recomputation and state of a peer.
  *
@@ -15,32 +17,49 @@ use crate::timestamp;
  */
 
 pub struct Peer {
-    pub idx   : usize,
+    // internal identifier
+    pub(crate) idx : usize,
 
     // mutable state
     state     : Mutex<State>,
     timestamp : Mutex<Option<timestamp::TAI64N>>,
 
     // constant state
-    pub pk  : PublicKey,     // public key of peer
-    pub ss  : SharedSecret,  // precomputed DH(static, static)
-    pub psk : Psk            // psk of peer
+    pub(crate) pk  : PublicKey,     // public key of peer
+    pub(crate) ss  : SharedSecret,  // precomputed DH(static, static)
+    pub(crate) psk : Psk            // psk of peer
 }
 
-#[derive(Debug, Copy, Clone)]
 pub enum State {
     Reset,
     InitiationSent{
-        hs : GenericArray<u8, U32>,
-        ck : GenericArray<u8, U32>
+        sender : u32, // assigned sender id
+        eph_sk : StaticSecret,
+        hs     : GenericArray<u8, U32>,
+        ck     : GenericArray<u8, U32>
     },
 }
 
+impl Clone for State {
+    fn clone(&self) -> State {
+        match self {
+            State::Reset => State::Reset,
+            State::InitiationSent{sender, eph_sk, hs, ck} =>
+                State::InitiationSent{
+                    sender : *sender,
+                    eph_sk : StaticSecret::from(eph_sk.to_bytes()),
+                    hs     : *hs,
+                    ck     : *ck
+                }
+        }
+    }
+}
+
 impl Peer {
     pub fn new(
-        idx : usize,
-        pk  : PublicKey,    // public key of peer
-        ss  : SharedSecret  // precomputed DH(static, static)
+        idx    : usize,
+        pk     : PublicKey,    // public key of peer
+        ss     : SharedSecret  // precomputed DH(static, static)
     ) -> Self {
         Self {
             idx       : idx,
@@ -56,7 +75,7 @@ impl Peer {
     ///
     /// # Arguments
     pub fn get_state(&self) -> State {
-        *self.state.lock().unwrap()
+        self.state.lock().unwrap().clone()
     }
 
     /// Set the state of the peer unconditionally
@@ -71,60 +90,45 @@ impl Peer {
         *state = state_new;
     }
 
-    /// # Arguments
-    ///
-    /// * ts_new - The timestamp
-    ///
-    /// # Returns
-    ///
-    /// A Boolean indicating if the state was updated
-    pub fn check_timestamp(&self,
-                           timestamp_new : &timestamp::TAI64N) -> Result<(), HandshakeError> {
-
-        let mut timestamp = self.timestamp.lock().unwrap();
-        match *timestamp {
-            None => Ok(()),
-            Some(timestamp_old) => if timestamp::compare(&timestamp_old, &timestamp_new) {
-                *timestamp = Some(*timestamp_new);
-                Ok(())
-            } else {
-                Err(HandshakeError::OldTimestamp)
-            }
-        }
-    }
-
     /// Set the mutable state of the peer conditioned on the timestamp being newer
     ///
     /// # Arguments
     ///
     /// * st_new - The updated state of the peer
     /// * ts_new - The associated timestamp
-    ///
-    /// # Returns
-    ///
-    /// A Boolean indicating if the state was updated
-    pub fn set_state_timestamp(
+    pub fn check_timestamp(
         &self,
-        state_new : State,
+        device : &Device,
         timestamp_new : &timestamp::TAI64N
     ) -> Result<(), HandshakeError> {
+
         let mut state = self.state.lock().unwrap();
         let mut timestamp = self.timestamp.lock().unwrap();
-        match *timestamp {
-            None => {
-                // no prior timestamp know
-                *state = state_new;
-                *timestamp = Some(*timestamp_new);
-                Ok(())
-            },
+
+        let update = match *timestamp {
+            None => true,
             Some(timestamp_old) => if timestamp::compare(&timestamp_old, &timestamp_new) {
-                // new timestamp is strictly greater
-                *state = state_new;
-                *timestamp = Some(*timestamp_new);
-                Ok(())
+                true
             } else {
-                Err(HandshakeError::OldTimestamp)
+                false
             }
+        };
+
+        if update {
+            // release existing identifier
+            match *state {
+                State::InitiationSent{sender, ..} => {
+                    device.release(sender)
+                },
+                _ => ()
+            }
+
+            // reset state and update timestamp
+            *state = State::Reset;
+            *timestamp = Some(*timestamp_new);
+            Ok(())
+        } else {
+            Err(HandshakeError::OldTimestamp)
         }
     }
 }
diff --git a/src/types.rs b/src/types.rs
index a2a7cdd..5500ffc 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -34,8 +34,10 @@ impl Error for ConfigError {
 pub enum HandshakeError {
     DecryptionFailure,
     UnknownPublicKey,
+    UnknownReceiverId,
     InvalidMessageFormat,
-    OldTimestamp
+    OldTimestamp,
+    InvalidState
 }
 
 impl fmt::Display for HandshakeError {
@@ -45,10 +47,14 @@ impl fmt::Display for HandshakeError {
                 write!(f, "Failed to AEAD:OPEN"),
             HandshakeError::UnknownPublicKey =>
                 write!(f, "Unknown public key"),
+            HandshakeError::UnknownReceiverId =>
+                write!(f, "Receiver id not allocated to any handshake"),
             HandshakeError::InvalidMessageFormat =>
                 write!(f, "Invalid handshake message format"),
             HandshakeError::OldTimestamp =>
-                write!(f, "Timestamp is less/equal to the newest")
+                write!(f, "Timestamp is less/equal to the newest"),
+            HandshakeError::InvalidState =>
+                write!(f, "Message does not apply to handshake state")
         }
     }
 }