Move to nested handshake message structure

Having the nested structure:

    Handshake Message:
        Noise part (zerocopy message)
        MAC footer part (zerocopy message)

Greatly simplifies processing the MAC fields,
since the MAC footer covers the noise part, which can
be accessed as bytes using AsBytes.

4 files changed, 242 insertions, 124 deletions

diff --git a/src/handshake/device.rs b/src/handshake/device.rs
index 04e00f9..b28613a 100644
--- a/src/handshake/device.rs
+++ b/src/handshake/device.rs
@@ -7,7 +7,8 @@ use rand::rngs::OsRng;
 use x25519_dalek::PublicKey;
 use x25519_dalek::StaticSecret;
 
-use super::messages;
+use super::messages::{CookieReply, Initiation, Response};
+use super::messages::{TYPE_COOKIEREPLY, TYPE_INITIATION, TYPE_RESPONSE};
 use super::noise;
 use super::peer::Peer;
 use super::types::*;
@@ -170,20 +171,40 @@ where
     /// * `msg` - Byte slice containing the message (untrusted input)
     pub fn process(&self, msg: &[u8]) -> Result<Output<T>, HandshakeError> {
         match msg.get(0) {
-            Some(&messages::TYPE_INITIATION) => {
+            Some(&TYPE_INITIATION) => {
+                let msg = Initiation::parse(msg)?;
+
+                // check mac footer and ratelimiter
+
                 // consume the initiation
-                let (peer, st) = noise::consume_initiation(self, msg)?;
+                let (peer, st) = noise::consume_initiation(self, &msg.noise)?;
 
                 // allocate new index for response
                 let sender = self.allocate(peer);
 
-                // create response (release id on error)
-                noise::create_response(peer, sender, st).map_err(|e| {
+                // create response (release id on error), TODO: take slice
+                let mut resp = Response::default();
+                noise::create_response(peer, sender, st, &mut resp.noise).map_err(|e| {
                     self.release(sender);
                     e
                 })
             }
-            Some(&messages::TYPE_RESPONSE) => noise::consume_response(self, msg),
+            Some(&TYPE_RESPONSE) => {
+                let msg = Response::parse(msg)?;
+
+                // check mac footer and ratelimiter
+
+                noise::consume_response(self, &msg.noise)
+            }
+            Some(&TYPE_COOKIEREPLY) => {
+                let msg = CookieReply::parse(msg)?;
+
+                // validate cookie reply
+
+                // update cookie generator for peer
+
+                unimplemented!()
+            }
             _ => Err(HandshakeError::InvalidMessageFormat),
         }
     }
@@ -235,9 +256,9 @@ where
 
 #[cfg(test)]
 mod tests {
+    use super::super::messages::*;
     use super::*;
     use hex;
-    use messages::*;
 
     #[test]
     fn handshake() {
diff --git a/src/handshake/macs.rs b/src/handshake/macs.rs
index c8b97e3..05546a7 100644
--- a/src/handshake/macs.rs
+++ b/src/handshake/macs.rs
@@ -4,7 +4,7 @@ use blake2::Blake2s;
 use subtle::ConstantTimeEq;
 use x25519_dalek::PublicKey;
 
-use zerocopy::{AsBytes, ByteSlice, FromBytes, LayoutVerified};
+use super::messages::{CookieReply, MacsFooter};
 
 const LABEL_MAC1: &[u8] = b"mac1----";
 const LABEL_COOKIE: &[u8] = b"cookie--";
@@ -38,22 +38,6 @@ macro_rules! MAC {
     }};
 }
 
-#[repr(C)]
-#[derive(Copy, Clone, FromBytes, AsBytes)]
-pub struct MacsFooter {
-    pub f_mac1: [u8; SIZE_MAC],
-    pub f_mac2: [u8; SIZE_MAC],
-}
-
-impl Default for MacsFooter {
-    fn default() -> Self {
-        Self {
-            f_mac1: [0u8; SIZE_MAC],
-            f_mac2: [0u8; SIZE_MAC],
-        }
-    }
-}
-
 struct Generator {
     mac1_key: [u8; 32],
     cookie_value: [u8; 16],
diff --git a/src/handshake/messages.rs b/src/handshake/messages.rs
index 8fa08b3..a2ff19a 100644
--- a/src/handshake/messages.rs
+++ b/src/handshake/messages.rs
@@ -8,10 +8,10 @@ use byteorder::LittleEndian;
 use zerocopy::byteorder::U32;
 use zerocopy::{AsBytes, ByteSlice, FromBytes, LayoutVerified};
 
-use super::macs::MacsFooter;
 use super::timestamp;
 use super::types::*;
 
+const SIZE_MAC: usize = 16; 
 const SIZE_TAG: usize = 16; // poly1305 tag
 const SIZE_NONCE: usize = 16; // xchacha20 nonce
 const SIZE_COOKIE: usize = 16; //
@@ -21,28 +21,20 @@ pub const TYPE_INITIATION: u8 = 1;
 pub const TYPE_RESPONSE: u8 = 2;
 pub const TYPE_COOKIEREPLY: u8 = 3;
 
+/* Handshake messsages */
+
 #[repr(C)]
 #[derive(Copy, Clone, FromBytes, AsBytes)]
-pub struct Initiation {
-    f_type: U32<LittleEndian>,
-    pub f_sender: U32<LittleEndian>,
-    pub f_ephemeral: [u8; SIZE_X25519_POINT],
-    pub f_static: [u8; SIZE_X25519_POINT],
-    pub f_static_tag: [u8; SIZE_TAG],
-    pub f_timestamp: timestamp::TAI64N,
-    pub f_timestamp_tag: [u8; SIZE_TAG],
-    pub f_macs: MacsFooter,
+pub struct Response {
+    pub noise: NoiseResponse, // inner message covered by macs
+    pub macs: MacsFooter,
 }
 
 #[repr(C)]
 #[derive(Copy, Clone, FromBytes, AsBytes)]
-pub struct Response {
-    f_type: U32<LittleEndian>,
-    pub f_sender: U32<LittleEndian>,
-    pub f_receiver: U32<LittleEndian>,
-    pub f_ephemeral: [u8; SIZE_X25519_POINT],
-    pub f_empty_tag: [u8; SIZE_TAG],
-    pub f_macs: MacsFooter,
+pub struct Initiation {
+    pub noise: NoiseInitiation, // inner message covered by macs
+    pub macs: MacsFooter,
 }
 
 #[repr(C)]
@@ -55,28 +47,45 @@ pub struct CookieReply {
     pub f_cookie_tag: [u8; SIZE_TAG],
 }
 
-impl Default for Initiation {
-    fn default() -> Self {
-        Self {
-            f_type: <U32<LittleEndian>>::new(TYPE_INITIATION as u32),
+/* Inner sub-messages */
 
-            f_sender: <U32<LittleEndian>>::ZERO,
-            f_ephemeral: [0u8; SIZE_X25519_POINT],
-            f_static: [0u8; SIZE_X25519_POINT],
-            f_static_tag: [0u8; SIZE_TAG],
-            f_timestamp: timestamp::ZERO,
-            f_timestamp_tag: [0u8; SIZE_TAG],
-            f_macs: Default::default(),
-        }
-    }
+#[repr(C)]
+#[derive(Copy, Clone, FromBytes, AsBytes)]
+pub struct MacsFooter {
+    pub f_mac1: [u8; SIZE_MAC],
+    pub f_mac2: [u8; SIZE_MAC],
 }
 
+#[repr(C)]
+#[derive(Copy, Clone, FromBytes, AsBytes)]
+pub struct NoiseInitiation {
+    f_type: U32<LittleEndian>,
+    pub f_sender: U32<LittleEndian>,
+    pub f_ephemeral: [u8; SIZE_X25519_POINT],
+    pub f_static: [u8; SIZE_X25519_POINT],
+    pub f_static_tag: [u8; SIZE_TAG],
+    pub f_timestamp: timestamp::TAI64N,
+    pub f_timestamp_tag: [u8; SIZE_TAG]
+}
+
+#[repr(C)]
+#[derive(Copy, Clone, FromBytes, AsBytes)]
+pub struct NoiseResponse {
+    f_type: U32<LittleEndian>,
+    pub f_sender: U32<LittleEndian>,
+    pub f_receiver: U32<LittleEndian>,
+    pub f_ephemeral: [u8; SIZE_X25519_POINT],
+    pub f_empty_tag: [u8; SIZE_TAG]
+}
+
+/* Zero copy parsing of handshake messages */
+
 impl Initiation {
     pub fn parse<B: ByteSlice>(bytes: B) -> Result<LayoutVerified<B, Self>, HandshakeError> {
         let msg: LayoutVerified<B, Self> =
             LayoutVerified::new(bytes).ok_or(HandshakeError::InvalidMessageFormat)?;
 
-        if msg.f_type.get() != (TYPE_INITIATION as u32) {
+        if msg.noise.f_type.get() != (TYPE_INITIATION as u32) {
             return Err(HandshakeError::InvalidMessageFormat);
         }
 
@@ -84,44 +93,25 @@ impl Initiation {
     }
 }
 
-#[cfg(test)]
-impl fmt::Debug for Initiation {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f,
-            "MessageInitiation {{ type = {}, sender = {}, ephemeral = {}, static = {}|{}, timestamp = {}|{} }}",
-            self.f_type.get(),
-            self.f_sender.get(),
-            hex::encode(self.f_ephemeral),
-            hex::encode(self.f_static),
-            hex::encode(self.f_static_tag),
-            hex::encode(self.f_timestamp),
-            hex::encode(self.f_timestamp_tag)
-        )
-    }
-}
+impl Response {
+    pub fn parse<B: ByteSlice>(bytes: B) -> Result<LayoutVerified<B, Self>, HandshakeError> {
+        let msg: LayoutVerified<B, Self> =
+            LayoutVerified::new(bytes).ok_or(HandshakeError::InvalidMessageFormat)?;
 
-#[cfg(test)]
-impl PartialEq for Initiation {
-    fn eq(&self, other: &Self) -> bool {
-        self.f_type.get() == other.f_type.get()
-            && self.f_sender.get() == other.f_sender.get()
-            && self.f_ephemeral[..] == other.f_ephemeral[..]
-            && self.f_static[..] == other.f_static[..]
-            && self.f_static_tag[..] == other.f_static_tag[..]
-            && self.f_timestamp[..] == other.f_timestamp
-            && self.f_timestamp_tag[..] == other.f_timestamp_tag
+        if msg.noise.f_type.get() != (TYPE_RESPONSE as u32) {
+            return Err(HandshakeError::InvalidMessageFormat);
+        }
+
+        Ok(msg)
     }
 }
 
-#[cfg(test)]
-impl Eq for Initiation {}
-
-impl Response {
+impl CookieReply {
     pub fn parse<B: ByteSlice>(bytes: B) -> Result<LayoutVerified<B, Self>, HandshakeError> {
         let msg: LayoutVerified<B, Self> =
             LayoutVerified::new(bytes).ok_or(HandshakeError::InvalidMessageFormat)?;
 
-        if msg.f_type.get() != (TYPE_RESPONSE as u32) {
+        if msg.f_type.get() != (TYPE_COOKIEREPLY as u32) {
             return Err(HandshakeError::InvalidMessageFormat);
         }
 
@@ -129,24 +119,125 @@ impl Response {
     }
 }
 
+/* Default values */
+
 impl Default for Response {
     fn default() -> Self {
         Self {
+        noise: Default::default(),
+        macs: Default::default(),
+        }
+    }
+}
+
+impl Default for Initiation {
+    fn default() -> Self {
+        Self {
+        noise: Default::default(),
+        macs: Default::default(),
+        }
+    }
+}
+
+impl Default for CookieReply {
+    fn default() -> Self {
+        Self {
+        f_type: <U32<LittleEndian>>::new(TYPE_COOKIEREPLY as u32),
+        f_receiver: <U32<LittleEndian>>::ZERO,
+        f_nonce: [0u8; SIZE_NONCE],
+        f_cookie: [0u8; SIZE_COOKIE],
+        f_cookie_tag: [0u8; SIZE_TAG],
+        }
+    }
+}
+
+impl Default for MacsFooter {
+    fn default() -> Self {
+        Self {
+            f_mac1: [0u8; SIZE_MAC],
+            f_mac2: [0u8; SIZE_MAC],
+        }
+    }
+}
+
+impl Default for NoiseInitiation {
+      fn default() -> Self {
+        Self {
+            f_type: <U32<LittleEndian>>::new(TYPE_INITIATION as u32),
+
+            f_sender: <U32<LittleEndian>>::ZERO,
+            f_ephemeral: [0u8; SIZE_X25519_POINT],
+            f_static: [0u8; SIZE_X25519_POINT],
+            f_static_tag: [0u8; SIZE_TAG],
+            f_timestamp: timestamp::ZERO,
+            f_timestamp_tag: [0u8; SIZE_TAG]
+        }
+    }
+}
+
+impl Default for NoiseResponse {
+    fn default() -> Self {
+        Self {
             f_type: <U32<LittleEndian>>::new(TYPE_RESPONSE as u32),
             f_sender: <U32<LittleEndian>>::ZERO,
             f_receiver: <U32<LittleEndian>>::ZERO,
             f_ephemeral: [0u8; SIZE_X25519_POINT],
             f_empty_tag: [0u8; SIZE_TAG],
-            f_macs: Default::default(),
         }
     }
 }
 
+/* Debug formatting (for testing purposes) */
+
+#[cfg(test)]
+impl fmt::Debug for Initiation {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "Initiation {{ {:?} || {:?} }}", self.noise, self.macs)
+    }
+}
+
 #[cfg(test)]
 impl fmt::Debug for Response {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "Response {{ {:?} || {:?} }}", self.noise, self.macs)
+    }
+}
+
+#[cfg(test)]
+impl fmt::Debug for CookieReply {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f,
-            "MessageResponse {{ type = {}, sender = {}, receiver = {}, ephemeral = {}, empty = |{}  }}",
+            "CookieReply {{ type = {}, receiver = {}, nonce = {}, cookie = {}|{}  }}",
+            self.f_type,
+            self.f_receiver,
+            hex::encode(self.f_nonce),
+            hex::encode(self.f_cookie),
+            hex::encode(self.f_cookie_tag)
+        )
+    }
+}
+
+#[cfg(test)]
+impl fmt::Debug for NoiseInitiation {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f,
+            "NoiseInitiation {{ type = {}, sender = {}, ephemeral = {}, static = {}|{}, timestamp = {}|{} }}",
+            self.f_type.get(),
+            self.f_sender.get(),
+            hex::encode(self.f_ephemeral),
+            hex::encode(self.f_static),
+            hex::encode(self.f_static_tag),
+            hex::encode(self.f_timestamp),
+            hex::encode(self.f_timestamp_tag)
+        )
+    }
+}
+
+#[cfg(test)]
+impl fmt::Debug for NoiseResponse {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f,
+            "NoiseResponse {{ type = {}, sender = {}, receiver = {}, ephemeral = {}, empty = |{}  }}",
             self.f_type,
             self.f_sender,
             self.f_receiver,
@@ -157,17 +248,48 @@ impl fmt::Debug for Response {
 }
 
 #[cfg(test)]
-impl PartialEq for Response {
-    fn eq(&self, other: &Self) -> bool {
-        self.f_type == other.f_type
-            && self.f_sender == other.f_sender
-            && self.f_receiver == other.f_receiver
-            && self.f_ephemeral == other.f_ephemeral
-            && self.f_empty_tag == other.f_empty_tag
+impl fmt::Debug for MacsFooter {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f,
+            "Macs {{ mac1 = {}, mac2 = {} }}",
+            hex::encode(self.f_mac1),
+            hex::encode(self.f_mac2)
+        )
+    }
+}
+
+/* Equality (for testing purposes) */
+
+macro_rules! eq_as_bytes {
+    ($type:path) => {
+        impl PartialEq for $type {
+            fn eq(&self, other: &Self) -> bool {
+                self.as_bytes() == other.as_bytes()
+            }
+        }
+        impl Eq for $type {}
     }
 }
 
 #[cfg(test)]
+eq_as_bytes!(Initiation);
+
+#[cfg(test)]
+eq_as_bytes!(Response);
+
+#[cfg(test)]
+eq_as_bytes!(CookieReply);
+
+#[cfg(test)]
+eq_as_bytes!(MacsFooter);
+
+#[cfg(test)]
+eq_as_bytes!(NoiseInitiation);
+
+#[cfg(test)]
+eq_as_bytes!(NoiseResponse);
+
+#[cfg(test)]
 mod tests {
     use super::*;
 
@@ -175,22 +297,22 @@ mod tests {
     fn message_response_identity() {
         let mut msg: Response = Default::default();
 
-        msg.f_sender.set(146252);
-        msg.f_receiver.set(554442);
-        msg.f_ephemeral = [
+        msg.noise.f_sender.set(146252);
+        msg.noise.f_receiver.set(554442);
+        msg.noise.f_ephemeral = [
             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,
         ];
-        msg.f_empty_tag = [
+        msg.noise.f_empty_tag = [
             0x60, 0x0e, 0x1e, 0x95, 0x41, 0x6b, 0x52, 0x05, 0xa2, 0x09, 0xe1, 0xbf, 0x40, 0x05,
             0x2f, 0xde,
         ];
-        msg.f_macs.f_mac1 = [
+        msg.macs.f_mac1 = [
             0xf2, 0xad, 0x40, 0xb5, 0xf7, 0xde, 0x77, 0x35, 0x89, 0x19, 0xb7, 0x5c, 0xf9, 0x54,
             0x69, 0x29,
         ];
-        msg.f_macs.f_mac2 = [
+        msg.macs.f_mac2 = [
             0x4f, 0xd2, 0x1b, 0xfe, 0x77, 0xe6, 0x2e, 0xc9, 0x07, 0xe2, 0x87, 0x17, 0xbb, 0xe5,
             0xdf, 0xbb,
         ];
@@ -204,33 +326,33 @@ mod tests {
     fn message_initiate_identity() {
         let mut msg: Initiation = Default::default();
 
-        msg.f_sender.set(575757);
-        msg.f_ephemeral = [
+        msg.noise.f_sender.set(575757);
+        msg.noise.f_ephemeral = [
             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,
         ];
-        msg.f_static = [
+        msg.noise.f_static = [
             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,
         ];
-        msg.f_static_tag = [
+        msg.noise.f_static_tag = [
             0x45, 0xac, 0x8d, 0x43, 0xea, 0x1b, 0x2f, 0x02, 0x45, 0x5d, 0x86, 0x37, 0xee, 0x83,
             0x6b, 0x42,
         ];
-        msg.f_timestamp = [
+        msg.noise.f_timestamp = [
             0x4f, 0x1c, 0x60, 0xec, 0x0e, 0xf6, 0x36, 0xf0, 0x78, 0x28, 0x57, 0x42,
         ];
-        msg.f_timestamp_tag = [
+        msg.noise.f_timestamp_tag = [
             0x60, 0x0e, 0x1e, 0x95, 0x41, 0x6b, 0x52, 0x05, 0xa2, 0x09, 0xe1, 0xbf, 0x40, 0x05,
             0x2f, 0xde,
         ];
-        msg.f_macs.f_mac1 = [
+        msg.macs.f_mac1 = [
             0xf2, 0xad, 0x40, 0xb5, 0xf7, 0xde, 0x77, 0x35, 0x89, 0x19, 0xb7, 0x5c, 0xf9, 0x54,
             0x69, 0x29,
         ];
-        msg.f_macs.f_mac2 = [
+        msg.macs.f_mac2 = [
             0x4f, 0xd2, 0x1b, 0xfe, 0x77, 0xe6, 0x2e, 0xc9, 0x07, 0xe2, 0x87, 0x17, 0xbb, 0xe5,
             0xdf, 0xbb,
         ];
diff --git a/src/handshake/noise.rs b/src/handshake/noise.rs
index a03cae2..72e0f47 100644
--- a/src/handshake/noise.rs
+++ b/src/handshake/noise.rs
@@ -18,7 +18,7 @@ use generic_array::GenericArray;
 use zerocopy::AsBytes;
 
 use super::device::Device;
-use super::messages::{Initiation, Response};
+use super::messages::{NoiseInitiation, NoiseResponse};
 use super::peer::{Peer, State};
 use super::timestamp;
 use super::types::*;
@@ -142,7 +142,7 @@ pub fn create_initiation<T: Copy>(
     sender: u32,
 ) -> Result<Vec<u8>, HandshakeError> {
     let mut rng = OsRng::new().unwrap();
-    let mut msg: Initiation = Default::default();
+    let mut msg: NoiseInitiation = Default::default();
 
     // initialize state
 
@@ -221,12 +221,8 @@ pub fn create_initiation<T: Copy>(
 
 pub fn consume_initiation<'a, T: Copy>(
     device: &'a Device<T>,
-    msg: &[u8],
+    msg: &NoiseInitiation,
 ) -> Result<(&'a Peer<T>, TemporaryState), HandshakeError> {
-    // parse message
-
-    let msg = Initiation::parse(msg)?;
-
     // initialize state
 
     let ck = INITIAL_CK;
@@ -295,12 +291,11 @@ pub fn consume_initiation<'a, T: Copy>(
 
 pub fn create_response<T: Copy>(
     peer: &Peer<T>,
-    sender: u32,           // sending identifier
-    state: TemporaryState, // state from "consume_initiation"
+    sender: u32,             // sending identifier
+    state: TemporaryState,   // state from "consume_initiation"
+    msg: &mut NoiseResponse, // resulting response
 ) -> Result<Output<T>, HandshakeError> {
     let mut rng = OsRng::new().unwrap();
-    let mut msg: Response = Default::default();
-
     let (receiver, eph_r_pk, hs, ck) = state;
 
     msg.f_sender.set(sender);
@@ -380,12 +375,8 @@ pub fn create_response<T: Copy>(
 
 pub fn consume_response<T: Copy>(
     device: &Device<T>,
-    msg: &[u8],
+    msg: &NoiseResponse,
 ) -> Result<Output<T>, HandshakeError> {
-    // parse message
-
-    let msg = Response::parse(msg)?;
-
     // retrieve peer and associated state
 
     let peer = device.lookup_id(msg.f_receiver.get())?;