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.

1 files changed, 24 insertions, 19 deletions

diff --git a/tokio-test/src/io.rs b/tokio-test/src/io.rs
index 26ef57e4..f1ce77aa 100644
--- 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)),
             }
         }
     }