diff --git a/src/buf/buf.rs b/src/buf/buf.rs index b72c8d91c..0b81e35b7 100644 --- a/src/buf/buf.rs +++ b/src/buf/buf.rs @@ -1120,6 +1120,30 @@ impl> Buf for io::Cursor { } } +impl<'a> Buf for io::Take<&'a [u8]> { + #[inline] + fn remaining(&self) -> usize { + let len = self.get_ref().len() as u64; + let limit = self.limit(); + // The smaller of the two values will always fit in usize. + cmp::min(len, limit) as usize + } + + fn bytes(&self) -> &[u8] { + &self.get_ref()[..self.remaining()] + } + + fn advance(&mut self, cnt: usize) { + let remaining = self.remaining(); + assert!(cnt <= remaining); + // Use the actual number of remaining bytes as new limit, even if limit + // was greater than remaining before. + self.set_limit((remaining - cnt) as u64); + let slice = self.get_mut(); + *slice = &slice[cnt..]; + } +} + impl Buf for Option<[u8; 1]> { fn remaining(&self) -> usize { if self.is_some() { diff --git a/tests/test_buf.rs b/tests/test_buf.rs index f25c25f2b..0de0ea417 100644 --- a/tests/test_buf.rs +++ b/tests/test_buf.rs @@ -56,3 +56,41 @@ fn test_bufs_vec() { assert_eq!(1, buf.bytes_vec(&mut dst[..])); } + +#[test] +fn test_take() { + // Pulling Read into the scope would result in a conflict between + // Buf::bytes() from Read::bytes(). + let mut buf = std::io::Read::take(&b"hello world"[..], 5); + assert_eq!(buf.bytes(), b"hello"); + assert_eq!(buf.remaining(), 5); + + buf.advance(3); + assert_eq!(buf.bytes(), b"lo"); + assert_eq!(buf.remaining(), 2); + + buf.advance(2); + assert_eq!(buf.bytes(), b""); + assert_eq!(buf.remaining(), 0); +} + +#[test] +#[should_panic] +fn test_take_advance_too_far() { + let mut buf = std::io::Read::take(&b"hello world"[..], 5); + buf.advance(10); +} + +#[test] +fn test_take_limit_gt_length() { + // The byte array has only 11 bytes, but we take 15 bytes. + let mut buf = std::io::Read::take(&b"hello world"[..], 15); + assert_eq!(buf.remaining(), 11); + assert_eq!(buf.limit(), 15); + + buf.advance(5); + assert_eq!(buf.remaining(), 6); + // The limit is reduced my more than the number of bytes we advanced, to + // the actual number of remaining bytes in the buffer. + assert_eq!(buf.limit(), 6); +}