From fdb97db0ffb76ade9e0aea887e08d56582677d4a Mon Sep 17 00:00:00 2001 From: Nick Johnson Date: Thu, 3 Oct 2024 10:00:47 -0700 Subject: [PATCH] Add back garbage tests --- protocol/src/lib.rs | 102 +++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 95 insertions(+), 7 deletions(-) diff --git a/protocol/src/lib.rs b/protocol/src/lib.rs index 3cd1c41..7b3cb7b 100644 --- a/protocol/src/lib.rs +++ b/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, ) -> Result { + 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]];