From c393236dfd12c13e82badd631d3a3a90481c6f95 Mon Sep 17 00:00:00 2001
From: Sean McArthur <sean@seanmonstar.com>
Date: Thu, 13 Aug 2020 20:15:01 -0700
Subject: 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-util/src/compat.rs             | 26 +++++++++++++++++++++-----
 tokio-util/tests/framed.rs           |  4 ++--
 tokio-util/tests/framed_read.rs      | 18 +++++++++---------
 tokio-util/tests/length_delimited.rs | 14 +++++++-------
 4 files changed, 39 insertions(+), 23 deletions(-)

(limited to 'tokio-util')

diff --git a/tokio-util/src/compat.rs b/tokio-util/src/compat.rs
index 769e30c2..34120d43 100644
--- 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
index d7ee3ef5..4c5f8418 100644
--- 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
index 27bb298a..da38c432 100644
--- 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
index 734cd834..9f615412 100644
--- 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(())),
         }
     }
 }
-- 
cgit v1.2.3