Handle decoy packets in handshake

protocol/src/lib.rs

@@ -53,6 +53,14 @@ const TAG_BYTES: usize = 16;
 const MAX_GARBAGE_BYTES: usize = 4095;
 // Number of bytes for the garbage terminator.
 const GARBAGE_TERMINTOR_BYTES: usize = 16;
+// The max number of bytes a packet can be in the handshake. This allows
+// the buffer to be pre-allocated and for things to remain sans-io. The
+// only require packet in the hankshake is the small version packet,
+// however, any number of any sized decoy packets can optionally
+// be sent. These packets shouldn't be crazy big because the point
+// is to resemble actual communication, but we need to cap the size
+// somewhere to allow for pre-allocation.
+const MAX_HANDSHAKE_PACKET_BYTES: usize = 4096;
 
 /// Errors encountered throughout the lifetime of a V2 connection.
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
@@ -219,7 +227,7 @@ impl PacketReader {
     ///
     /// - `ciphertext` - The message from the peer.
     /// - `contents`   - Mutable buffer to write plaintext.
-    /// - `aad`        - Optional authentication for the peer, currently only used for the first round of messages.
+    /// - `aad`        - Optional authentication for the peer.
     ///
     /// # Errors
     ///
@@ -248,7 +256,7 @@ impl PacketReader {
     /// # Arguments
     ///
     /// - `ciphertext` - The message from the peer.
-    /// - `aad`        - Optional authentication for the peer, currently only used for the first round of messages.
+    /// - `aad`        - Optional authentication for the peer.
     ///
     /// # Errors
     ///
@@ -564,8 +572,10 @@ pub struct Handshake<'a> {
     remote_garbage_terminator: Option<[u8; 16]>,
     /// Packet handler output.
     packet_handler: Option<PacketHandler>,
-    /// Stored state between authentication attempts, decrypted length for next packet.
-    authentication_packet_bytes: Option<usize>,
+    /// Decrypted length for next packet. Store state between authentication attempts to avoid resetting ciphers.
+    current_packet_length_bytes: Option<usize>,
+    /// Processesed message index. Store state between authentication attempts to avoid resetting ciphers.
+    current_message_index: usize,
 }
 
 impl<'a> Handshake<'a> {
@@ -640,7 +650,8 @@ impl<'a> Handshake<'a> {
             garbage,
             remote_garbage_terminator: None,
             packet_handler: None,
-            authentication_packet_bytes: None,
+            current_packet_length_bytes: None,
+            current_message_index: 0,
         })
     }
 
