Skip to content

Commit c393236

Browse files
seanmonstarseanmonstar
authored
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 parent 71da060 commit c393236

40 files changed

+626
-544
lines changed

‎tokio-test/src/io.rs

+24-19
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
//! [`AsyncRead`]: tokio::io::AsyncRead
1919
//! [`AsyncWrite`]: tokio::io::AsyncWrite
2020
21-
use tokio::io::{AsyncRead, AsyncWrite};
21+
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
2222
use tokio::sync::mpsc;
2323
use tokio::time::{self, Delay, Duration, Instant};
2424

@@ -204,20 +204,19 @@ impl Inner {
204204
self.rx.poll_recv(cx)
205205
}
206206

207-
fn read(&mut self, dst: &mut [u8]) -> io::Result<usize> {
207+
fn read(&mut self, dst: &mut ReadBuf<'_>) -> io::Result<()> {
208208
match self.action() {
209209
Some(&mut Action::Read(ref mut data)) => {
210210
// Figure out how much to copy
211-
let n = cmp::min(dst.len(), data.len());
211+
let n = cmp::min(dst.remaining(), data.len());
212212

213213
// Copy the data into the `dst` slice
214-
(&mut dst[..n]).copy_from_slice(&data[..n]);
214+
dst.append(&data[..n]);
215215

216216
// Drain the data from the source
217217
data.drain(..n);
218218

219-
// Return the number of bytes read
220-
Ok(n)
219+
Ok(())
221220
}
222221
Some(&mut Action::ReadError(ref mut err)) => {
223222
// As the
@@ -229,7 +228,7 @@ impl Inner {
229228
// Either waiting or expecting a write
230229
Err(io::ErrorKind::WouldBlock.into())
231230
}
232-
None => Ok(0),
231+
None => Ok(()),
233232
}
234233
}
235234

@@ -348,8 +347,8 @@ impl AsyncRead for Mock {
348347
fn poll_read(
349348
mut self: Pin<&mut Self>,
350349
cx: &mut task::Context<'_>,
351-
buf: &mut [u8],
352-
) -> Poll<io::Result<usize>> {
350+
buf: &mut ReadBuf<'_>,
351+
) -> Poll<io::Result<()>> {
353352
loop {
354353
if let Some(ref mut sleep) = self.inner.sleep {
355354
ready!(Pin::new(sleep).poll(cx));
@@ -358,6 +357,9 @@ impl AsyncRead for Mock {
358357
// If a sleep is set, it has already fired
359358
self.inner.sleep = None;
360359

360+
// Capture 'filled' to monitor if it changed
361+
let filled = buf.filled().len();
362+
361363
match self.inner.read(buf) {
362364
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
363365
if let Some(rem) = self.inner.remaining_wait() {
@@ -368,19 +370,22 @@ impl AsyncRead for Mock {
368370
return Poll::Pending;
369371
}
370372
}
371-
Ok(0) => {
372-
// TODO: Extract
373-
match ready!(self.inner.poll_action(cx)) {
374-
Some(action) => {
375-
self.inner.actions.push_back(action);
376-
continue;
377-
}
378-
None => {
379-
return Poll::Ready(Ok(0));
373+
Ok(()) => {
374+
if buf.filled().len() == filled {
375+
match ready!(self.inner.poll_action(cx)) {
376+
Some(action) => {
377+
self.inner.actions.push_back(action);
378+
continue;
379+
}
380+
None => {
381+
return Poll::Ready(Ok(()));
382+
}
380383
}
384+
} else {
385+
return Poll::Ready(Ok(()));
381386
}
382387
}
383-
ret => return Poll::Ready(ret),
388+
Err(e) => return Poll::Ready(Err(e)),
384389
}
385390
}
386391
}

‎tokio-util/src/compat.rs

+21-5
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
//! Compatibility between the `tokio::io` and `futures-io` versions of the
22
//! `AsyncRead` and `AsyncWrite` traits.
3+
use futures_core::ready;
34
use pin_project_lite::pin_project;
45
use std::io;
56
use std::pin::Pin;
@@ -107,9 +108,18 @@ where
107108
fn poll_read(
108109
self: Pin<&mut Self>,
109110
cx: &mut Context<'_>,
110-
buf: &mut [u8],
111-
) -> Poll<io::Result<usize>> {
112-
futures_io::AsyncRead::poll_read(self.project().inner, cx, buf)
111+
buf: &mut tokio::io::ReadBuf<'_>,
112+
) -> Poll<io::Result<()>> {
113+
// We can't trust the inner type to not peak at the bytes,
114+
// so we must defensively initialize the buffer.
115+
let slice = buf.initialize_unfilled();
116+
let n = ready!(futures_io::AsyncRead::poll_read(
117+
self.project().inner,
118+
cx,
119+
slice
120+
))?;
121+
buf.add_filled(n);
122+
Poll::Ready(Ok(()))
113123
}
114124
}
115125

@@ -120,9 +130,15 @@ where
120130
fn poll_read(
121131
self: Pin<&mut Self>,
122132
cx: &mut Context<'_>,
123-
buf: &mut [u8],
133+
slice: &mut [u8],
124134
) -> Poll<io::Result<usize>> {
125-
tokio::io::AsyncRead::poll_read(self.project().inner, cx, buf)
135+
let mut buf = tokio::io::ReadBuf::new(slice);
136+
ready!(tokio::io::AsyncRead::poll_read(
137+
self.project().inner,
138+
cx,
139+
&mut buf
140+
))?;
141+
Poll::Ready(Ok(buf.filled().len()))
126142
}
127143
}
128144

‎tokio-util/tests/framed.rs

+2-2
Original file line numberDiff line numberDiff line change
@@ -55,8 +55,8 @@ impl AsyncRead for DontReadIntoThis {
5555
fn poll_read(
5656
self: Pin<&mut Self>,
5757
_cx: &mut Context<'_>,
58-
_buf: &mut [u8],
59-
) -> Poll<io::Result<usize>> {
58+
_buf: &mut tokio::io::ReadBuf<'_>,
59+
) -> Poll<io::Result<()>> {
6060
unreachable!()
6161
}
6262
}

‎tokio-util/tests/framed_read.rs

+9-9
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
#![warn(rust_2018_idioms)]
22

3-
use tokio::io::AsyncRead;
3+
use tokio::io::{AsyncRead, ReadBuf};
44
use tokio_test::assert_ready;
55
use tokio_test::task;
66
use tokio_util::codec::{Decoder, FramedRead};
@@ -264,19 +264,19 @@ impl AsyncRead for Mock {
264264
fn poll_read(
265265
mut self: Pin<&mut Self>,
266266
_cx: &mut Context<'_>,
267-
buf: &mut [u8],
268-
) -> Poll<io::Result<usize>> {
267+
buf: &mut ReadBuf<'_>,
268+
) -> Poll<io::Result<()>> {
269269
use io::ErrorKind::WouldBlock;
270270

271271
match self.calls.pop_front() {
272272
Some(Ok(data)) => {
273-
debug_assert!(buf.len() >= data.len());
274-
buf[..data.len()].copy_from_slice(&data[..]);
275-
Ready(Ok(data.len()))
273+
debug_assert!(buf.remaining() >= data.len());
274+
buf.append(&data);
275+
Ready(Ok(()))
276276
}
277277
Some(Err(ref e)) if e.kind() == WouldBlock => Pending,
278278
Some(Err(e)) => Ready(Err(e)),
279-
None => Ready(Ok(0)),
279+
None => Ready(Ok(())),
280280
}
281281
}
282282
}
@@ -288,8 +288,8 @@ impl AsyncRead for Slice<'_> {
288288
fn poll_read(
289289
mut self: Pin<&mut Self>,
290290
cx: &mut Context<'_>,
291-
buf: &mut [u8],
292-
) -> Poll<io::Result<usize>> {
291+
buf: &mut ReadBuf<'_>,
292+
) -> Poll<io::Result<()>> {
293293
Pin::new(&mut self.0).poll_read(cx, buf)
294294
}
295295
}

‎tokio-util/tests/length_delimited.rs

+7-7
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
#![warn(rust_2018_idioms)]
22

3-
use tokio::io::{AsyncRead, AsyncWrite};
3+
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
44
use tokio_test::task;
55
use tokio_test::{
66
assert_err, assert_ok, assert_pending, assert_ready, assert_ready_err, assert_ready_ok,
@@ -707,18 +707,18 @@ impl AsyncRead for Mock {
707707
fn poll_read(
708708
mut self: Pin<&mut Self>,
709709
_cx: &mut Context<'_>,
710-
dst: &mut [u8],
711-
) -> Poll<io::Result<usize>> {
710+
dst: &mut ReadBuf<'_>,
711+
) -> Poll<io::Result<()>> {
712712
match self.calls.pop_front() {
713713
Some(Ready(Ok(Op::Data(data)))) => {
714-
debug_assert!(dst.len() >= data.len());
715-
dst[..data.len()].copy_from_slice(&data[..]);
716-
Ready(Ok(data.len()))
714+
debug_assert!(dst.remaining() >= data.len());
715+
dst.append(&data);
716+
Ready(Ok(()))
717717
}
718718
Some(Ready(Ok(_))) => panic!(),
719719
Some(Ready(Err(e))) => Ready(Err(e)),
720720
Some(Pending) => Pending,
721-
None => Ready(Ok(0)),
721+
None => Ready(Ok(())),
722722
}
723723
}
724724
}

‎tokio/src/fs/file.rs

+7-12
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
use self::State::*;
66
use crate::fs::{asyncify, sys};
77
use crate::io::blocking::Buf;
8-
use crate::io::{AsyncRead, AsyncSeek, AsyncWrite};
8+
use crate::io::{AsyncRead, AsyncSeek, AsyncWrite, ReadBuf};
99

1010
use std::fmt;
1111
use std::fs::{Metadata, Permissions};
@@ -537,25 +537,20 @@ impl File {
537537
}
538538

539539
impl AsyncRead for File {
540-
unsafe fn prepare_uninitialized_buffer(&self, _buf: &mut [std::mem::MaybeUninit<u8>]) -> bool {
541-
// https://github.com/rust-lang/rust/blob/09c817eeb29e764cfc12d0a8d94841e3ffe34023/src/libstd/fs.rs#L668
542-
false
543-
}
544-
545540
fn poll_read(
546541
mut self: Pin<&mut Self>,
547542
cx: &mut Context<'_>,
548-
dst: &mut [u8],
549-
) -> Poll<io::Result<usize>> {
543+
dst: &mut ReadBuf<'_>,
544+
) -> Poll<io::Result<()>> {
550545
loop {
551546
match self.state {
552547
Idle(ref mut buf_cell) => {
553548
let mut buf = buf_cell.take().unwrap();
554549

555550
if !buf.is_empty() {
556-
let n = buf.copy_to(dst);
551+
buf.copy_to(dst);
557552
*buf_cell = Some(buf);
558-
return Ready(Ok(n));
553+
return Ready(Ok(()));
559554
}
560555

561556
buf.ensure_capacity_for(dst);
@@ -571,9 +566,9 @@ impl AsyncRead for File {
571566

572567
match op {
573568
Operation::Read(Ok(_)) => {
574-
let n = buf.copy_to(dst);
569+
buf.copy_to(dst);
575570
self.state = Idle(Some(buf));
576-
return Ready(Ok(n));
571+
return Ready(Ok(()));
577572
}
578573
Operation::Read(Err(e)) => {
579574
assert!(buf.is_empty());

0 commit comments

Comments
 (0)