Prepare for resuse of message buffers for response

4 files changed, 31 insertions, 22 deletions

diff --git a/src/handshake/device.rs b/src/handshake/device.rs
index 7417949..809d7a3 100644
--- a/src/handshake/device.rs
+++ b/src/handshake/device.rs
@@ -10,7 +10,7 @@ use x25519_dalek::StaticSecret;
 
 use super::macs;
 use super::messages::{CookieReply, Initiation, Response};
-use super::messages::{TYPE_COOKIEREPLY, TYPE_INITIATION, TYPE_RESPONSE};
+use super::messages::{TYPE_COOKIE_REPLY, TYPE_INITIATION, TYPE_RESPONSE};
 use super::noise;
 use super::peer::Peer;
 use super::types::*;
@@ -271,7 +271,7 @@ where
                 // consume inner playload
                 noise::consume_response(self, &msg.noise)
             }
-            Some(&TYPE_COOKIEREPLY) => {
+            Some(&TYPE_COOKIE_REPLY) => {
                 let msg = CookieReply::parse(msg)?;
 
                 // lookup peer
diff --git a/src/handshake/macs.rs b/src/handshake/macs.rs
index f465099..d5dd95d 100644
--- a/src/handshake/macs.rs
+++ b/src/handshake/macs.rs
@@ -9,7 +9,7 @@ use x25519_dalek::PublicKey;
 
 use std::net::SocketAddr;
 
-use super::messages::{CookieReply, MacsFooter};
+use super::messages::{CookieReply, MacsFooter, TYPE_COOKIE_REPLY};
 use super::types::HandshakeError;
 
 const LABEL_MAC1: &[u8] = b"mac1----";
@@ -219,21 +219,25 @@ impl Validator {
 
     fn get_set_tau<R: RngCore + CryptoRng>(&self, rng: &mut R, src: &[u8]) -> [u8; SIZE_COOKIE] {
         // check if current value is still valid
-        let secret = self.secret.read();
-        if secret.birth.elapsed() < Duration::from_secs(SECS_COOKIE_UPDATE) {
-            return MAC!(&secret.value, src);
-        };
+        {
+            let secret = self.secret.read();
+            if secret.birth.elapsed() < Duration::from_secs(SECS_COOKIE_UPDATE) {
+                return MAC!(&secret.value, src);
+            };
+        }
 
         // take write lock, check again
-        let mut secret = self.secret.write();
-        if secret.birth.elapsed() < Duration::from_secs(SECS_COOKIE_UPDATE) {
-            return MAC!(&secret.value, src);
-        };
-
-        // set new random cookie secret
-        rng.fill_bytes(&mut secret.value);
-        secret.birth = Instant::now();
-        MAC!(&secret.value, src)
+        {
+            let mut secret = self.secret.write();
+            if secret.birth.elapsed() < Duration::from_secs(SECS_COOKIE_UPDATE) {
+                return MAC!(&secret.value, src);
+            };
+
+            // set new random cookie secret
+            rng.fill_bytes(&mut secret.value);
+            secret.birth = Instant::now();
+            MAC!(&secret.value, src)
+        }
     }
 
     pub fn create_cookie_reply<R: RngCore + CryptoRng>(
@@ -245,6 +249,7 @@ impl Validator {
         msg: &mut CookieReply, // resulting cookie reply
     ) {
         let src = addr_to_mac_bytes(src);
+        msg.f_type.set(TYPE_COOKIE_REPLY as u32);
         msg.f_receiver.set(receiver);
         rng.fill_bytes(&mut msg.f_nonce);
         XSEAL!(
diff --git a/src/handshake/messages.rs b/src/handshake/messages.rs
index 52ddac1..d068f26 100644
--- a/src/handshake/messages.rs
+++ b/src/handshake/messages.rs
@@ -19,7 +19,7 @@ const SIZE_X25519_POINT: usize = 32; // x25519 public key
 
 pub const TYPE_INITIATION: u8 = 1;
 pub const TYPE_RESPONSE: u8 = 2;
-pub const TYPE_COOKIEREPLY: u8 = 3;
+pub const TYPE_COOKIE_REPLY: u8 = 3;
 
 /* Handshake messsages */
 
@@ -40,7 +40,7 @@ pub struct Initiation {
 #[repr(packed)]
 #[derive(Copy, Clone, FromBytes, AsBytes)]
 pub struct CookieReply {
-    f_type: U32<LittleEndian>,
+    pub f_type: U32<LittleEndian>,
     pub f_receiver: U32<LittleEndian>,
     pub f_nonce: [u8; SIZE_XNONCE],
     pub f_cookie: [u8; SIZE_COOKIE],
@@ -59,7 +59,7 @@ pub struct MacsFooter {
 #[repr(packed)]
 #[derive(Copy, Clone, FromBytes, AsBytes)]
 pub struct NoiseInitiation {
-    f_type: U32<LittleEndian>,
+    pub f_type: U32<LittleEndian>,
     pub f_sender: U32<LittleEndian>,
     pub f_ephemeral: [u8; SIZE_X25519_POINT],
     pub f_static: [u8; SIZE_X25519_POINT],
@@ -71,7 +71,7 @@ pub struct NoiseInitiation {
 #[repr(packed)]
 #[derive(Copy, Clone, FromBytes, AsBytes)]
 pub struct NoiseResponse {
-    f_type: U32<LittleEndian>,
+    pub f_type: U32<LittleEndian>,
     pub f_sender: U32<LittleEndian>,
     pub f_receiver: U32<LittleEndian>,
     pub f_ephemeral: [u8; SIZE_X25519_POINT],
@@ -111,7 +111,7 @@ impl CookieReply {
         let msg: LayoutVerified<B, Self> =
             LayoutVerified::new(bytes).ok_or(HandshakeError::InvalidMessageFormat)?;
 
-        if msg.f_type.get() != (TYPE_COOKIEREPLY as u32) {
+        if msg.f_type.get() != (TYPE_COOKIE_REPLY as u32) {
             return Err(HandshakeError::InvalidMessageFormat);
         }
 
@@ -142,7 +142,7 @@ impl Default for Initiation {
 impl Default for CookieReply {
     fn default() -> Self {
         Self {
-            f_type: <U32<LittleEndian>>::new(TYPE_COOKIEREPLY as u32),
+            f_type: <U32<LittleEndian>>::new(TYPE_COOKIE_REPLY as u32),
             f_receiver: <U32<LittleEndian>>::ZERO,
             f_nonce: [0u8; SIZE_XNONCE],
             f_cookie: [0u8; SIZE_COOKIE],
diff --git a/src/handshake/noise.rs b/src/handshake/noise.rs
index d66304c..6532f4d 100644
--- a/src/handshake/noise.rs
+++ b/src/handshake/noise.rs
@@ -16,6 +16,7 @@ use generic_array::GenericArray;
 
 use super::device::Device;
 use super::messages::{NoiseInitiation, NoiseResponse};
+use super::messages::{TYPE_INITIATION, TYPE_RESPONSE};
 use super::peer::{Peer, State};
 use super::timestamp;
 use super::types::*;
@@ -178,6 +179,7 @@ pub fn create_initiation<T: Copy, R: RngCore + CryptoRng>(
     let hs = INITIAL_HS;
     let hs = HASH!(&hs, peer.pk.as_bytes());
 
+    msg.f_type.set(TYPE_INITIATION as u32);
     msg.f_sender.set(sender);
 
     // (E_priv, E_pub) := DH-Generate()
@@ -325,6 +327,8 @@ pub fn create_response<T: Copy, R: RngCore + CryptoRng>(
     // unpack state
 
     let (receiver, eph_r_pk, hs, ck) = state;
+
+    msg.f_type.set(TYPE_RESPONSE as u32);
     msg.f_sender.set(sender);
     msg.f_receiver.set(receiver);