Add back garbage tests

protocol/src/lib.rs

@@ -67,7 +67,7 @@ pub const NUM_INITIAL_HANDSHAKE_BUFFER_BYTES: usize = 4096;
 const VERSION_CONTENT: [u8; 0] = [];
 // Number of bytes for the authentication tag of a packet.
 const NUM_TAG_BYTES: usize = 16;
-// Maximum number of garbage bytes to read before the terminator.
+// Maximum number of garbage bytes before the terminator.
 const MAX_NUM_GARBAGE_BYTES: usize = 4095;
 // Number of bytes for the garbage terminator.
 const NUM_GARBAGE_TERMINTOR_BYTES: usize = 16;
@@ -85,8 +85,11 @@ pub enum Error {
     /// total required bytes for the failed packet so the
     /// caller can re-allocate and re-attempt.
     BufferTooSmall { required_bytes: usize },
-    /// The maximum amount of garbage bytes was exceeded in the handshake.
-    MaxGarbageLength,
+    /// Tried to send more garbage bytes before terminator than allowed by spec.
+    TooMuchGarbage,
+    /// The remote sent the maximum amount of garbage bytes without
+    /// a garbage terminator in the handshake.
+    NoGarbageTerminator,
     /// A handshake step was not completed in the proper order.
     HandshakeOutOfOrder,
     /// Not able to generate secret material.
@@ -109,12 +112,16 @@ impl fmt::Display for Error {
                 "Buffer memory allocation too small, need at least {} bytes.",
                 required_bytes
             ),
-            Error::MaxGarbageLength => {
-                write!(f, "More than 4095 bytes of garbage in the handshake.")
+            Error::NoGarbageTerminator => {
+                write!(f, "More than 4095 bytes of garbage recieved in the handshake before a terminator was sent.")
             }
             Error::HandshakeOutOfOrder => write!(f, "Handshake flow out of sequence."),
             Error::SecretGeneration(e) => write!(f, "Cannot generate secrets: {:?}.", e),
             Error::Decryption(e) => write!(f, "Decrytion error: {:?}.", e),
+            Error::TooMuchGarbage => write!(
+                f,
+                "Tried to send more than 4095 bytes of garbage in handshake."
+            ),
         }
     }
 }
@@ -624,6 +631,13 @@ impl<'a> Handshake<'a> {
         rng: &mut impl Rng,
         curve: &Secp256k1<C>,
     ) -> Result<Self, Error> {
+        if garbage
+            .as_ref()
+            .map_or(false, |g| g.len() > MAX_NUM_GARBAGE_BYTES)
+        {
+            return Err(Error::TooMuchGarbage);
+        };
+
         let mut secret_key_buffer = [0u8; 32];
         rng.fill(&mut secret_key_buffer[..]);
         let sk = SecretKey::from_slice(&secret_key_buffer)?;
@@ -900,7 +914,7 @@ impl<'a> Handshake<'a> {
         {
             Ok((&buffer[..index], &buffer[(index + garbage_term.len())..]))
         } else if buffer.len() >= (MAX_NUM_GARBAGE_BYTES + NUM_GARBAGE_TERMINTOR_BYTES) {
-            Err(Error::MaxGarbageLength)
+            Err(Error::NoGarbageTerminator)
         } else {
             // Terminator not found, the buffer needs more information.
             Err(Error::CiphertextTooSmall)
@@ -1432,7 +1446,81 @@ mod tests {
 
     #[test]
     #[cfg(feature = "std")]
-    fn test_full_handshake_with_garbage_and_decoys() {
+    fn test_handshake_garbage_length_check() {
+        let mut rng = rand::thread_rng();
+        let curve = Secp256k1::new();
+        let mut handshake_buffer = [0u8; NUM_ELLIGATOR_SWIFT_BYTES + MAX_NUM_GARBAGE_BYTES];
+
+        // Test with valid garbage length.
+        let valid_garbage = vec![0u8; MAX_NUM_GARBAGE_BYTES];
+        let result = Handshake::new_with_rng(
+            Network::Bitcoin,
+            Role::Initiator,
+            Some(&valid_garbage),
+            &mut handshake_buffer,
+            &mut rng,
+            &curve,
+        );
+        assert!(result.is_ok());
+
+        // Test with garbage length exceeding MAX_NUM_GARBAGE_BYTES.
+        let invalid_garbage = vec![0u8; MAX_NUM_GARBAGE_BYTES + 1];
+        let result = Handshake::new_with_rng(
+            Network::Bitcoin,
+            Role::Initiator,
+            Some(&invalid_garbage),
+            &mut handshake_buffer,
+            &mut rng,
+            &curve,
+        );
+        assert!(matches!(result, Err(Error::TooMuchGarbage)));
+
+        // Test with no garbage.
+        let result = Handshake::new_with_rng(
+            Network::Bitcoin,
+            Role::Initiator,
+            None,
+            &mut handshake_buffer,
+            &mut rng,
+            &curve,
+        );
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    #[cfg(feature = "std")]
+    fn test_handshake_no_garbage_terminator() {
+        let mut handshake_buffer = [0u8; NUM_ELLIGATOR_SWIFT_BYTES];
+        let mut rng = rand::thread_rng();
+        let curve = Secp256k1::signing_only();
+
+        let mut handshake = Handshake::new_with_rng(
+            Network::Bitcoin,
+            Role::Initiator,
+            None,
+            &mut handshake_buffer,
+            &mut rng,
+            &curve,
+        )
+        .expect("Handshake creation should succeed");
+
+        // Skipping material creation and just placing a mock terminator.
+        handshake.remote_garbage_terminator = Some([0xFF; NUM_GARBAGE_TERMINTOR_BYTES]);
+
+        // Test with a buffer that is too long.
+        let test_buffer = vec![0; MAX_NUM_GARBAGE_BYTES + NUM_GARBAGE_TERMINTOR_BYTES];
+        let result = handshake.split_garbage(&test_buffer);
+        assert!(matches!(result, Err(Error::NoGarbageTerminator)));
+
+        // Test with a buffer that's just short of the required length.
+        let short_buffer = vec![0; MAX_NUM_GARBAGE_BYTES + NUM_GARBAGE_TERMINTOR_BYTES - 1];
+        let result = handshake.split_garbage(&short_buffer);
+        assert!(matches!(result, Err(Error::CiphertextTooSmall)));
+    }
+
+    #[test]
+    #[cfg(feature = "std")]
+    fn test_handshake_with_garbage_and_decoys() {
         // Define the garbage and decoys that the initiator is sending to the responder.
         let initiator_garbage = vec![1u8, 2u8, 3u8];
         let initiator_decoys: Vec<&[u8]> = vec![&[6u8, 7u8], &[8u8, 0u8]];