io: change AsyncRead to use a ReadBuf (#2758)

Works towards #2716. Changes the argument to `AsyncRead::poll_read` to take a `ReadBuf` struct that safely manages writes to uninitialized memory.
tokio-rs · Aug 14, 2020 · c393236 · c393236
1 parent 71da060
commit c393236
Show file tree

Hide file tree

Showing 40 changed files with 626 additions and 544 deletions.
diff --git a/tokio-test/src/io.rs b/tokio-test/src/io.rs
@@ -18,7 +18,7 @@
 //! [`AsyncRead`]: tokio::io::AsyncRead
 //! [`AsyncWrite`]: tokio::io::AsyncWrite
 
-use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 use tokio::sync::mpsc;
 use tokio::time::{self, Delay, Duration, Instant};
 
@@ -204,20 +204,19 @@ impl Inner {
         self.rx.poll_recv(cx)
     }
 
-    fn read(&mut self, dst: &mut [u8]) -> io::Result<usize> {
+    fn read(&mut self, dst: &mut ReadBuf<'_>) -> io::Result<()> {
         match self.action() {
             Some(&mut Action::Read(ref mut data)) => {
                 // Figure out how much to copy
-                let n = cmp::min(dst.len(), data.len());
+                let n = cmp::min(dst.remaining(), data.len());
 
                 // Copy the data into the `dst` slice
-                (&mut dst[..n]).copy_from_slice(&data[..n]);
+                dst.append(&data[..n]);
 
                 // Drain the data from the source
                 data.drain(..n);
 
-                // Return the number of bytes read
-                Ok(n)
+                Ok(())
             }
             Some(&mut Action::ReadError(ref mut err)) => {
                 // As the
@@ -229,7 +228,7 @@ impl Inner {
                 // Either waiting or expecting a write
                 Err(io::ErrorKind::WouldBlock.into())
             }
-            None => Ok(0),
+            None => Ok(()),
         }
     }
 
@@ -348,8 +347,8 @@ impl AsyncRead for Mock {
     fn poll_read(
         mut self: Pin<&mut Self>,
         cx: &mut task::Context<'_>,
-        buf: &mut [u8],
-    ) -> Poll<io::Result<usize>> {
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
         loop {
             if let Some(ref mut sleep) = self.inner.sleep {
                 ready!(Pin::new(sleep).poll(cx));
@@ -358,6 +357,9 @@ impl AsyncRead for Mock {
             // If a sleep is set, it has already fired
             self.inner.sleep = None;
 
+            // Capture 'filled' to monitor if it changed
+            let filled = buf.filled().len();
+
             match self.inner.read(buf) {
                 Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                     if let Some(rem) = self.inner.remaining_wait() {
@@ -368,19 +370,22 @@ impl AsyncRead for Mock {
                         return Poll::Pending;
                     }
                 }
-                Ok(0) => {
-                    // TODO: Extract
-                    match ready!(self.inner.poll_action(cx)) {
-                        Some(action) => {
-                            self.inner.actions.push_back(action);
-                            continue;
-                        }
-                        None => {
-                            return Poll::Ready(Ok(0));
+                Ok(()) => {
+                    if buf.filled().len() == filled {
+                        match ready!(self.inner.poll_action(cx)) {
+                            Some(action) => {
+                                self.inner.actions.push_back(action);
+                                continue;
+                            }
+                            None => {
+                                return Poll::Ready(Ok(()));
+                            }
                         }
+                    } else {
+                        return Poll::Ready(Ok(()));
                     }
                 }
-                ret => return Poll::Ready(ret),
+                Err(e) => return Poll::Ready(Err(e)),
             }
         }
     }

diff --git a/tokio-util/src/compat.rs b/tokio-util/src/compat.rs
@@ -1,5 +1,6 @@
 //! Compatibility between the `tokio::io` and `futures-io` versions of the
 //! `AsyncRead` and `AsyncWrite` traits.
+use futures_core::ready;
 use pin_project_lite::pin_project;
 use std::io;
 use std::pin::Pin;
@@ -107,9 +108,18 @@ where
     fn poll_read(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
-        buf: &mut [u8],
-    ) -> Poll<io::Result<usize>> {
-        futures_io::AsyncRead::poll_read(self.project().inner, cx, buf)
+        buf: &mut tokio::io::ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        // We can't trust the inner type to not peak at the bytes,
+        // so we must defensively initialize the buffer.
+        let slice = buf.initialize_unfilled();
+        let n = ready!(futures_io::AsyncRead::poll_read(
+            self.project().inner,
+            cx,
+            slice
+        ))?;
+        buf.add_filled(n);
+        Poll::Ready(Ok(()))
     }
 }
 
@@ -120,9 +130,15 @@ where
     fn poll_read(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
-        buf: &mut [u8],
+        slice: &mut [u8],
     ) -> Poll<io::Result<usize>> {
-        tokio::io::AsyncRead::poll_read(self.project().inner, cx, buf)
+        let mut buf = tokio::io::ReadBuf::new(slice);
+        ready!(tokio::io::AsyncRead::poll_read(
+            self.project().inner,
+            cx,
+            &mut buf
+        ))?;
+        Poll::Ready(Ok(buf.filled().len()))
     }
 }
 

diff --git a/tokio-util/tests/framed.rs b/tokio-util/tests/framed.rs
@@ -55,8 +55,8 @@ impl AsyncRead for DontReadIntoThis {
     fn poll_read(
         self: Pin<&mut Self>,
         _cx: &mut Context<'_>,
-        _buf: &mut [u8],
-    ) -> Poll<io::Result<usize>> {
+        _buf: &mut tokio::io::ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
         unreachable!()
     }
 }

diff --git a/tokio-util/tests/framed_read.rs b/tokio-util/tests/framed_read.rs
@@ -1,6 +1,6 @@
 #![warn(rust_2018_idioms)]
 
-use tokio::io::AsyncRead;
+use tokio::io::{AsyncRead, ReadBuf};
 use tokio_test::assert_ready;
 use tokio_test::task;
 use tokio_util::codec::{Decoder, FramedRead};
@@ -264,19 +264,19 @@ impl AsyncRead for Mock {
     fn poll_read(
         mut self: Pin<&mut Self>,
         _cx: &mut Context<'_>,
-        buf: &mut [u8],
-    ) -> Poll<io::Result<usize>> {
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
         use io::ErrorKind::WouldBlock;
 
         match self.calls.pop_front() {
             Some(Ok(data)) => {
-                debug_assert!(buf.len() >= data.len());
-                buf[..data.len()].copy_from_slice(&data[..]);
-                Ready(Ok(data.len()))
+                debug_assert!(buf.remaining() >= data.len());
+                buf.append(&data);
+                Ready(Ok(()))
             }
             Some(Err(ref e)) if e.kind() == WouldBlock => Pending,
             Some(Err(e)) => Ready(Err(e)),
-            None => Ready(Ok(0)),
+            None => Ready(Ok(())),
         }
     }
 }
@@ -288,8 +288,8 @@ impl AsyncRead for Slice<'_> {
     fn poll_read(
         mut self: Pin<&mut Self>,
         cx: &mut Context<'_>,
-        buf: &mut [u8],
-    ) -> Poll<io::Result<usize>> {
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
         Pin::new(&mut self.0).poll_read(cx, buf)
     }
 }
diff --git a/tokio-util/tests/length_delimited.rs b/tokio-util/tests/length_delimited.rs
@@ -1,6 +1,6 @@
 #![warn(rust_2018_idioms)]
 
-use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 use tokio_test::task;
 use tokio_test::{
     assert_err, assert_ok, assert_pending, assert_ready, assert_ready_err, assert_ready_ok,
@@ -707,18 +707,18 @@ impl AsyncRead for Mock {
     fn poll_read(
         mut self: Pin<&mut Self>,
         _cx: &mut Context<'_>,
-        dst: &mut [u8],
-    ) -> Poll<io::Result<usize>> {
+        dst: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
         match self.calls.pop_front() {
             Some(Ready(Ok(Op::Data(data)))) => {
-                debug_assert!(dst.len() >= data.len());
-                dst[..data.len()].copy_from_slice(&data[..]);
-                Ready(Ok(data.len()))
+                debug_assert!(dst.remaining() >= data.len());
+                dst.append(&data);
+                Ready(Ok(()))
             }
             Some(Ready(Ok(_))) => panic!(),
             Some(Ready(Err(e))) => Ready(Err(e)),
             Some(Pending) => Pending,
-            None => Ready(Ok(0)),
+            None => Ready(Ok(())),
         }
     }
 }

diff --git a/tokio/src/fs/file.rs b/tokio/src/fs/file.rs
@@ -5,7 +5,7 @@
 use self::State::*;
 use crate::fs::{asyncify, sys};
 use crate::io::blocking::Buf;
-use crate::io::{AsyncRead, AsyncSeek, AsyncWrite};
+use crate::io::{AsyncRead, AsyncSeek, AsyncWrite, ReadBuf};
 
 use std::fmt;
 use std::fs::{Metadata, Permissions};
@@ -537,25 +537,20 @@ impl File {
 }
 
 impl AsyncRead for File {
-    unsafe fn prepare_uninitialized_buffer(&self, _buf: &mut [std::mem::MaybeUninit<u8>]) -> bool {
-        // https://github.com/rust-lang/rust/blob/09c817eeb29e764cfc12d0a8d94841e3ffe34023/src/libstd/fs.rs#L668
-        false
-    }
-
     fn poll_read(
         mut self: Pin<&mut Self>,
         cx: &mut Context<'_>,
-        dst: &mut [u8],
-    ) -> Poll<io::Result<usize>> {
+        dst: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
         loop {
             match self.state {
                 Idle(ref mut buf_cell) => {
                     let mut buf = buf_cell.take().unwrap();
 
                     if !buf.is_empty() {
-                        let n = buf.copy_to(dst);
+                        buf.copy_to(dst);
                         *buf_cell = Some(buf);
-                        return Ready(Ok(n));
+                        return Ready(Ok(()));
                     }
 
                     buf.ensure_capacity_for(dst);
@@ -571,9 +566,9 @@ impl AsyncRead for File {
 
                     match op {
                         Operation::Read(Ok(_)) => {
-                            let n = buf.copy_to(dst);
+                            buf.copy_to(dst);
                             self.state = Idle(Some(buf));
-                            return Ready(Ok(n));
+                            return Ready(Ok(()));
                         }
                         Operation::Read(Err(e)) => {
                             assert!(buf.is_empty());