io: fix unsound Buf::ensure_capacity_for

Author: paolobarbolini

tokio/src/io/blocking.rs

@@ -21,6 +21,7 @@ pub(crate) struct Blocking<T> {
 pub(crate) struct Buf {
     buf: Vec<u8>,
     pos: usize,
+    init_len: usize,
 }
 
 pub(crate) const DEFAULT_MAX_BUF_SIZE: usize = 2 * 1024 * 1024;
@@ -190,6 +191,7 @@ impl Buf {
         Buf {
             buf: Vec::with_capacity(n),
             pos: 0,
+            init_len: 0,
         }
     }
 
@@ -220,6 +222,7 @@ impl Buf {
         let n = cmp::min(src.len(), max_buf_size);
 
         self.buf.extend_from_slice(&src[..n]);
+        self.init_len = cmp::max(self.init_len, self.buf.len());
         n
     }
 
@@ -236,6 +239,27 @@ impl Buf {
             self.buf.reserve(len - self.buf.len());
         }
 
+        if self.init_len < len {
+            debug_assert!(
+                self.init_len < self.buf.capacity(),
+                "init_len of Vec is bigger than the capacity"
+            );
+            debug_assert!(
+                len <= self.buf.capacity(),
+                "uninit area of Vec is bigger than the capacity"
+            );
+
+            let uninit_len = len - self.init_len;
+            // SAFETY: the area is within the allocation of the Vec
+            unsafe {
+                self.buf
+                    .as_mut_ptr()
+                    .add(self.init_len)
+                    .write_bytes(0, uninit_len);
+            }
+        }
+
+        // SAFETY: `len` is within the capacity and is init
         unsafe {
             self.buf.set_len(len);
         }
@@ -287,6 +311,7 @@ cfg_fs! {
                 self.buf.extend_from_slice(&buf[..len]);
                 rem -= len;
             }
+            self.init_len = cmp::max(self.init_len, self.buf.len());
 
             max_buf_size - rem
         }

tokio/src/io/mod.rs

@@ -293,3 +293,6 @@ cfg_io_blocking! {
         pub(crate) use crate::blocking::JoinHandle as Blocking;
     }
 }
+
+#[cfg(test)]
+mod tests;

tokio/src/io/tests.rs

@@ -0,0 +1,34 @@
+use std::{
+    io::{self, Read},
+    mem::MaybeUninit,
+};
+
+use super::blocking::Buf;
+use crate::io::ReadBuf;
+
+// Have miri check that `Buf::ensure_capacity_for` initializes its length.
+#[test]
+fn buf_ensure_capacity_for_len_is_init() {
+    const MAX_BUF_SIZE: usize = 128;
+
+    let mut buf = Buf::with_capacity(0);
+
+    let mut dst = [MaybeUninit::uninit(); 96];
+    buf.ensure_capacity_for(&ReadBuf::uninit(&mut dst), MAX_BUF_SIZE);
+    let res = buf
+        .read_from(&mut EnsureInitReader(&[123u8; 64][..]))
+        .unwrap();
+    assert_eq!(res, 64);
+    assert_eq!(buf.bytes(), [123u8; 64]);
+}
+
+struct EnsureInitReader<R>(R);
+
+impl<R: Read> Read for EnsureInitReader<R> {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        // ensure `buf` is init when run under miri
+        for &_b in &*buf {}
+
+        self.0.read(buf)
+    }
+}