Skip to content

Commit

Permalink
fix(dht): fix over allocation for encrypted messages (#4832)
Browse files Browse the repository at this point in the history
* fix(dht): fix over allocation for encryupted messages

* fix: always add padding bytes

Co-authored-by: Cayle Sharrock <CjS77@users.noreply.github.com>
  • Loading branch information
sdbondi and CjS77 committed Oct 20, 2022
1 parent 8289493 commit d29a64c
Showing 1 changed file with 9 additions and 6 deletions.
15 changes: 9 additions & 6 deletions comms/dht/src/crypt.rs
Expand Up @@ -86,11 +86,12 @@ fn get_message_padding_length(message_length: usize) -> usize {
}
}

/// Pads a message to a multiple of MESSAGE_BASE_LENGTH excluding the additional prefix space
/// Pads a message to a multiple of MESSAGE_BASE_LENGTH excluding the additional prefix space.
/// This function returns the number of additional padding bytes appended to the message.
fn pad_message_to_base_length_multiple(
message: &mut BytesMut,
additional_prefix_space: usize,
) -> Result<(), DhtEncryptError> {
) -> Result<usize, DhtEncryptError> {
// We require a 32-bit length representation, and also don't want to overflow after including this encoding
if message.len() > u32::MAX as usize {
return Err(DhtEncryptError::PaddingError("Message is too long".to_string()));
Expand All @@ -99,10 +100,10 @@ fn pad_message_to_base_length_multiple(
get_message_padding_length(message.len().checked_sub(additional_prefix_space).ok_or_else(|| {
DhtEncryptError::PaddingError("Message length shorter than the additional_prefix_space".to_string())
})?);
message.reserve(message.len() + padding_length);
message.extend(iter::repeat(0u8).take(padding_length));

Ok(())
message.resize(message.len() + padding_length, 0);

Ok(padding_length)
}

/// Returns the unpadded message. The messages must have the length prefixed to it and the nonce is removec.
Expand Down Expand Up @@ -445,7 +446,9 @@ mod test {
.collect::<Vec<_>>();

let mut pad_message = BytesMut::from(message);
pad_message_to_base_length_multiple(&mut pad_message, 0).unwrap();
let pad_len = pad_message_to_base_length_multiple(&mut pad_message, 0).unwrap();
// For small messages less than MESSAGE_BASE_LENGTH we can expect an exact capacity
assert_eq!(pad_message.capacity(), message.len() + pad_len);

// padded message is of correct length
assert_eq!(pad_message.len(), MESSAGE_BASE_LENGTH);
Expand Down

0 comments on commit d29a64c

Please sign in to comment.