1 files changed, 206 insertions, 84 deletions

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,
         ];