From 8135c1f60692eff20cbc3f3d297704aa7600dbb7 Mon Sep 17 00:00:00 2001 From: Carl Lerche Date: Wed, 20 Nov 2019 12:11:40 -0800 Subject: [PATCH] implicitly grow BytesMut; add BufMutExt::chain_mut (#316) This brings `BytesMut` in line with `Vec` behavior. This also fixes an existing bug in BytesMut::bytes_mut that exposes invalid slices. The bug was recently introduced and was only on master and never released to `crates.io`. In order to fix a test, `BufMutExt::chain_mut` is provided. Withou this, it is not possible to chain two `&mut [u8]`. Closes #170 --- src/buf/buf_mut.rs | 2 +- src/buf/ext/mod.rs | 26 ++++++++++++++++++++++++++ src/bytes_mut.rs | 35 ++++++++++++++++++++--------------- tests/test_buf_mut.rs | 2 +- tests/test_bytes.rs | 37 +++++++++++++++++++++++++++++++++++-- tests/test_chain.rs | 13 +++++-------- 6 files changed, 88 insertions(+), 27 deletions(-) diff --git a/src/buf/buf_mut.rs b/src/buf/buf_mut.rs index 8ec568aae..d29591010 100644 --- a/src/buf/buf_mut.rs +++ b/src/buf/buf_mut.rs @@ -270,7 +270,7 @@ pub trait BufMut { fn put_slice(&mut self, src: &[u8]) { let mut off = 0; - assert!(self.remaining_mut() >= src.len(), "buffer overflow"); + assert!(self.remaining_mut() >= src.len(), "buffer overflow; remaining = {}; src = {}", self.remaining_mut(), src.len()); while off < src.len() { let cnt; diff --git a/src/buf/ext/mod.rs b/src/buf/ext/mod.rs index 10b9a34c0..2fefde90b 100644 --- a/src/buf/ext/mod.rs +++ b/src/buf/ext/mod.rs @@ -124,6 +124,32 @@ pub trait BufMutExt: BufMut { fn writer(self) -> Writer where Self: Sized { writer::new(self) } + + /// Creates an adapter which will chain this buffer with another. + /// + /// The returned `BufMut` instance will first write to all bytes from + /// `self`. Afterwards, it will write to `next`. + /// + /// # Examples + /// + /// ``` + /// use bytes::{BufMut, buf::BufMutExt}; + /// + /// let mut a = [0u8; 5]; + /// let mut b = [0u8; 6]; + /// + /// let mut chain = (&mut a[..]).chain_mut(&mut b[..]); + /// + /// chain.put_slice(b"hello world"); + /// + /// assert_eq!(&a[..], b"hello"); + /// assert_eq!(&b[..], b" world"); + /// ``` + fn chain_mut(self, next: U) -> Chain + where Self: Sized + { + Chain::new(self, next) + } } impl BufMutExt for B {} diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index 6cf791a01..ea71a9aaf 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -1,4 +1,5 @@ -use core::{cmp, fmt, hash, isize, mem, slice, usize}; +use core::{cmp, fmt, hash, isize, slice, usize}; +use core::mem::{self, ManuallyDrop}; use core::ops::{Deref, DerefMut}; use core::ptr::{self, NonNull}; use core::iter::{FromIterator, Iterator}; @@ -559,17 +560,15 @@ impl BytesMut { self.cap += off; } else { // No space - allocate more - let mut v = rebuild_vec(self.ptr.as_ptr(), self.len, self.cap, off); + let mut v = ManuallyDrop::new(rebuild_vec(self.ptr.as_ptr(), self.len, self.cap, off)); v.reserve(additional); // Update the info self.ptr = vptr(v.as_mut_ptr().offset(off as isize)); self.len = v.len() - off; self.cap = v.capacity() - off; - - // Drop the vec reference - mem::forget(v); } + return; } } @@ -582,7 +581,8 @@ impl BytesMut { // allocating a new vector with the requested capacity. // // Compute the new capacity - let mut new_cap = len + additional; + let mut new_cap = len.checked_add(additional).expect("overflow"); + let original_capacity; let original_capacity_repr; @@ -618,8 +618,10 @@ impl BytesMut { // There are some situations, using `reserve_exact` that the // buffer capacity could be below `original_capacity`, so do a // check. + let double = v.capacity().checked_shl(1).unwrap_or(new_cap); + new_cap = cmp::max( - cmp::max(v.capacity() << 1, new_cap), + cmp::max(double, new_cap), original_capacity); } else { new_cap = cmp::max(new_cap, original_capacity); @@ -627,7 +629,7 @@ impl BytesMut { } // Create a new vector to store the data - let mut v = Vec::with_capacity(new_cap); + let mut v = ManuallyDrop::new(Vec::with_capacity(new_cap)); // Copy the bytes v.extend_from_slice(self.as_ref()); @@ -642,10 +644,6 @@ impl BytesMut { self.ptr = vptr(v.as_mut_ptr()); self.len = v.len(); self.cap = v.capacity(); - - // Forget the vector handle - mem::forget(v); - } /// Appends given bytes to this object. /// @@ -924,20 +922,27 @@ impl Buf for BytesMut { impl BufMut for BytesMut { #[inline] fn remaining_mut(&self) -> usize { - self.capacity() - self.len() + usize::MAX - self.len() } #[inline] unsafe fn advance_mut(&mut self, cnt: usize) { let new_len = self.len() + cnt; - assert!(new_len <= self.cap); + assert!(new_len <= self.cap, "new_len = {}; capacity = {}", new_len, self.cap); self.len = new_len; } #[inline] fn bytes_mut(&mut self) -> &mut [mem::MaybeUninit] { + if self.capacity() == self.len() { + self.reserve(64); + } + unsafe { - slice::from_raw_parts_mut(self.ptr.as_ptr().offset(self.len as isize) as *mut mem::MaybeUninit, self.cap) + let ptr = self.ptr.as_ptr().offset(self.len as isize); + let len = self.cap - self.len; + + slice::from_raw_parts_mut(ptr as *mut mem::MaybeUninit, len) } } } diff --git a/tests/test_buf_mut.rs b/tests/test_buf_mut.rs index 0d372d1ec..dd414e5a2 100644 --- a/tests/test_buf_mut.rs +++ b/tests/test_buf_mut.rs @@ -74,7 +74,7 @@ fn test_bufs_vec_mut() { // with no capacity let mut buf = BytesMut::new(); assert_eq!(buf.capacity(), 0); - assert_eq!(0, buf.bytes_vectored_mut(&mut dst[..])); + assert_eq!(1, buf.bytes_vectored_mut(&mut dst[..])); // with capacity let mut buf = BytesMut::with_capacity(64); diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs index 9fa8019ee..6e3c4b6da 100644 --- a/tests/test_bytes.rs +++ b/tests/test_bytes.rs @@ -2,6 +2,8 @@ use bytes::{Bytes, BytesMut, Buf, BufMut}; +use std::usize; + const LONG: &'static [u8] = b"mary had a little lamb, little lamb, little lamb"; const SHORT: &'static [u8] = b"hello world"; @@ -93,8 +95,8 @@ fn fmt_write() { let mut c = BytesMut::with_capacity(64); - write!(c, "{}", s).unwrap_err(); - assert!(c.is_empty()); + write!(c, "{}", s).unwrap(); + assert_eq!(c, s[..].as_bytes()); } #[test] @@ -820,3 +822,34 @@ fn empty_slice_ref_catches_not_an_empty_subset() { bytes.slice_ref(slice); } + +#[test] +fn bytes_buf_mut_advance() { + let mut bytes = BytesMut::with_capacity(1024); + + unsafe { + let ptr = bytes.bytes_mut().as_ptr(); + assert_eq!(1024, bytes.bytes_mut().len()); + + bytes.advance_mut(10); + + let next = bytes.bytes_mut().as_ptr(); + assert_eq!(1024 - 10, bytes.bytes_mut().len()); + assert_eq!(ptr.offset(10), next); + + // advance to the end + bytes.advance_mut(1024 - 10); + + // The buffer size is doubled + assert_eq!(1024, bytes.bytes_mut().len()); + } +} + +#[test] +#[should_panic] +fn bytes_reserve_overflow() { + let mut bytes = BytesMut::with_capacity(1024); + bytes.put_slice(b"hello world"); + + bytes.reserve(usize::MAX); +} diff --git a/tests/test_chain.rs b/tests/test_chain.rs index 877618843..332571d8b 100644 --- a/tests/test_chain.rs +++ b/tests/test_chain.rs @@ -1,7 +1,7 @@ #![deny(warnings, rust_2018_idioms)] -use bytes::{Buf, BufMut, Bytes, BytesMut}; -use bytes::buf::BufExt; +use bytes::{Buf, BufMut, Bytes}; +use bytes::buf::{BufExt, BufMutExt}; use std::io::IoSlice; #[test] @@ -15,20 +15,17 @@ fn collect_two_bufs() { #[test] fn writing_chained() { - let mut a = BytesMut::with_capacity(64); - let mut b = BytesMut::with_capacity(64); + let mut a = [0u8; 64]; + let mut b = [0u8; 64]; { - let mut buf = (&mut a).chain(&mut b); + let mut buf = (&mut a[..]).chain_mut(&mut b[..]); for i in 0u8..128 { buf.put_u8(i); } } - assert_eq!(64, a.len()); - assert_eq!(64, b.len()); - for i in 0..64 { let expect = i as u8; assert_eq!(expect, a[i]);