Skip to content

Commit

Permalink
fix(dht): fix over allocation for encryupted messages
Browse files Browse the repository at this point in the history
  • Loading branch information
sdbondi committed Oct 20, 2022
1 parent 659f4b8 commit 3983a66
Showing 1 changed file with 7 additions and 5 deletions.
12 changes: 7 additions & 5 deletions comms/dht/src/crypt.rs
Original file line number Diff line number Diff line change
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.reserve(padding_length);
message.extend(iter::repeat(0u8).take(padding_length));

Ok(())
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,8 @@ 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();
assert!(pad_message.capacity() >= message.len() + paf_len);

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

0 comments on commit 3983a66

Please sign in to comment.