Skip to content

Commit

Permalink
Fix chain remaining_mut(), allowing to chain growing buffer (tokio-rs…
Browse files Browse the repository at this point in the history
  • Loading branch information
Zettroke authored and lelongg committed Jan 9, 2023
1 parent a1e804d commit 211daff
Show file tree
Hide file tree
Showing 3 changed files with 27 additions and 2 deletions.
4 changes: 4 additions & 0 deletions src/buf/buf_mut.rs
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,10 @@ pub unsafe trait BufMut {
/// Implementations of `remaining_mut` should ensure that the return value
/// does not change unless a call is made to `advance_mut` or any other
/// function that is documented to change the `BufMut`'s current position.
///
/// # Note
///
/// `remaining_mut` may return value smaller than actual available space.
fn remaining_mut(&self) -> usize;

/// Advance the internal cursor of the BufMut
Expand Down
3 changes: 1 addition & 2 deletions src/buf/chain.rs
Original file line number Diff line number Diff line change
Expand Up @@ -198,8 +198,7 @@ where
fn remaining_mut(&self) -> usize {
self.a
.remaining_mut()
.checked_add(self.b.remaining_mut())
.unwrap()
.saturating_add(self.b.remaining_mut())
}

fn chunk_mut(&mut self) -> &mut UninitSlice {
Expand Down
22 changes: 22 additions & 0 deletions tests/test_chain.rs
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,28 @@ fn vectored_read() {
}
}

#[test]
fn chain_growing_buffer() {
let mut buff = [' ' as u8; 10];
let mut vec = b"wassup".to_vec();

let mut chained = (&mut buff[..]).chain_mut(&mut vec).chain_mut(Vec::new()); // Required for potential overflow because remaining_mut for Vec is isize::MAX - vec.len(), but for chain_mut is usize::MAX

chained.put_slice(b"hey there123123");

assert_eq!(&buff, b"hey there1");
assert_eq!(&vec, b"wassup23123");
}

#[test]
fn chain_overflow_remaining_mut() {
let mut chained = Vec::<u8>::new().chain_mut(Vec::new()).chain_mut(Vec::new());

assert_eq!(chained.remaining_mut(), usize::MAX);
chained.put_slice(&[0; 256]);
assert_eq!(chained.remaining_mut(), usize::MAX);
}

#[test]
fn chain_get_bytes() {
let mut ab = Bytes::copy_from_slice(b"ab");
Expand Down

0 comments on commit 211daff

Please sign in to comment.