1 files changed, 17 insertions, 68 deletions

diff --git a/tokio/src/io/util/read_to_end.rs b/tokio/src/io/util/read_to_end.rs
index 29b8b811..609af28e 100644
--- a/tokio/src/io/util/read_to_end.rs
+++ b/tokio/src/io/util/read_to_end.rs
@@ -1,4 +1,4 @@
-use crate::io::AsyncRead;
+use crate::io::{AsyncRead, ReadBuf};
 
 use std::future::Future;
 use std::io;
@@ -21,7 +21,6 @@ pub(crate) fn read_to_end<'a, R>(reader: &'a mut R, buffer: &'a mut Vec<u8>) ->
 where
     R: AsyncRead + Unpin + ?Sized,
 {
-    prepare_buffer(buffer, reader);
     ReadToEnd {
         reader,
         buf: buffer,
@@ -29,12 +28,7 @@ where
     }
 }
 
-/// # Safety
-///
-/// Before first calling this method, the unused capacity must have been
-/// prepared for use with the provided AsyncRead. This can be done using the
-/// `prepare_buffer` function later in this file.
-pub(super) unsafe fn read_to_end_internal<R: AsyncRead + ?Sized>(
+pub(super) fn read_to_end_internal<R: AsyncRead + ?Sized>(
     buf: &mut Vec<u8>,
     mut reader: Pin<&mut R>,
     num_read: &mut usize,
@@ -56,13 +50,7 @@ pub(super) unsafe fn read_to_end_internal<R: AsyncRead + ?Sized>(
 /// Tries to read from the provided AsyncRead.
 ///
 /// The length of the buffer is increased by the number of bytes read.
-///
-/// # Safety
-///
-/// The caller ensures that the buffer has been prepared for use with the
-/// AsyncRead before calling this function. This can be done using the
-/// `prepare_buffer` function later in this file.
-unsafe fn poll_read_to_end<R: AsyncRead + ?Sized>(
+fn poll_read_to_end<R: AsyncRead + ?Sized>(
     buf: &mut Vec<u8>,
     read: Pin<&mut R>,
     cx: &mut Context<'_>,
@@ -73,70 +61,32 @@ unsafe fn poll_read_to_end<R: AsyncRead + ?Sized>(
     // of data to return. Simply tacking on an extra DEFAULT_BUF_SIZE space every
     // time is 4,500 times (!) slower than this if the reader has a very small
     // amount of data to return.
-    reserve(buf, &*read, 32);
-
-    let unused_capacity: &mut [MaybeUninit<u8>] = get_unused_capacity(buf);
-
-    // safety: The buffer has been prepared for use with the AsyncRead before
-    // calling this function.
-    let slice: &mut [u8] = &mut *(unused_capacity as *mut [MaybeUninit<u8>] as *mut [u8]);
-
-    let res = ready!(read.poll_read(cx, slice));
-    if let Ok(num) = res {
-        // safety: There are two situations:
-        //
-        // 1. The AsyncRead has not overriden `prepare_uninitialized_buffer`.
-        //
-        // In this situation, the default implementation of that method will have
-        // zeroed the unused capacity. This means that setting the length will
-        // never expose uninitialized memory in the vector.
-        //
-        // Note that the assert! below ensures that we don't set the length to
-        // something larger than the capacity, which malicious implementors might
-        // try to have us do.
-        //
-        // 2. The AsyncRead has overriden `prepare_uninitialized_buffer`.
-        //
-        // In this case, the safety of the `set_len` call below relies on this
-        // guarantee from the documentation on `prepare_uninitialized_buffer`:
-        //
-        // > This function isn't actually unsafe to call but unsafe to implement.
-        // > The implementer must ensure that either the whole buf has been zeroed
-        // > or poll_read() overwrites the buffer without reading it and returns
-        // > correct value.
-        //
-        // Note that `prepare_uninitialized_buffer` is unsafe to implement, so this
-        // is a guarantee we can rely on in unsafe code.
-        //
-        // The assert!() is technically only necessary in the first case.
-        let new_len = buf.len() + num;
-        assert!(new_len <= buf.capacity());
+    reserve(buf, 32);
 
-        buf.set_len(new_len);
-    }
-    Poll::Ready(res)
-}
+    let mut unused_capacity = ReadBuf::uninit(get_unused_capacity(buf));
 
-/// This function prepares the unused capacity for use with the provided AsyncRead.
-pub(super) fn prepare_buffer<R: AsyncRead + ?Sized>(buf: &mut Vec<u8>, read: &R) {
-    let buffer = get_unused_capacity(buf);
+    ready!(read.poll_read(cx, &mut unused_capacity))?;
 
-    // safety: This function is only unsafe to implement.
+    let n = unused_capacity.filled().len();
+    let new_len = buf.len() + n;
+
+    // This should no longer even be possible in safe Rust. An implementor
+    // would need to have unsafely *replaced* the buffer inside `ReadBuf`,
+    // which... yolo?
+    assert!(new_len <= buf.capacity());
     unsafe {
-        read.prepare_uninitialized_buffer(buffer);
+        buf.set_len(new_len);
     }
+    Poll::Ready(Ok(n))
 }
 
 /// Allocates more memory and ensures that the unused capacity is prepared for use
 /// with the `AsyncRead`.
-fn reserve<R: AsyncRead + ?Sized>(buf: &mut Vec<u8>, read: &R, bytes: usize) {
+fn reserve(buf: &mut Vec<u8>, bytes: usize) {
     if buf.capacity() - buf.len() >= bytes {
         return;
     }
     buf.reserve(bytes);
-    // The call above has reallocated the buffer, so we must reinitialize the entire
-    // unused capacity, even if we already initialized some of it before the resize.
-    prepare_buffer(buf, read);
 }
 
 /// Returns the unused capacity of the provided vector.
@@ -153,8 +103,7 @@ where
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let Self { reader, buf, read } = &mut *self;
 
-        // safety: The constructor of ReadToEnd calls `prepare_buffer`
-        unsafe { read_to_end_internal(buf, Pin::new(*reader), read, cx) }
+        read_to_end_internal(buf, Pin::new(*reader), read, cx)
     }
 }