diff --git a/protocol/src/lib.rs b/protocol/src/lib.rs index 2e3d248..36c82ef 100644 --- a/protocol/src/lib.rs +++ b/protocol/src/lib.rs @@ -168,7 +168,7 @@ pub struct SessionKeyMaterial { } /// Your role in the handshake. -#[derive(Clone, Debug)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum Role { /// You started the handshake with a peer. Initiator, @@ -176,24 +176,24 @@ pub enum Role { Responder, } -/// A message or decoy packet from a connected peer. -#[cfg(feature = "alloc")] -#[derive(Clone, Debug)] -pub struct ReceivedMessage { - /// A message to handle or `None` if the peer sent a decoy and the message may be safely ignored. - pub message: Option>, +/// A decoy packet contains bogus information, but can be +/// used to hide the shape of the data being communicated +/// over an encrypted channel. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum PacketType { + /// Genuine packet contains information. + Genuine, + /// Docoy packet contians bogus information. + Decoy, } -#[cfg(feature = "alloc")] -impl ReceivedMessage { - pub fn new(msg_bytes: &[u8]) -> Result { - let header = msg_bytes.first().ok_or(Error::CiphertextTooSmall)?; - if header.eq(&DECOY_BYTE) { - Ok(ReceivedMessage { message: None }) +impl PacketType { + /// Check if plaintext packet packets are a decoy. + pub fn from_bytes(plaintext: &[u8]) -> Self { + if plaintext.first() == Some(&DECOY_BYTE) { + PacketType::Decoy } else { - Ok(ReceivedMessage { - message: Some(msg_bytes[1..].to_vec()), - }) + PacketType::Genuine } } } @@ -240,15 +240,22 @@ impl PacketReader { /// - `contents` - Mutable buffer to write plaintext. /// - `aad` - Optional authentication for the peer. /// + /// # Returns + /// + /// A `Result` containing: + /// - `Ok(PacketType)`: A flag indicating if the decoded packet is a decoy or not. + /// - `Err(Error)`: An error that occurred during decryption. + /// /// # Errors /// - /// Fails if the packet was not decrypted or authenticated properly. + /// - `CiphertextTooSmall` - Ciphertext argument does not contain a whole packet. + /// - `BufferTooSmall ` - Contents buffer argument is not large enough for plaintext. pub fn decrypt_contents( &mut self, ciphertext: &[u8], contents: &mut [u8], aad: Option<&[u8]>, - ) -> Result<(), Error> { + ) -> Result { let auth = aad.unwrap_or_default(); // Check minimum size of ciphertext. if ciphertext.len() < TAG_BYTES { @@ -268,7 +275,7 @@ impl PacketReader { tag.try_into().expect("16 byte tag"), )?; - Ok(()) + Ok(PacketType::from_bytes(contents)) } /// Decrypt the rest of the message from the peer, excluding the 3 length bytes. This method should only be called after @@ -279,18 +286,30 @@ impl PacketReader { /// - `ciphertext` - The message from the peer. /// - `aad` - Optional authentication for the peer. /// + /// # Returns + /// + /// A `Result` containing: + /// - `Ok(Some(Vec))`: The plaintext in a byte vector if it is not a decoy packet. + /// - `Err(Error)`: An error that occurred during decryption. + /// /// # Errors /// - /// Fails if the packet was not decrypted or authenticated properly. + /// - `CiphertextTooSmall` - Ciphertext argument does not contain a whole packet. #[cfg(feature = "alloc")] pub fn decrypt_contents_with_alloc( &mut self, ciphertext: &[u8], aad: Option<&[u8]>, - ) -> Result, Error> { + ) -> Result>, Error> { let mut contents = vec![0u8; ciphertext.len() - TAG_BYTES]; - self.decrypt_contents(ciphertext, &mut contents, aad)?; - Ok(contents) + match self.decrypt_contents(ciphertext, &mut contents, aad)? { + PacketType::Decoy => Ok(None), + PacketType::Genuine => { + // Drop the decoy byte flag. + contents.remove(0); + Ok(Some(contents)) + } + } } } @@ -492,14 +511,9 @@ impl PacketHandler { &mut self, ciphertext: &[u8], aad: Option<&[u8]>, - ) -> Result { - let contents = self - .packet_reader - .decrypt_contents_with_alloc(ciphertext, aad)?; - - let message = ReceivedMessage::new(&contents)?; - - Ok(message) + ) -> Result>, Error> { + self.packet_reader + .decrypt_contents_with_alloc(ciphertext, aad) } } @@ -722,7 +736,7 @@ impl<'a> Handshake<'a> { } }; - let mut packet_handler = PacketHandler::new(materials, self.role.clone()); + let mut packet_handler = PacketHandler::new(materials, self.role); // TODO: Support sending decoy packets before the version packet. @@ -1063,7 +1077,7 @@ mod tests { let dec = bob_packet_handler .decrypt_contents_with_alloc(&enc_packet[LENGTH_BYTES..], None) .unwrap(); - assert_eq!(None, dec.message); + assert_eq!(None, dec); let message = b"Windows sox!".to_vec(); let enc_packet = bob_packet_handler .prepare_packet_with_alloc(&message, None, false) @@ -1071,7 +1085,7 @@ mod tests { let dec = alice_packet_handler .decrypt_contents_with_alloc(&enc_packet[LENGTH_BYTES..], None) .unwrap(); - assert_eq!(message, dec.message.unwrap()); + assert_eq!(message, dec.unwrap()); } #[test] @@ -1102,7 +1116,7 @@ mod tests { let dec_packet = bob_packet_handler .decrypt_contents_with_alloc(&enc_packet[LENGTH_BYTES..], None) .unwrap(); - assert_eq!(message, dec_packet.message.unwrap()); + assert_eq!(message, dec_packet.unwrap()); let message = gen_garbage(420, &mut rng); let enc_packet = bob_packet_handler .prepare_packet_with_alloc(&message, None, false) @@ -1110,7 +1124,7 @@ mod tests { let dec_packet = alice_packet_handler .decrypt_contents_with_alloc(&enc_packet[LENGTH_BYTES..], None) .unwrap(); - assert_eq!(message, dec_packet.message.unwrap()); + assert_eq!(message, dec_packet.unwrap()); } } @@ -1193,7 +1207,7 @@ mod tests { let dec = alice .decrypt_contents_with_alloc(&encrypted_message_to_alice[LENGTH_BYTES..], None) .unwrap(); - assert_eq!(message, dec.message.unwrap()); + assert_eq!(message, dec.unwrap()); let message = b"g!".to_vec(); let encrypted_message_to_bob = alice .prepare_packet_with_alloc(&message, None, false) @@ -1201,7 +1215,7 @@ mod tests { let dec = bob .decrypt_contents_with_alloc(&encrypted_message_to_bob[LENGTH_BYTES..], None) .unwrap(); - assert_eq!(message, dec.message.unwrap()); + assert_eq!(message, dec.unwrap()); } #[test] @@ -1248,11 +1262,10 @@ mod tests { let contents = bob .decrypt_contents_with_alloc(&message_to_bob[3..3 + alice_message_len], None) .unwrap(); - assert_eq!(contents.message.unwrap(), message); + assert_eq!(contents.unwrap(), message); } // The rest are sourced from [the BIP324 test vectors](https://github.com/bitcoin/bips/blob/master/bip-0324/packet_encoding_test_vectors.csv). - // #[test] #[cfg(feature = "std")] @@ -1280,7 +1293,7 @@ mod tests { let dec_packet = bob_packet_handler .decrypt_contents_with_alloc(&enc[LENGTH_BYTES..], None) .unwrap(); - assert_eq!(first, dec_packet.message.unwrap()); + assert_eq!(first, dec_packet.unwrap()); let contents: Vec = vec![0x8e]; let enc = alice_packet_handler .prepare_packet_with_alloc(&contents, None, false) diff --git a/protocol/tests/round_trips.rs b/protocol/tests/round_trips.rs index 9cf60f8..51ed95c 100644 --- a/protocol/tests/round_trips.rs +++ b/protocol/tests/round_trips.rs @@ -48,7 +48,7 @@ fn hello_world_happy_path() { let messages = alice .decrypt_contents_with_alloc(&encrypted_message_to_alice[3..], None) .unwrap(); - assert_eq!(message, messages.message.unwrap()); + assert_eq!(message, messages.unwrap()); let message = b"Goodbye!".to_vec(); let encrypted_message_to_bob = alice .prepare_packet_with_alloc(&message, None, false) @@ -56,7 +56,7 @@ fn hello_world_happy_path() { let messages = bob .decrypt_contents_with_alloc(&encrypted_message_to_bob[3..], None) .unwrap(); - assert_eq!(message, messages.message.unwrap()); + assert_eq!(message, messages.unwrap()); } #[test] @@ -71,7 +71,7 @@ fn regtest_handshake() { use bip324::{ serde::{deserialize, serialize, NetworkMessage}, - Handshake, ReceivedMessage, + Handshake, }; use bitcoincore_rpc::{ bitcoin::p2p::{message_network::VersionMessage, Address, ServiceFlags}, @@ -148,8 +148,7 @@ fn regtest_handshake() { let msg = decrypter .decrypt_contents_with_alloc(&response_message, None) .unwrap(); - let message = ReceivedMessage::new(&msg.clone()).unwrap(); - let message = deserialize(&message.message.unwrap()).unwrap(); + let message = deserialize(&msg.unwrap()).unwrap(); dbg!("{}", message.cmd()); assert_eq!(message.cmd(), "version"); rpc.stop().unwrap(); diff --git a/proxy/src/lib.rs b/proxy/src/lib.rs index 69925a4..bb75559 100644 --- a/proxy/src/lib.rs +++ b/proxy/src/lib.rs @@ -12,7 +12,6 @@ use std::fmt; use std::net::SocketAddr; use bip324::serde::{deserialize, serialize}; -use bip324::ReceivedMessage; use bip324::{PacketReader, PacketWriter}; use bitcoin::consensus::{Decodable, Encodable}; use bitcoin::p2p::message::{NetworkMessage, RawNetworkMessage}; @@ -142,10 +141,7 @@ pub async fn read_v2( .decrypt_contents_with_alloc(&packet_bytes, None) .expect("decrypt"); - let contents = ReceivedMessage::new(&raw) - .expect("some bytes") - .message - .expect("not a decoy"); + let contents = raw.expect("not a decoy"); let message = deserialize(&contents).map_err(|_| Error::Serde)?; Ok(message)