Skip to content

Commit

Permalink
Get rid of a lot of message usage
Browse files Browse the repository at this point in the history
  • Loading branch information
@nyonson
nyonson committed Sep 25, 2024
1 parent 0a61881 commit 8c9832f
Showing 1 changed file with 26 additions and 30 deletions.
56 changes: 26 additions & 30 deletions protocol/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -79,8 +79,6 @@ pub enum Error {
/// total required bytes for the failed packet so the
/// caller can re-allocate and re-attempt.
BufferTooSmall { required_bytes: usize },
/// There is a mismatch in the encoding of a message.
IncompatableV1Message,
/// The maximum amount of garbage bytes was exceeded in the handshake.
MaxGarbageLength,
/// A handshake step was not completed in the proper order.
Expand All @@ -89,7 +87,7 @@ pub enum Error {
SecretMaterialsGeneration(secp256k1::Error),
/// Deriving the shared secrets was unsuccessful.
SecretExpansion,
/// The authentication data was not correct when decoding a message.
/// The authentication data was not correct when decoding a packet.
Cipher(fschacha20poly1305::Error),
/// The internal counters of the ciphers are not in sync.
OutOfSync,
Expand All @@ -112,7 +110,6 @@ impl fmt::Display for Error {
"Buffer memory allocation too small, need at least {} bytes.",
required_bytes
),
Error::IncompatableV1Message => write!(f, "Incompatable V1 message."),
Error::MaxGarbageLength => {
write!(f, "More than 4095 bytes of garbage in the handshake.")
}
Expand All @@ -130,7 +127,6 @@ impl std::error::Error for Error {
match self {
Error::SecretMaterialsGeneration(e) => Some(e),
Error::CiphertextTooSmall => None,
Error::IncompatableV1Message => None,
Error::MaxGarbageLength => None,
Error::HandshakeOutOfOrder => None,
Error::Cipher(e) => Some(e),
Expand Down Expand Up @@ -227,10 +223,10 @@ pub struct PacketReader {
}

impl PacketReader {
/// Decode the length, in bytes, of the rest of the inbound message.
/// Decode the length, in bytes, of the rest of the inbound packet.
///
/// Note that this does not decode to the length of contents described
/// in BIP324, and is meant to represent the entire inbound message
/// in BIP324, and is meant to represent the rest of the inbound packet
/// which includes the header byte and the 16-byte authentication tag.
///
/// # Arguments
Expand All @@ -255,7 +251,7 @@ impl PacketReader {
///
/// # Arguments
///
/// * `ciphertext` - The message from the peer excluding the first 3 length bytes. It should contain
/// * `ciphertext` - The packet from the peer excluding the first 3 length bytes. It should contain
/// the header, contents, and authentication tag.
/// * `contents` - Mutable buffer to write plaintext. Note that the first byte is the header byte
/// containing protocol flags.
Expand Down Expand Up @@ -304,7 +300,7 @@ impl PacketReader {
///
/// # Arguments
///
/// * `ciphertext` - The message from the peer excluding the first 3 length bytes. It should contain
/// * `ciphertext` - The packet from the peer excluding the first 3 length bytes. It should contain
/// the header, contents, and authentication tag.
/// * `aad` - Optional associated authenticated data.
///
Expand Down Expand Up @@ -335,7 +331,7 @@ impl PacketReader {
}
}

/// Prepare messages to be sent over a byte stream.
/// Prepare packets to be sent over encrypted channel to a peer.
#[derive(Clone)]
pub struct PacketWriter {
length_encoding_cipher: FSChaCha20,
Expand Down Expand Up @@ -405,7 +401,7 @@ impl PacketWriter {
/// and handle necessary memory allocation.
///
/// * `plaintext` - Plaintext content to be encrypted.
/// * `aad` - Optional authentication for the peer, currently only used for the first round of messages.
/// * `aad` - Optional associated authenticated data.
/// * `packet_type` - Is this a genuine packet or a decoy.
#[cfg(feature = "alloc")]
pub fn encrypt_packet_with_alloc(
Expand All @@ -421,7 +417,7 @@ impl PacketWriter {
}
}

/// Encrypt and decrypt messages with a peer.
/// Encrypt and decrypt packets with a peer.
#[derive(Clone)]
pub struct PacketHandler {
/// A unique identifier for the communication session.
Expand Down Expand Up @@ -558,7 +554,7 @@ fn initialize_session_key_material(
/// channel between an *initiator* and a *responder*. The next step is to call `complete_materials`
/// no matter if initiator or responder, however the responder should already have the
/// necessary materials from their peers request. `complete_materials` creates the response
/// message to be sent from each peer and `authenticate_garbage_and_version` is then used
/// packet to be sent from each peer and `authenticate_garbage_and_version` is then used
/// to verify the handshake. Finally, the `finalized` method is used to consumer the handshake
/// and return a packet handler for further communication on the channel.
pub struct Handshake<'a> {
Expand All @@ -576,18 +572,18 @@ pub struct Handshake<'a> {
packet_handler: Option<PacketHandler>,
/// 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,
/// Processesed buffer index. Store state between authentication attempts to avoid resetting ciphers.
current_buffer_index: usize,
}

impl<'a> Handshake<'a> {
/// Initialize a V2 transport handshake with a peer.
///
/// # Arguements
/// # Arguments
///
/// * `network` - The bitcoin network which both peers operate on.
/// * `garbage` - Optional garbage to send in handshake.
/// * `buffer` - Message buffer to send to peer which will include initial materials for handshake + garbage.
/// * `buffer` - Packet buffer to send to peer which will include initial materials for handshake + garbage.
///
/// # Returns
///
Expand All @@ -614,7 +610,7 @@ impl<'a> Handshake<'a> {
///
/// * `network` - The bitcoin network which both peers operate on.
/// * `garbage` - Optional garbage to send in handshake.
/// * `buffer` - Message buffer to send to peer which will include initial materials for handshake + garbage.
/// * `buffer` - Packet buffer to send to peer which will include initial materials for handshake + garbage.
/// * `rng` - Supplied Random Number Generator.
/// * `curve` - Supplied secp256k1 context.
///
Expand Down Expand Up @@ -653,7 +649,7 @@ impl<'a> Handshake<'a> {
remote_garbage_terminator: None,
packet_handler: None,
current_packet_length_bytes: None,
current_message_index: 0,
current_buffer_index: 0,
})
}

Expand Down Expand Up @@ -762,7 +758,7 @@ impl<'a> Handshake<'a> {
///
/// # 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.
/// * `CiphertextTooSmall` - The buffer did not contain all required information and should be extended (e.g. read more off a socket) and authentication re-tried.
/// * `BufferTooSmall` - The supplied packet_buffer is not large enough for decrypting the decoy and version packets.
/// * `HandshakeOutOfOrder` - The handshake sequence is in a bad state and should be restarted.
/// * `MaxGarbageLength` - Buffer did not contain the garbage terminator, should not be retried.
Expand All @@ -784,7 +780,7 @@ impl<'a> Handshake<'a> {
// 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_buffer_index == 0 {
found_version_packet = self.decrypt_packet(message, packet_buffer, Some(garbage))?;
}

Expand All @@ -802,7 +798,7 @@ impl<'a> Handshake<'a> {
Ok(())
}

/// Decrypt the next packet in the message buffer while
/// Decrypt the next packet in the buffer while
/// book keeping relevant lengths and indexes. This allows
/// the buffer to be re-processed without throwing off
/// the state of the ciphers.
Expand All @@ -823,11 +819,11 @@ impl<'a> Handshake<'a> {

if self.current_packet_length_bytes.is_none() {
// Bounds check on the input buffer.
if message.len() < self.current_message_index + NUM_LENGTH_BYTES {
if message.len() < self.current_buffer_index + NUM_LENGTH_BYTES {
return Err(Error::CiphertextTooSmall);
}
let packet_length = packet_handler.packet_reader.decypt_len(
message[self.current_message_index..NUM_LENGTH_BYTES]
message[self.current_buffer_index..NUM_LENGTH_BYTES]
.try_into()
.expect("Buffer slice must be exactly 3 bytes long"),
);
Expand All @@ -842,19 +838,19 @@ impl<'a> Handshake<'a> {
.ok_or(Error::HandshakeOutOfOrder)?;

// Bounds check on input buffer.
if message.len() < self.current_message_index + NUM_LENGTH_BYTES + packet_length {
if message.len() < self.current_buffer_index + NUM_LENGTH_BYTES + packet_length {
return Err(Error::CiphertextTooSmall);
}
packet_handler.packet_reader.decrypt_contents(
&message[self.current_message_index + NUM_LENGTH_BYTES
..self.current_message_index + NUM_LENGTH_BYTES + packet_length],
&message[self.current_buffer_index + NUM_LENGTH_BYTES
..self.current_buffer_index + NUM_LENGTH_BYTES + packet_length],
packet_buffer,
garbage,
)?;

// Mark current decryption point in the buffer.
self.current_message_index =
self.current_message_index + NUM_LENGTH_BYTES + packet_length + 1;
self.current_buffer_index =
self.current_buffer_index + NUM_LENGTH_BYTES + packet_length + 1;
self.current_packet_length_bytes = None;

// The version packet is currently just an empty packet.
Expand All @@ -881,7 +877,7 @@ impl<'a> Handshake<'a> {
///
/// # Error
///
/// * `MessageLengthTooSmall` - Buffer did not contain a garbage terminator.
/// * `CiphertextTooSmall` - Buffer did not contain a garbage terminator.
/// * `MaxGarbageLength` - Buffer did not contain the garbage terminator and contains too much garbage, should not be retried.
fn split_garbage_and_message(
message: &[u8],
Expand Down

0 comments on commit 8c9832f

Please sign in to comment.