Change BufMut methods that expose maybe-uninitialized bytes (#305)

- The return type of `BufMut::bytes_mut` is now
  `&mut [MaybeUninit<u8>]`.
- The argument type of `BufMut::bytes_vectored_mut` is now
  `&mut [bytes::buf::IoSliceMut]`.
- `bytes::buf::IoSliceMut` is a `repr(transparent)` wrapper around an
  `std::io::IoSliceMut`, but does not expose the inner bytes with a safe
  API, since they might be uninitialized.
- `BufMut::bytesMut` and `BufMut::bytes_vectored_mut` are no longer
  `unsafe fn`, since the types encapsulate the unsafety instead.

Author: seanmonstar

src/buf/buf_mut.rs

@@ -1,10 +1,10 @@
 #[cfg(feature = "std")]
 use super::Writer;
 
-use core::{mem, cmp, ptr, usize};
+use core::{cmp, mem::{self, MaybeUninit}, ptr, usize};
 
 #[cfg(feature = "std")]
-use std::io::IoSliceMut;
+use std::fmt;
 
 use alloc::{vec::Vec, boxed::Box};
 
@@ -72,13 +72,15 @@ pub trait BufMut {
     /// let mut buf = Vec::with_capacity(16);
     ///
     /// unsafe {
-    ///     buf.bytes_mut()[0] = b'h';
-    ///     buf.bytes_mut()[1] = b'e';
+    ///     // MaybeUninit::as_mut_ptr
+    ///     buf.bytes_mut()[0].as_mut_ptr().write(b'h');
+    ///     buf.bytes_mut()[1].as_mut_ptr().write(b'e');
     ///
     ///     buf.advance_mut(2);
     ///
-    ///     buf.bytes_mut()[0] = b'l';
-    ///     buf.bytes_mut()[1..3].copy_from_slice(b"lo");
+    ///     buf.bytes_mut()[0].as_mut_ptr().write(b'l');
+    ///     buf.bytes_mut()[1].as_mut_ptr().write(b'l');
+    ///     buf.bytes_mut()[2].as_mut_ptr().write(b'o');
     ///
     ///     buf.advance_mut(3);
     /// }
@@ -139,13 +141,15 @@ pub trait BufMut {
     /// let mut buf = Vec::with_capacity(16);
     ///
     /// unsafe {
-    ///     buf.bytes_mut()[0] = b'h';
-    ///     buf.bytes_mut()[1] = b'e';
+    ///     // MaybeUninit::as_mut_ptr
+    ///     buf.bytes_mut()[0].as_mut_ptr().write(b'h');
+    ///     buf.bytes_mut()[1].as_mut_ptr().write(b'e');
     ///
     ///     buf.advance_mut(2);
     ///
-    ///     buf.bytes_mut()[0] = b'l';
-    ///     buf.bytes_mut()[1..3].copy_from_slice(b"lo");
+    ///     buf.bytes_mut()[0].as_mut_ptr().write(b'l');
+    ///     buf.bytes_mut()[1].as_mut_ptr().write(b'l');
+    ///     buf.bytes_mut()[2].as_mut_ptr().write(b'o');
     ///
     ///     buf.advance_mut(3);
     /// }
@@ -161,7 +165,7 @@ pub trait BufMut {
     /// `bytes_mut` returning an empty slice implies that `remaining_mut` will
     /// return 0 and `remaining_mut` returning 0 implies that `bytes_mut` will
     /// return an empty slice.
-    unsafe fn bytes_mut(&mut self) -> &mut [u8];
+    fn bytes_mut(&mut self) -> &mut [MaybeUninit<u8>];
 
     /// Fills `dst` with potentially multiple mutable slices starting at `self`'s
     /// current position.
@@ -192,13 +196,13 @@ pub trait BufMut {
     ///
     /// [`readv`]: http://man7.org/linux/man-pages/man2/readv.2.html
     #[cfg(feature = "std")]
-    unsafe fn bytes_vectored_mut<'a>(&'a mut self, dst: &mut [IoSliceMut<'a>]) -> usize {
+    fn bytes_vectored_mut<'a>(&'a mut self, dst: &mut [IoSliceMut<'a>]) -> usize {
         if dst.is_empty() {
             return 0;
         }
 
         if self.has_remaining_mut() {
-            dst[0] = IoSliceMut::new(self.bytes_mut());
+            dst[0] = IoSliceMut::from(self.bytes_mut());
             1
         } else {
             0
@@ -238,7 +242,7 @@ pub trait BufMut {
 
                 ptr::copy_nonoverlapping(
                     s.as_ptr(),
-                    d.as_mut_ptr(),
+                    d.as_mut_ptr() as *mut u8,
                     l);
             }
 
@@ -280,7 +284,7 @@ pub trait BufMut {
 
                 ptr::copy_nonoverlapping(
                     src[off..].as_ptr(),
-                    dst.as_mut_ptr(),
+                    dst.as_mut_ptr() as *mut u8,
                     cnt);
 
                 off += cnt;
@@ -931,12 +935,12 @@ impl<T: BufMut + ?Sized> BufMut for &mut T {
         (**self).remaining_mut()
     }
 
-    unsafe fn bytes_mut(&mut self) -> &mut [u8] {
+    fn bytes_mut(&mut self) -> &mut [MaybeUninit<u8>] {
         (**self).bytes_mut()
     }
 
     #[cfg(feature = "std")]
-    unsafe fn bytes_vectored_mut<'b>(&'b mut self, dst: &mut [IoSliceMut<'b>]) -> usize {
+    fn bytes_vectored_mut<'b>(&'b mut self, dst: &mut [IoSliceMut<'b>]) -> usize {
         (**self).bytes_vectored_mut(dst)
     }
 
@@ -950,12 +954,12 @@ impl<T: BufMut + ?Sized> BufMut for Box<T> {
         (**self).remaining_mut()
     }
 
-    unsafe fn bytes_mut(&mut self) -> &mut [u8] {
+    fn bytes_mut(&mut self) -> &mut [MaybeUninit<u8>] {
         (**self).bytes_mut()
     }
 
     #[cfg(feature = "std")]
-    unsafe fn bytes_vectored_mut<'b>(&'b mut self, dst: &mut [IoSliceMut<'b>]) -> usize {
+    fn bytes_vectored_mut<'b>(&'b mut self, dst: &mut [IoSliceMut<'b>]) -> usize {
         (**self).bytes_vectored_mut(dst)
     }
 
@@ -971,8 +975,9 @@ impl BufMut for &mut [u8] {
     }
 
     #[inline]
-    unsafe fn bytes_mut(&mut self) -> &mut [u8] {
-        self
+    fn bytes_mut(&mut self) -> &mut [MaybeUninit<u8>] {
+        // MaybeUninit is repr(transparent), so safe to transmute
+        unsafe { mem::transmute(&mut **self) }
     }
 
     #[inline]
@@ -1003,7 +1008,7 @@ impl BufMut for Vec<u8> {
     }
 
     #[inline]
-    unsafe fn bytes_mut(&mut self) -> &mut [u8] {
+    fn bytes_mut(&mut self) -> &mut [MaybeUninit<u8>] {
         use core::slice;
 
         if self.capacity() == self.len() {
@@ -1013,11 +1018,54 @@ impl BufMut for Vec<u8> {
         let cap = self.capacity();
         let len = self.len();
 
-        let ptr = self.as_mut_ptr();
-        &mut slice::from_raw_parts_mut(ptr, cap)[len..]
+        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
+        unsafe {
+            &mut slice::from_raw_parts_mut(ptr, cap)[len..]
+        }
     }
 }
 
 // The existence of this function makes the compiler catch if the BufMut
 // trait is "object-safe" or not.
 fn _assert_trait_object(_b: &dyn BufMut) {}
+
+// ===== impl IoSliceMut =====
+
+/// A buffer type used for `readv`.
+///
+/// This is a wrapper around an `std::io::IoSliceMut`, but does not expose
+/// the inner bytes in a safe API, as they may point at uninitialized memory.
+///
+/// This is `repr(transparent)` of the `std::io::IoSliceMut`, so it is valid to
+/// transmute them. However, as the memory might be uninitialized, care must be
+/// taken to not *read* the internal bytes, only *write* to them.
+#[repr(transparent)]
+#[cfg(feature = "std")]
+pub struct IoSliceMut<'a>(std::io::IoSliceMut<'a>);
+
+#[cfg(feature = "std")]
+impl fmt::Debug for IoSliceMut<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("IoSliceMut")
+            .field("len", &self.0.len())
+            .finish()
+    }
+}
+
+#[cfg(feature = "std")]
+impl<'a> From<&'a mut [u8]> for IoSliceMut<'a> {
+    fn from(buf: &'a mut [u8]) -> IoSliceMut<'a> {
+        IoSliceMut(std::io::IoSliceMut::new(buf))
+    }
+}
+
+#[cfg(feature = "std")]
+impl<'a> From<&'a mut [MaybeUninit<u8>]> for IoSliceMut<'a> {
+    fn from(buf: &'a mut [MaybeUninit<u8>]) -> IoSliceMut<'a> {
+        IoSliceMut(std::io::IoSliceMut::new(unsafe {
+            // We don't look at the contents, and `std::io::IoSliceMut`
+            // doesn't either.
+            mem::transmute::<&'a mut [MaybeUninit<u8>], &'a mut [u8]>(buf)
+        }))
+    }
+}

src/buf/chain.rs

@@ -1,8 +1,12 @@
 use crate::{Buf, BufMut};
 use crate::buf::IntoIter;
 
+use core::mem::MaybeUninit;
+
+#[cfg(feature = "std")]
+use std::io::{IoSlice};
 #[cfg(feature = "std")]
-use std::io::{IoSlice, IoSliceMut};
+use crate::buf::IoSliceMut;
 
 /// A `Chain` sequences two buffers.
 ///
@@ -196,7 +200,7 @@ impl<T, U> BufMut for Chain<T, U>
         self.a.remaining_mut() + self.b.remaining_mut()
     }
 
-    unsafe fn bytes_mut(&mut self) -> &mut [u8] {
+    fn bytes_mut(&mut self) -> &mut [MaybeUninit<u8>] {
         if self.a.has_remaining_mut() {
             self.a.bytes_mut()
         } else {
@@ -223,7 +227,7 @@ impl<T, U> BufMut for Chain<T, U>
     }
 
     #[cfg(feature = "std")]
-    unsafe fn bytes_vectored_mut<'a>(&'a mut self, dst: &mut [IoSliceMut<'a>]) -> usize {
+    fn bytes_vectored_mut<'a>(&'a mut self, dst: &mut [IoSliceMut<'a>]) -> usize {
         let mut n = self.a.bytes_vectored_mut(dst);
         n += self.b.bytes_vectored_mut(&mut dst[n..]);
         n

src/buf/mod.rs

@@ -31,6 +31,8 @@ mod writer;
 
 pub use self::buf_impl::Buf;
 pub use self::buf_mut::BufMut;
+#[cfg(feature = "std")]
+pub use self::buf_mut::IoSliceMut;
 pub use self::chain::Chain;
 pub use self::iter::IntoIter;
 #[cfg(feature = "std")]

src/bytes_mut.rs

@@ -926,19 +926,9 @@ impl BufMut for BytesMut {
     }
 
     #[inline]
-    unsafe fn bytes_mut(&mut self) -> &mut [u8] {
-        slice::from_raw_parts_mut(self.ptr.as_ptr().offset(self.len as isize), self.cap)
-    }
-
-    #[inline]
-    fn put_slice(&mut self, src: &[u8]) {
-        assert!(self.remaining_mut() >= src.len());
-
-        let len = src.len();
-
+    fn bytes_mut(&mut self) -> &mut [mem::MaybeUninit<u8>] {
         unsafe {
-            self.bytes_mut()[..len].copy_from_slice(src);
-            self.advance_mut(len);
+            slice::from_raw_parts_mut(self.ptr.as_ptr().offset(self.len as isize) as *mut mem::MaybeUninit<u8>, self.cap)
         }
     }
 }
@@ -1085,11 +1075,12 @@ impl Extend<u8> for BytesMut {
         let (lower, _) = iter.size_hint();
         self.reserve(lower);
 
+        // TODO: optimize
+        // 1. If self.kind() == KIND_VEC, use Vec::extend
+        // 2. Make `reserve` inline-able
         for b in iter {
-            unsafe {
-                self.bytes_mut()[0] = b;
-                self.advance_mut(1);
-            }
+            self.reserve(1);
+            self.put_u8(b);
         }
     }
 }

src/either.rs

@@ -4,7 +4,9 @@ use either::Either;
 use either::Either::*;
 
 #[cfg(feature = "std")]
-use std::io::{IoSlice, IoSliceMut};
+use std::io::IoSlice;
+#[cfg(feature = "std")]
+use crate::buf::IoSliceMut;
 
 impl<L, R> Buf for Either<L, R>
 where
@@ -60,15 +62,15 @@ where
         }
     }
 
-    unsafe fn bytes_mut(&mut self) -> &mut [u8] {
+    fn bytes_mut(&mut self) -> &mut [core::mem::MaybeUninit<u8>] {
         match *self {
             Left(ref mut b) => b.bytes_mut(),
             Right(ref mut b) => b.bytes_mut(),
         }
     }
 
     #[cfg(feature = "std")]
-    unsafe fn bytes_vectored_mut<'a>(&'a mut self, dst: &mut [IoSliceMut<'a>]) -> usize {
+    fn bytes_vectored_mut<'a>(&'a mut self, dst: &mut [IoSliceMut<'a>]) -> usize {
         match *self {
             Left(ref mut b) => b.bytes_vectored_mut(dst),
             Right(ref mut b) => b.bytes_vectored_mut(dst),

tests/test_buf_mut.rs

@@ -1,19 +1,16 @@
 #![deny(warnings, rust_2018_idioms)]
 
-use bytes::{BufMut, BytesMut};
+use bytes::{buf::IoSliceMut, BufMut, BytesMut};
 use std::usize;
 use std::fmt::Write;
-use std::io::IoSliceMut;
 
 #[test]
 fn test_vec_as_mut_buf() {
     let mut buf = Vec::with_capacity(64);
 
     assert_eq!(buf.remaining_mut(), usize::MAX);
 
-    unsafe {
-        assert!(buf.bytes_mut().len() >= 64);
-    }
+    assert!(buf.bytes_mut().len() >= 64);
 
     buf.put(&b"zomg"[..]);
 
@@ -72,20 +69,16 @@ fn test_clone() {
 fn test_bufs_vec_mut() {
     let b1: &mut [u8] = &mut [];
     let b2: &mut [u8] = &mut [];
-    let mut dst = [IoSliceMut::new(b1), IoSliceMut::new(b2)];
+    let mut dst = [IoSliceMut::from(b1), IoSliceMut::from(b2)];
 
     // with no capacity
     let mut buf = BytesMut::new();
     assert_eq!(buf.capacity(), 0);
-    unsafe {
-        assert_eq!(0, buf.bytes_vectored_mut(&mut dst[..]));
-    }
+    assert_eq!(0, buf.bytes_vectored_mut(&mut dst[..]));
 
     // with capacity
     let mut buf = BytesMut::with_capacity(64);
-    unsafe {
-        assert_eq!(1, buf.bytes_vectored_mut(&mut dst[..]));
-    }
+    assert_eq!(1, buf.bytes_vectored_mut(&mut dst[..]));
 }
 
 #[test]

tests/test_bytes.rs

@@ -464,6 +464,16 @@ fn extend_from_slice_mut() {
     }
 }
 
+#[test]
+fn extend_mut_without_size_hint() {
+    let mut bytes = BytesMut::with_capacity(0);
+    let mut long_iter = LONG.iter();
+
+    // Use iter::from_fn since it doesn't know a size_hint
+    bytes.extend(std::iter::from_fn(|| long_iter.next()));
+    assert_eq!(*bytes, LONG[..]);
+}
+
 #[test]
 fn from_static() {
     let mut a = Bytes::from_static(b"ab");