diff --git a/tokio/src/io/async_read_ext.rs b/tokio/src/io/async_read_ext.rs index 2fc2cbb4064..f04935c51a6 100644 --- a/tokio/src/io/async_read_ext.rs +++ b/tokio/src/io/async_read_ext.rs @@ -2,6 +2,7 @@ use crate::io::copy::{copy, Copy}; use crate::io::read::{read, Read}; use crate::io::read_exact::{read_exact, ReadExact}; use crate::io::read_to_end::{read_to_end, ReadToEnd}; +use crate::io::read_to_string::{read_to_string, ReadToString}; use tokio_io::{AsyncRead, AsyncWrite}; @@ -63,12 +64,36 @@ pub trait AsyncReadExt: AsyncRead { } /// Read all bytes until EOF in this source, placing them into `dst`. + /// + /// On success the total number of bytes read is returned. + /// + /// # Examples + /// + /// ``` + /// unimplemented!(); + /// ``` fn read_to_end<'a>(&'a mut self, dst: &'a mut Vec) -> ReadToEnd<'a, Self> where Self: Unpin, { read_to_end(self, dst) } + + /// Read all bytes until EOF in this source, placing them into `dst`. + /// + /// On success the total number of bytes read is returned. + /// + /// # Examples + /// + /// ``` + /// unimplemented!(); + /// ``` + fn read_to_string<'a>(&'a mut self, dst: &'a mut String) -> ReadToString<'a, Self> + where + Self: Unpin, + { + read_to_string(self, dst) + } } impl AsyncReadExt for R {} diff --git a/tokio/src/io/mod.rs b/tokio/src/io/mod.rs index 837ab5aa8bf..d1fce45ef4e 100644 --- a/tokio/src/io/mod.rs +++ b/tokio/src/io/mod.rs @@ -45,6 +45,7 @@ mod read; mod read_exact; mod read_line; mod read_to_end; +mod read_to_string; mod read_until; mod write; mod write_all; diff --git a/tokio/src/io/read_to_string.rs b/tokio/src/io/read_to_string.rs new file mode 100644 index 00000000000..b9037443efb --- /dev/null +++ b/tokio/src/io/read_to_string.rs @@ -0,0 +1,72 @@ +use super::read_to_end::read_to_end_internal; +use futures_core::ready; +use std::future::Future; +use std::pin::Pin; +use std::task::{Context, Poll}; +use std::{io, mem, str}; +use tokio_io::AsyncRead; + +/// Future for the [`read_to_string`](super::AsyncReadExt::read_to_string) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct ReadToString<'a, R: ?Sized + Unpin> { + reader: &'a mut R, + buf: &'a mut String, + bytes: Vec, + start_len: usize, +} + +impl Unpin for ReadToString<'_, R> {} + +pub(crate) fn read_to_string<'a, R>(reader: &'a mut R, buf: &'a mut String) -> ReadToString<'a, R> +where + R: AsyncRead + ?Sized + Unpin, +{ + let start_len = buf.len(); + ReadToString { + reader, + bytes: unsafe { mem::replace(buf.as_mut_vec(), Vec::new()) }, + buf, + start_len, + } +} + +fn read_to_string_internal( + reader: Pin<&mut R>, + cx: &mut Context<'_>, + buf: &mut String, + bytes: &mut Vec, + start_len: usize, +) -> Poll> { + let ret = ready!(read_to_end_internal(reader, cx, bytes, start_len)); + if str::from_utf8(&bytes).is_err() { + Poll::Ready(ret.and_then(|_| { + Err(io::Error::new( + io::ErrorKind::InvalidData, + "stream did not contain valid UTF-8", + )) + })) + } else { + debug_assert!(buf.is_empty()); + // Safety: `bytes` is a valid UTF-8 because `str::from_utf8` returned `Ok`. + mem::swap(unsafe { buf.as_mut_vec() }, bytes); + Poll::Ready(ret) + } +} + +impl Future for ReadToString<'_, A> +where + A: AsyncRead + ?Sized + Unpin, +{ + type Output = io::Result; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let Self { + reader, + buf, + bytes, + start_len, + } = &mut *self; + read_to_string_internal(Pin::new(reader), cx, buf, bytes, *start_len) + } +} diff --git a/tokio/tests/io_read_to_string.rs b/tokio/tests/io_read_to_string.rs new file mode 100644 index 00000000000..4bbee177e68 --- /dev/null +++ b/tokio/tests/io_read_to_string.rs @@ -0,0 +1,40 @@ +#![deny(warnings, rust_2018_idioms)] +#![feature(async_await)] + +use tokio::io::{AsyncRead, AsyncReadExt}; +use tokio_test::assert_ok; + +use std::pin::Pin; +use std::task::{Context, Poll}; +use std::{cmp, io}; + +#[tokio::test] +async fn read_to_string() { + struct Rd { + val: &'static [u8], + } + + impl AsyncRead for Rd { + fn poll_read( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &mut [u8], + ) -> Poll> { + let me = &mut *self; + let len = cmp::min(buf.len(), me.val.len()); + + buf[..len].copy_from_slice(&me.val[..len]); + me.val = &me.val[len..]; + Poll::Ready(Ok(len)) + } + } + + let mut buf = String::new(); + let mut rd = Rd { + val: b"hello world", + }; + + let n = assert_ok!(rd.read_to_string(&mut buf).await); + assert_eq!(n, 11); + assert_eq!(buf[..], "hello world"[..]); +}