@@ -708,69 +719,131 @@ impl<'a> Handshake<'a> {
     /// Authenticate the channel.
     ///
     /// Designed to be called multiple times until succesful in order to flush
-    /// garbage and decoy packets from channel.
+    /// garbage and decoy packets from channel. If a `MessageLengthTooSmall` is
+    /// returned, the buffer should be extended until `MessageLengthTooSmall` is
+    /// not returned. All other errors are fatal for the handshake and it should
+    /// be completely restarted.
     ///
     /// # Arguments
     ///
-    /// - `buffer` - Should contain all garbage, the garbage terminator, and the version packet received from peer.
+    /// - `buffer` - Should contain all garbage, the garbage terminator, any decoy packets, and finally the version packet received from peer.
     ///
     /// # Error    
     ///
     /// - `MessageLengthTooSmall` - The buffer did not contain all required information and should be extended (e.g. read more off a socket) and authentication re-tried.
     /// - `HandshakeOutOfOrder`   - The handshake sequence is in a bad state and should be restarted.
     /// - `MaxGarbageLength`      - Buffer did not contain the garbage terminator and contains too much garbage, should not be retried.
     pub fn authenticate_garbage_and_version(&mut self, buffer: &[u8]) -> Result<(), Error> {
-        // Find the end of the garbage
+        // Find the end of the garbage.
         let (garbage, message) = split_garbage_and_message(
             buffer,
             self.remote_garbage_terminator
                 .ok_or(Error::HandshakeOutOfOrder)?,
         )?;
 
-        // Quickly fail if the message doesn't even have enough bytes for a length packet.
-        if message.len() < LENGTH_BYTES {
-            return Err(Error::MessageLengthTooSmall);
-        }
-
         let packet_handler = self
             .packet_handler
             .as_mut()
             .ok_or(Error::HandshakeOutOfOrder)?;
 
-        // TODO: Drain decoy packets, will require some more state to be store between attempts, like a message index.
+        // The buffer used to decode any decoy packets and the version packet.
+        let mut packet_buffer = [0u8; MAX_HANDSHAKE_PACKET_BYTES];
+        // Initialize the version packet to a decoy packet in order
+        // to keep flushing the decoy packets on follow up authentication attempts.
+        let mut version_packet: [u8; NUM_DECOY_BYTES + TAG_BYTES] =
+            [DECOY_BYTE, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
+
+        // The first packet, even if it is a decoy packet,
+        // is used to authenticate the received garbage through
+        // the AAD.
+        if self.current_message_index == 0 {
+            if self.current_packet_length_bytes.is_none() {
+                // Bounds check on the message buffer.
+                if message.len() < LENGTH_BYTES {
+                    return Err(Error::MessageLengthTooSmall);
+                }
+                let packet_length = packet_handler.decypt_len(
+                    message[0..LENGTH_BYTES]
+                        .try_into()
+                        .map_err(|_| Error::MessageLengthTooSmall)?,
+                );
+                // Hang on to decrypted length incase next steps fail to avoid using
+                // the cipher again and throwing off the state.
+                self.current_packet_length_bytes = Some(packet_length);
+            }
 
-        // Grab the packet length from internal statem, else decrypt it and store incase of failure.
-        let packet_length = if self.authentication_packet_bytes.is_some() {
-            self.authentication_packet_bytes
-                .ok_or(Error::HandshakeOutOfOrder)
-        } else {
-            let packet_length = packet_handler.decypt_len(
-                message[0..LENGTH_BYTES]
-                    .try_into()
-                    .map_err(|_| Error::MessageLengthTooSmall)?,
-            );
-            // Hang on to decrypted length incase next steps fail to avoid using the cipher again re-attempting authentication.
-            self.authentication_packet_bytes = Some(packet_length);
-            Ok(packet_length)
-        }?;
-
-        // Fail if there is not enough bytes to parse the message.
-        if message.len() < LENGTH_BYTES + packet_length {
-            return Err(Error::MessageLengthTooSmall);
+            let first_packet_length = self
+                .current_packet_length_bytes
+                .ok_or(Error::HandshakeOutOfOrder)?;
+
+            // Bounds check on the message buffer.
+            if message.len() < LENGTH_BYTES + first_packet_length {
+                return Err(Error::MessageLengthTooSmall);
+            }
+
+            // Authenticate received garbage with the first packet AAD.
+            packet_handler.packet_reader.decrypt_contents(
+                &message[LENGTH_BYTES..first_packet_length + LENGTH_BYTES],
+                &mut packet_buffer,
+                Some(garbage),
+            )?;
+
+            // Mark current decryption point in the buffer.
+            self.current_message_index = LENGTH_BYTES + first_packet_length + 1;
+            self.current_packet_length_bytes = None;
+
+            // Assume that the first packet is also the version packet,
+            // If it is a decoy packet then the version packet
+            // needs to be found and the decoys will be flushed below.
+            version_packet = packet_buffer[..NUM_DECOY_BYTES + TAG_BYTES]
+                .try_into()
+                .unwrap();
         }
 
-        // Authenticate received garbage and get version packet.
-        // Assuming no decoy packets so AAD is set on version packet.
-        // The version packet is ignored in this version of the protocol, but
-        // moves along state in the ciphers.
-
-        // Version packets have 0 contents.
-        let mut version_packet = [0u8; NUM_DECOY_BYTES + TAG_BYTES];
-        packet_handler.packet_reader.decrypt_contents(
-            &message[LENGTH_BYTES..packet_length + LENGTH_BYTES],
-            &mut version_packet,
-            Some(garbage),
-        )?;
+        // If the first packet is a decoy, or if this is a follow up
+        // authentication attempt, the decoys need to be flushed and
+        // the version packet found.
+        //
+        // The version packet is essentially ignored in the current
+        // version of the protocol, it just moves the cipher
+        // states forward, but could be extended in the future. It is
+        // just an empty packet in the current version.
+        while version_packet[0] == DECOY_BYTE {
+            if self.current_packet_length_bytes.is_none() {
+                if message.len() < self.current_message_index + LENGTH_BYTES {
+                    return Err(Error::MessageLengthTooSmall);
+                }
+                let packet_length = packet_handler.decypt_len(
+                    message[self.current_message_index..LENGTH_BYTES]
+                        .try_into()
+                        .map_err(|_| Error::MessageLengthTooSmall)?,
+                );
+                self.current_packet_length_bytes = Some(packet_length);
+            }
+
+            let next_packet_length = self
+                .current_packet_length_bytes
+                .ok_or(Error::HandshakeOutOfOrder)?;
+
+            if message.len() < self.current_message_index + LENGTH_BYTES + next_packet_length {
+                return Err(Error::MessageLengthTooSmall);
+            }
+            packet_handler.packet_reader.decrypt_contents(
+                &message[self.current_message_index + LENGTH_BYTES
+                    ..self.current_message_index + LENGTH_BYTES + next_packet_length],
+                &mut packet_buffer,
+                None,
+            )?;
+
+            // Mark current state.
+            self.current_message_index =
+                self.current_message_index + LENGTH_BYTES + next_packet_length + 1;
+            self.current_packet_length_bytes = None;
+
+            version_packet = packet_buffer[..NUM_DECOY_BYTES + TAG_BYTES]
+                .try_into()
+                .unwrap();
+        }
 
         Ok(())
     }