io: extend via Vec::set_len only after having read into it

tokio/src/fs/file.rs

@@ -604,7 +604,7 @@ impl AsyncRead for File {
                     let std = me.std.clone();
 
                     inner.state = State::Busy(spawn_blocking(move || {
-                        let res = buf.read_from(&mut &*std);
+                        let res = unsafe { buf.read_from(&mut &*std) };
                         (Operation::Read(res), buf)
                     }));
                 }

tokio/src/io/blocking.rs

@@ -5,6 +5,7 @@ use std::cmp;
 use std::future::Future;
 use std::io;
 use std::io::prelude::*;
+use std::mem::MaybeUninit;
 use std::pin::Pin;
 use std::task::{ready, Context, Poll};
 
@@ -21,6 +22,7 @@ pub(crate) struct Blocking<T> {
 pub(crate) struct Buf {
     buf: Vec<u8>,
     pos: usize,
+    capacity: usize,
 }
 
 pub(crate) const DEFAULT_MAX_BUF_SIZE: usize = 2 * 1024 * 1024;
@@ -68,7 +70,7 @@ where
                     let mut inner = self.inner.take().unwrap();
 
                     self.state = State::Busy(sys::run(move || {
-                        let res = buf.read_from(&mut inner);
+                        let res = unsafe { buf.read_from(&mut inner) };
                         (res, buf, inner)
                     }));
                 }
@@ -190,6 +192,7 @@ impl Buf {
         Buf {
             buf: Vec::with_capacity(n),
             pos: 0,
+            capacity: 0,
         }
     }
 
@@ -230,22 +233,20 @@ impl Buf {
     pub(crate) fn ensure_capacity_for(&mut self, bytes: &ReadBuf<'_>, max_buf_size: usize) {
         assert!(self.is_empty());
 
-        let len = cmp::min(bytes.remaining(), max_buf_size);
-
-        if self.buf.len() < len {
-            self.buf.reserve(len - self.buf.len());
-        }
-
-        unsafe {
-            self.buf.set_len(len);
-        }
+        let capacity = cmp::min(bytes.remaining(), max_buf_size);
+        self.buf.reserve(capacity);
+        self.capacity = capacity;
     }
 
-    pub(crate) fn read_from<T: Read>(&mut self, rd: &mut T) -> io::Result<usize> {
-        let res = uninterruptibly!(rd.read(&mut self.buf));
+    /// Safety: `rd` must not read from the buffer passed to `read` and
+    /// must correctly report the length of the written portion of the buffer.
+    pub(crate) unsafe fn read_from<T: Read>(&mut self, rd: &mut T) -> io::Result<usize> {
+        let buf = &mut self.buf.spare_capacity_mut()[..self.capacity - self.pos];
+        let buf = unsafe { &mut *(buf as *mut [MaybeUninit<u8>] as *mut [u8]) };
+        let res = uninterruptibly!(rd.read(buf));
 
         if let Ok(n) = res {
-            self.buf.truncate(n);
+            unsafe { self.buf.set_len(n) }
         } else {
             self.buf.clear();
         }