Checking of mac2 fields on initiation & response

In addition, moved the rng out.
This will allow allocating one instance per worker,
rather than every call.

5 files changed, 119 insertions, 76 deletions

diff --git a/src/handshake/device.rs b/src/handshake/device.rs
index a1f882d..7417949 100644
--- a/src/handshake/device.rs
+++ b/src/handshake/device.rs
@@ -1,9 +1,9 @@
 use spin::RwLock;
 use std::collections::HashMap;
+use std::net::SocketAddr;
 use zerocopy::AsBytes;
 
 use rand::prelude::*;
-use rand::rngs::OsRng;
 
 use x25519_dalek::PublicKey;
 use x25519_dalek::StaticSecret;
@@ -159,15 +159,19 @@ where
     /// # Arguments
     ///
     /// * `pk` - Public key of peer to initiate handshake for
-    pub fn begin(&self, pk: &PublicKey) -> Result<Vec<u8>, HandshakeError> {
+    pub fn begin<R: RngCore + CryptoRng>(
+        &self,
+        rng: &mut R,
+        pk: &PublicKey,
+    ) -> Result<Vec<u8>, HandshakeError> {
         match self.pk_map.get(pk.as_bytes()) {
             None => Err(HandshakeError::UnknownPublicKey),
             Some(peer) => {
-                let sender = self.allocate(peer);
+                let sender = self.allocate(rng, peer);
 
                 let mut msg = Initiation::default();
 
-                noise::create_initiation(self, peer, sender, &mut msg.noise)?;
+                noise::create_initiation(rng, self, peer, sender, &mut msg.noise)?;
 
                 // add macs to initation
 
@@ -185,30 +189,51 @@ where
     /// # Arguments
     ///
     /// * `msg` - Byte slice containing the message (untrusted input)
-    pub fn process(&self, msg: &[u8]) -> Result<Output<T>, HandshakeError> {
+    pub fn process<R: RngCore + CryptoRng>(
+        &self,
+        rng: &mut R,
+        msg: &[u8],               // message buffer
+        src: Option<&SocketAddr>, // optional source address, set when "under load"
+    ) -> Result<Output<T>, HandshakeError> {
         match msg.get(0) {
             Some(&TYPE_INITIATION) => {
+                // parse message
                 let msg = Initiation::parse(msg)?;
 
                 // check mac1 field
                 self.macs.check_mac1(msg.noise.as_bytes(), &msg.macs)?;
 
-                // check ratelimiter
+                // check mac2 field
+                if let Some(src) = src {
+                    if !self.macs.check_mac2(msg.noise.as_bytes(), src, &msg.macs) {
+                        let mut reply = Default::default();
+                        self.macs.create_cookie_reply(
+                            rng,
+                            msg.noise.f_sender.get(),
+                            src,
+                            &msg.macs,
+                            &mut reply,
+                        );
+                        return Ok((None, Some(reply.as_bytes().to_owned()), None));
+                    }
+                }
+
                 // consume the initiation
                 let (peer, st) = noise::consume_initiation(self, &msg.noise)?;
 
                 // allocate new index for response
-                let sender = self.allocate(peer);
+                let sender = self.allocate(rng, peer);
 
                 // prepare memory for response, TODO: take slice for zero allocation
                 let mut resp = Response::default();
 
                 // create response (release id on error)
-                let keys =
-                    noise::create_response(peer, sender, st, &mut resp.noise).map_err(|e| {
+                let keys = noise::create_response(rng, peer, sender, st, &mut resp.noise).map_err(
+                    |e| {
                         self.release(sender);
                         e
-                    })?;
+                    },
+                )?;
 
                 // add macs to response
                 peer.macs
@@ -217,7 +242,7 @@ where
 
                 // return unconfirmed keypair and the response as vector
                 Ok((
-                    peer.identifier,
+                    Some(peer.identifier),
                     Some(resp.as_bytes().to_owned()),
                     Some(keys),
                 ))
@@ -228,6 +253,22 @@ where
                 // check mac1 field
                 self.macs.check_mac1(msg.noise.as_bytes(), &msg.macs)?;
 
+                // check mac2 field
+                if let Some(src) = src {
+                    if !self.macs.check_mac2(msg.noise.as_bytes(), src, &msg.macs) {
+                        let mut reply = Default::default();
+                        self.macs.create_cookie_reply(
+                            rng,
+                            msg.noise.f_sender.get(),
+                            src,
+                            &msg.macs,
+                            &mut reply,
+                        );
+                        return Ok((None, Some(reply.as_bytes().to_owned()), None));
+                    }
+                }
+
+                // consume inner playload
                 noise::consume_response(self, &msg.noise)
             }
             Some(&TYPE_COOKIEREPLY) => {
@@ -239,8 +280,9 @@ where
                 // validate cookie reply
                 peer.macs.lock().process(&msg)?;
 
-                // this prompts no new message
-                Ok((peer.identifier, None, None))
+                // this prompts no new message and
+                // DOES NOT cryptographically verify the peer
+                Ok((None, None, None))
             }
             _ => Err(HandshakeError::InvalidMessageFormat),
         }
@@ -270,9 +312,7 @@ where
     // Internal function
     //
     // Allocated a new receiver identifier for the peer
-    fn allocate(&self, peer: &Peer<T>) -> u32 {
-        let mut rng = OsRng::new().unwrap();
-
+    fn allocate<R: RngCore + CryptoRng>(&self, rng: &mut R, peer: &Peer<T>) -> u32 {
         loop {
             let id = rng.gen();
 
@@ -296,6 +336,7 @@ mod tests {
     use super::super::messages::*;
     use super::*;
     use hex;
+    use rand::rngs::OsRng;
 
     #[test]
     fn handshake() {
@@ -332,14 +373,14 @@ mod tests {
 
             // create initiation
 
-            let msg1 = dev1.begin(&pk2).unwrap();
+            let msg1 = dev1.begin(&mut rng, &pk2).unwrap();
 
             println!("msg1 = {} : {} bytes", hex::encode(&msg1[..]), msg1.len());
             println!("msg1 = {:?}", Initiation::parse(&msg1[..]).unwrap());
 
             // process initiation and create response
 
-            let (_, msg2, ks_r) = dev2.process(&msg1).unwrap();
+            let (_, msg2, ks_r) = dev2.process(&mut rng, &msg1, None).unwrap();
 
             let ks_r = ks_r.unwrap();
             let msg2 = msg2.unwrap();
@@ -351,7 +392,7 @@ mod tests {
 
             // process response and obtain confirmed key-pair
 
-            let (_, msg3, ks_i) = dev1.process(&msg2).unwrap();
+            let (_, msg3, ks_i) = dev1.process(&mut rng, &msg2, None).unwrap();
             let ks_i = ks_i.unwrap();
 
             assert!(msg3.is_none(), "Returned message after response");
diff --git a/src/handshake/macs.rs b/src/handshake/macs.rs
index d95489f..65fd7fa 100644
--- a/src/handshake/macs.rs
+++ b/src/handshake/macs.rs
@@ -1,15 +1,14 @@
-use std::time::{Duration, Instant};
-
-use rand::CryptoRng;
-use rand::RngCore;
-
+use rand::{CryptoRng, RngCore};
 use spin::Mutex;
+use std::time::{Duration, Instant};
 
 use blake2::Blake2s;
+use sodiumoxide::crypto::aead::xchacha20poly1305_ietf;
 use subtle::ConstantTimeEq;
 use x25519_dalek::PublicKey;
 
-use sodiumoxide::crypto::aead::xchacha20poly1305_ietf;
+use std::net::SocketAddr;
+use zerocopy::AsBytes;
 
 use super::messages::{CookieReply, MacsFooter};
 use super::types::HandshakeError;
@@ -100,6 +99,23 @@ pub struct Generator {
     cookie: Option<Cookie>,
 }
 
+fn addr_to_mac_bytes(addr: &SocketAddr) -> Vec<u8> {
+    match addr {
+        SocketAddr::V4(addr) => {
+            let mut res = Vec::with_capacity(4 + 2);
+            res.extend(&addr.ip().octets());
+            res.extend(&addr.port().to_le_bytes());
+            res
+        }
+        SocketAddr::V6(addr) => {
+            let mut res = Vec::with_capacity(16 + 2);
+            res.extend(&addr.ip().octets());
+            res.extend(&addr.port().to_le_bytes());
+            res
+        }
+    }
+}
+
 impl Generator {
     /// Initalize a new mac field generator
     ///
@@ -193,12 +209,12 @@ impl Validator {
         }
     }
 
-    fn get_tau(&self, src: &[u8]) -> Result<[u8; SIZE_COOKIE], HandshakeError> {
+    fn get_tau(&self, src: &[u8]) -> Option<[u8; SIZE_COOKIE]> {
         let secret = self.secret.lock();
         if secret.birth.elapsed() < Duration::from_secs(SECS_COOKIE_UPDATE) {
-            Ok(MAC!(&secret.value, src))
+            Some(MAC!(&secret.value, src))
         } else {
-            Err(HandshakeError::InvalidMac2)
+            None
         }
     }
 
@@ -219,25 +235,26 @@ impl Validator {
         MAC!(&secret.value, src)
     }
 
-    fn create_cookie_reply<T>(
+    pub fn create_cookie_reply<T>(
         &self,
         rng: &mut T,
         receiver: u32,         // receiver id of incoming message
-        src: &[u8],            // source address of incoming message
+        src: &SocketAddr,      // source address of incoming message
         macs: &MacsFooter,     // footer of incoming message
         msg: &mut CookieReply, // resulting cookie reply
     ) where
         T: RngCore + CryptoRng,
     {
+        let src = addr_to_mac_bytes(src);
         msg.f_receiver.set(receiver);
         rng.fill_bytes(&mut msg.f_nonce);
         XSEAL!(
-            &self.cookie_key,            // key
-            &msg.f_nonce,                // nonce
-            &macs.f_mac1,                // ad
-            &self.get_set_tau(rng, src), // pt
-            &mut msg.f_cookie,           // ct
-            &mut msg.f_cookie_tag        // tag
+            &self.cookie_key,             // key
+            &msg.f_nonce,                 // nonce
+            &macs.f_mac1,                 // ad
+            &self.get_set_tau(rng, &src), // pt
+            &mut msg.f_cookie,            // ct
+            &mut msg.f_cookie_tag         // tagf
         );
     }
 
@@ -256,25 +273,11 @@ impl Validator {
         }
     }
 
-    /// Check the mac2 field against the inner message
-    ///
-    /// # Arguments
-    ///
-    /// - inner: The inner message covered by the mac1 field
-    /// - src: Source address
-    /// - macs: The mac footer
-    pub fn check_mac2(
-        &self,
-        inner: &[u8],
-        src: &[u8],
-        macs: &MacsFooter,
-    ) -> Result<(), HandshakeError> {
-        let tau = self.get_tau(src)?;
-        let valid_mac2: bool = MAC!(&tau, inner, macs.f_mac1).ct_eq(&macs.f_mac2).into();
-        if !valid_mac2 {
-            Err(HandshakeError::InvalidMac2)
-        } else {
-            Ok(())
+    pub fn check_mac2(&self, inner: &[u8], src: &SocketAddr, macs: &MacsFooter) -> bool {
+        let src = addr_to_mac_bytes(src);
+        match self.get_tau(&src) {
+            Some(tau) => MAC!(&tau, inner, macs.f_mac1).ct_eq(&macs.f_mac2).into(),
+            None => false,
         }
     }
 }
@@ -295,10 +298,11 @@ mod tests {
 
     proptest! {
         #[test]
-        fn test_cookie_reply(inner1 : Vec<u8>, inner2 : Vec<u8>, src: Vec<u8>, receiver : u32) {
+        fn test_cookie_reply(inner1 : Vec<u8>, inner2 : Vec<u8>, receiver : u32) {
             let mut msg = CookieReply::default();
-            let mut rng = OsRng::new().unwrap();
+            let mut rng = OsRng::new().expect("failed to create rng");
             let mut macs = MacsFooter::default();
+            let src = "127.0.0.1:8080".parse().unwrap();
             let (validator, mut generator) = new_validator_generator();
 
             // generate mac1 for first message
@@ -308,10 +312,8 @@ mod tests {
 
             // check validity of mac1
             validator.check_mac1(&inner1[..], &macs).expect("mac1 of inner1 did not validate");
-
-            // generate cookie reply in response
-            validator.create_cookie_reply(&mut rng, receiver, &src[..], &macs, &mut msg);
-            assert_eq!(msg.f_receiver.get(), receiver);
+            assert_eq!(validator.check_mac2(&inner1[..], &src, &macs), false, "mac2 of inner2 did not validate");
+            validator.create_cookie_reply(&mut rng, receiver, &src, &macs, &mut msg);
 
             // consume cookie reply
             generator.process(&msg).expect("failed to process CookieReply");
@@ -323,7 +325,7 @@ mod tests {
 
             // check validity of mac1 and mac2
             validator.check_mac1(&inner2[..], &macs).expect("mac1 of inner2 did not validate");
-            validator.check_mac2(&inner2[..], &src[..], &macs).expect("mac2 of inner2 did not validate");
+            assert!(validator.check_mac2(&inner2[..], &src, &macs), "mac2 of inner2 did not validate");
         }
     }
 }
diff --git a/src/handshake/noise.rs b/src/handshake/noise.rs
index 0534e97..a93d544 100644
--- a/src/handshake/noise.rs
+++ b/src/handshake/noise.rs
@@ -10,6 +10,7 @@ use hmac::Hmac;
 use sodiumoxide::crypto::aead::chacha20poly1305;
 
 use rand::rngs::OsRng;
+use rand::{CryptoRng, RngCore};
 
 use generic_array::typenum::*;
 use generic_array::GenericArray;
@@ -165,14 +166,13 @@ mod tests {
     }
 }
 
-pub fn create_initiation<T: Copy>(
+pub fn create_initiation<T: Copy, R: RngCore + CryptoRng>(
+    rng: &mut R,
     device: &Device<T>,
     peer: &Peer<T>,
     sender: u32,
     msg: &mut NoiseInitiation,
 ) -> Result<(), HandshakeError> {
-    let mut rng = OsRng::new().unwrap();
-
     // initialize state
 
     let ck = INITIAL_CK;
@@ -183,7 +183,7 @@ pub fn create_initiation<T: Copy>(
 
     // (E_priv, E_pub) := DH-Generate()
 
-    let eph_sk = StaticSecret::new(&mut rng);
+    let eph_sk = StaticSecret::new(rng);
     let eph_pk = PublicKey::from(&eph_sk);
 
     // C := Kdf(C, E_pub)
@@ -316,20 +316,23 @@ pub fn consume_initiation<'a, T: Copy>(
     Ok((peer, (msg.f_sender.get(), eph_r_pk, hs, ck)))
 }
 
-pub fn create_response<T: Copy>(
+pub fn create_response<T: Copy, R: RngCore + CryptoRng>(
+    rng: &mut R,
     peer: &Peer<T>,
     sender: u32,             // sending identifier
     state: TemporaryState,   // state from "consume_initiation"
     msg: &mut NoiseResponse, // resulting response
 ) -> Result<KeyPair, HandshakeError> {
+
+    // unpack state
+
     let (receiver, eph_r_pk, hs, ck) = state;
-    let mut rng = OsRng::new().unwrap();
     msg.f_sender.set(sender);
     msg.f_receiver.set(receiver);
 
     // (E_priv, E_pub) := DH-Generate()
 
-    let eph_sk = StaticSecret::new(&mut rng);
+    let eph_sk = StaticSecret::new(rng);
     let eph_pk = PublicKey::from(&eph_sk);
 
     // C := Kdf1(C, E_pub)
@@ -454,8 +457,8 @@ pub fn consume_response<T: Copy>(
     // return confirmed key-pair
 
     Ok((
-        peer.identifier,
-        None,
+        Some(peer.identifier), // proves overship of the public key (e.g. for updating the endpoint)
+        None,                  // no response message
         Some(KeyPair {
             confirmed: true,
             send: Key {
diff --git a/src/handshake/timestamp.rs b/src/handshake/timestamp.rs
index b78f1cd..b5bd9f0 100644
--- a/src/handshake/timestamp.rs
+++ b/src/handshake/timestamp.rs
@@ -2,7 +2,7 @@ use std::time::{SystemTime, UNIX_EPOCH};
 
 pub type TAI64N = [u8; 12];
 
-const TAI64_EPOCH: u64 = 0x4000000000000000;
+const TAI64_EPOCH: u64 = 0x400000000000000a;
 
 pub const ZERO: TAI64N = [0u8; 12];
 
diff --git a/src/handshake/types.rs b/src/handshake/types.rs
index 38b044e..08c43d0 100644
--- a/src/handshake/types.rs
+++ b/src/handshake/types.rs
@@ -43,7 +43,6 @@ pub enum HandshakeError {
     OldTimestamp,
     InvalidState,
     InvalidMac1,
-    InvalidMac2
 }
 
 impl fmt::Display for HandshakeError {
@@ -58,7 +57,6 @@ impl fmt::Display for HandshakeError {
             HandshakeError::OldTimestamp => write!(f, "Timestamp is less/equal to the newest"),
             HandshakeError::InvalidState => write!(f, "Message does not apply to handshake state"),
             HandshakeError::InvalidMac1 => write!(f, "Message has invalid mac1 field"),
-            HandshakeError::InvalidMac2 => write!(f, "Message has invalid mac2 field"),
         }
     }
 }
@@ -74,8 +72,7 @@ impl Error for HandshakeError {
 }
 
 pub type Output<T> = (
-    T, // external identifier associated with peer
-    // (e.g. a reference or vector index)
+    Option<T>,       // external identifier associated with peer
     Option<Vec<u8>>, // message to send
     Option<KeyPair>, // resulting key-pair of successful handshake
 );