use core::future::Future;
use core::pin::Pin;
use core::task::{Context, Poll};
use std::io;
use super::{AsyncRead, AsyncWrite, Stream};
const READ_TO_END_CHUNK: usize = 8192;
const COPY_BUF_SIZE: usize = 8192;
pub trait AsyncReadExt: AsyncRead {
fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> Read<'a, Self>
where
Self: Unpin,
{
Read { reader: self, buf }
}
fn read_exact<'a>(&'a mut self, buf: &'a mut [u8]) -> ReadExact<'a, Self>
where
Self: Unpin,
{
ReadExact {
reader: self,
buf,
filled: 0,
}
}
fn read_to_end<'a>(&'a mut self, buf: &'a mut Vec<u8>) -> ReadToEnd<'a, Self>
where
Self: Unpin,
{
ReadToEnd {
reader: self,
buf,
chunk: vec![0; READ_TO_END_CHUNK],
start_len: 0,
initialized: false,
}
}
fn lines(self) -> Lines<Self>
where
Self: Sized,
{
Lines {
reader: self,
buf: Vec::new(),
chunk: vec![0; READ_TO_END_CHUNK],
eof: false,
}
}
}
impl<R: AsyncRead + ?Sized> AsyncReadExt for R {}
pub trait AsyncWriteExt: AsyncWrite {
fn write<'a>(&'a mut self, buf: &'a [u8]) -> Write<'a, Self>
where
Self: Unpin,
{
Write { writer: self, buf }
}
fn write_all<'a>(&'a mut self, buf: &'a [u8]) -> WriteAll<'a, Self>
where
Self: Unpin,
{
WriteAll {
writer: self,
buf,
written: 0,
}
}
fn flush(&mut self) -> Flush<'_, Self>
where
Self: Unpin,
{
Flush { writer: self }
}
fn close(&mut self) -> Close<'_, Self>
where
Self: Unpin,
{
Close { writer: self }
}
}
impl<W: AsyncWrite + ?Sized> AsyncWriteExt for W {}
pub fn copy<'a, R, W>(reader: &'a mut R, writer: &'a mut W) -> Copy<'a, R, W>
where
R: AsyncRead + ?Sized,
W: AsyncWrite + ?Sized,
{
Copy {
reader,
writer,
buf: vec![0; COPY_BUF_SIZE],
pos: 0,
filled: 0,
copied: 0,
state: CopyState::Copying,
}
}
pub fn copy_bidirectional<'a, A, B>(a: &'a mut A, b: &'a mut B) -> CopyBidirectional<'a, A, B>
where
A: AsyncRead + AsyncWrite + Unpin + ?Sized,
B: AsyncRead + AsyncWrite + Unpin + ?Sized,
{
CopyBidirectional {
a,
b,
a_to_b: CopyHalf::new(),
b_to_a: CopyHalf::new(),
}
}
#[must_use = "futures do nothing unless awaited or polled"]
pub struct Read<'a, R: ?Sized> {
reader: &'a mut R,
buf: &'a mut [u8],
}
impl<R: AsyncRead + Unpin + ?Sized> Future for Read<'_, R> {
type Output = io::Result<usize>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = &mut *self;
Pin::new(&mut *this.reader).poll_read(cx, this.buf)
}
}
#[must_use = "futures do nothing unless awaited or polled"]
pub struct ReadExact<'a, R: ?Sized> {
reader: &'a mut R,
buf: &'a mut [u8],
filled: usize,
}
impl<R: AsyncRead + Unpin + ?Sized> Future for ReadExact<'_, R> {
type Output = io::Result<()>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = &mut *self;
while this.filled < this.buf.len() {
let read = match Pin::new(&mut *this.reader).poll_read(cx, &mut this.buf[this.filled..])
{
Poll::Ready(result) => result?,
Poll::Pending => return Poll::Pending,
};
if read == 0 {
return Poll::Ready(Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
"failed to fill whole buffer",
)));
}
this.filled += read;
}
Poll::Ready(Ok(()))
}
}
#[must_use = "futures do nothing unless awaited or polled"]
pub struct ReadToEnd<'a, R: ?Sized> {
reader: &'a mut R,
buf: &'a mut Vec<u8>,
chunk: Vec<u8>,
start_len: usize,
initialized: bool,
}
impl<R: AsyncRead + Unpin + ?Sized> Future for ReadToEnd<'_, R> {
type Output = io::Result<usize>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = &mut *self;
if !this.initialized {
this.start_len = this.buf.len();
this.initialized = true;
}
loop {
let read = match Pin::new(&mut *this.reader).poll_read(cx, &mut this.chunk) {
Poll::Ready(result) => result?,
Poll::Pending => return Poll::Pending,
};
if read == 0 {
return Poll::Ready(Ok(this.buf.len() - this.start_len));
}
this.buf.extend_from_slice(&this.chunk[..read]);
}
}
}
#[must_use = "futures do nothing unless awaited or polled"]
pub struct Write<'a, W: ?Sized> {
writer: &'a mut W,
buf: &'a [u8],
}
impl<W: AsyncWrite + Unpin + ?Sized> Future for Write<'_, W> {
type Output = io::Result<usize>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = &mut *self;
Pin::new(&mut *this.writer).poll_write(cx, this.buf)
}
}
#[must_use = "futures do nothing unless awaited or polled"]
pub struct WriteAll<'a, W: ?Sized> {
writer: &'a mut W,
buf: &'a [u8],
written: usize,
}
impl<W: AsyncWrite + Unpin + ?Sized> Future for WriteAll<'_, W> {
type Output = io::Result<()>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = &mut *self;
while this.written < this.buf.len() {
let written =
match Pin::new(&mut *this.writer).poll_write(cx, &this.buf[this.written..]) {
Poll::Ready(result) => result?,
Poll::Pending => return Poll::Pending,
};
if written == 0 {
return Poll::Ready(Err(io::Error::new(
io::ErrorKind::WriteZero,
"failed to write whole buffer",
)));
}
this.written += written;
}
Poll::Ready(Ok(()))
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum CopyState {
Copying,
Flushing,
}
#[must_use = "futures do nothing unless awaited or polled"]
pub struct Copy<'a, R: ?Sized, W: ?Sized> {
reader: &'a mut R,
writer: &'a mut W,
buf: Vec<u8>,
pos: usize,
filled: usize,
copied: u64,
state: CopyState,
}
impl<R, W> Future for Copy<'_, R, W>
where
R: AsyncRead + Unpin + ?Sized,
W: AsyncWrite + Unpin + ?Sized,
{
type Output = io::Result<u64>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = &mut *self;
loop {
match this.state {
CopyState::Copying if this.pos < this.filled => {
let written = match Pin::new(&mut *this.writer)
.poll_write(cx, &this.buf[this.pos..this.filled])
{
Poll::Ready(result) => result?,
Poll::Pending => return Poll::Pending,
};
if written == 0 {
return Poll::Ready(Err(io::Error::new(
io::ErrorKind::WriteZero,
"failed to write copied data",
)));
}
this.pos += written;
this.copied += written as u64;
}
CopyState::Copying => {
let read = match Pin::new(&mut *this.reader).poll_read(cx, &mut this.buf) {
Poll::Ready(result) => result?,
Poll::Pending => return Poll::Pending,
};
this.pos = 0;
this.filled = read;
if this.filled == 0 {
this.state = CopyState::Flushing;
}
}
CopyState::Flushing => {
return match Pin::new(&mut *this.writer).poll_flush(cx) {
Poll::Ready(Ok(())) => Poll::Ready(Ok(this.copied)),
Poll::Ready(Err(error)) => Poll::Ready(Err(error)),
Poll::Pending => Poll::Pending,
};
}
}
}
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum CopyHalfState {
Copying,
Closing,
Done,
}
struct CopyHalf {
buf: Vec<u8>,
pos: usize,
filled: usize,
copied: u64,
state: CopyHalfState,
}
impl CopyHalf {
fn new() -> Self {
Self {
buf: vec![0; COPY_BUF_SIZE],
pos: 0,
filled: 0,
copied: 0,
state: CopyHalfState::Copying,
}
}
fn is_done(&self) -> bool {
self.state == CopyHalfState::Done
}
}
fn poll_copy_half<R, W>(
reader: &mut R,
writer: &mut W,
half: &mut CopyHalf,
cx: &mut Context<'_>,
) -> Poll<io::Result<()>>
where
R: AsyncRead + Unpin + ?Sized,
W: AsyncWrite + Unpin + ?Sized,
{
loop {
match half.state {
CopyHalfState::Copying if half.pos < half.filled => {
let written =
match Pin::new(&mut *writer).poll_write(cx, &half.buf[half.pos..half.filled]) {
Poll::Ready(result) => result?,
Poll::Pending => return Poll::Pending,
};
if written == 0 {
return Poll::Ready(Err(io::Error::new(
io::ErrorKind::WriteZero,
"failed to write copied data",
)));
}
half.pos += written;
half.copied += written as u64;
}
CopyHalfState::Copying => {
let read = match Pin::new(&mut *reader).poll_read(cx, &mut half.buf) {
Poll::Ready(result) => result?,
Poll::Pending => return Poll::Pending,
};
half.pos = 0;
half.filled = read;
if half.filled == 0 {
half.state = CopyHalfState::Closing;
}
}
CopyHalfState::Closing => {
match Pin::new(&mut *writer).poll_close(cx) {
Poll::Ready(result) => result?,
Poll::Pending => return Poll::Pending,
}
half.state = CopyHalfState::Done;
return Poll::Ready(Ok(()));
}
CopyHalfState::Done => return Poll::Ready(Ok(())),
}
}
}
#[must_use = "futures do nothing unless awaited or polled"]
pub struct CopyBidirectional<'a, A: ?Sized, B: ?Sized> {
a: &'a mut A,
b: &'a mut B,
a_to_b: CopyHalf,
b_to_a: CopyHalf,
}
impl<A, B> Future for CopyBidirectional<'_, A, B>
where
A: AsyncRead + AsyncWrite + Unpin + ?Sized,
B: AsyncRead + AsyncWrite + Unpin + ?Sized,
{
type Output = io::Result<(u64, u64)>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = &mut *self;
if !this.a_to_b.is_done() {
match poll_copy_half(&mut *this.a, &mut *this.b, &mut this.a_to_b, cx) {
Poll::Ready(Ok(())) => {}
Poll::Ready(Err(error)) => return Poll::Ready(Err(error)),
Poll::Pending => {}
}
}
if !this.b_to_a.is_done() {
match poll_copy_half(&mut *this.b, &mut *this.a, &mut this.b_to_a, cx) {
Poll::Ready(Ok(())) => {}
Poll::Ready(Err(error)) => return Poll::Ready(Err(error)),
Poll::Pending => {}
}
}
if this.a_to_b.is_done() && this.b_to_a.is_done() {
Poll::Ready(Ok((this.a_to_b.copied, this.b_to_a.copied)))
} else {
Poll::Pending
}
}
}
#[must_use = "futures do nothing unless awaited or polled"]
pub struct Flush<'a, W: ?Sized> {
writer: &'a mut W,
}
impl<W: AsyncWrite + Unpin + ?Sized> Future for Flush<'_, W> {
type Output = io::Result<()>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
Pin::new(&mut *self.writer).poll_flush(cx)
}
}
#[must_use = "futures do nothing unless awaited or polled"]
pub struct Close<'a, W: ?Sized> {
writer: &'a mut W,
}
impl<W: AsyncWrite + Unpin + ?Sized> Future for Close<'_, W> {
type Output = io::Result<()>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
Pin::new(&mut *self.writer).poll_close(cx)
}
}
#[must_use = "streams do nothing unless polled"]
pub struct Lines<R> {
reader: R,
buf: Vec<u8>,
chunk: Vec<u8>,
eof: bool,
}
impl<R> Unpin for Lines<R> {}
impl<R: AsyncRead + Unpin> Stream for Lines<R> {
type Item = io::Result<String>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
loop {
if let Some(newline) = this.buf.iter().position(|byte| *byte == b'\n') {
let mut line = this.buf.drain(..=newline).collect::<Vec<_>>();
let _ = line.pop();
return Poll::Ready(Some(string_from_line(line)));
}
if this.eof {
if this.buf.is_empty() {
return Poll::Ready(None);
}
let line = core::mem::take(&mut this.buf);
return Poll::Ready(Some(string_from_line(line)));
}
match Pin::new(&mut this.reader).poll_read(cx, &mut this.chunk) {
Poll::Ready(Ok(0)) => this.eof = true,
Poll::Ready(Ok(read)) => this.buf.extend_from_slice(&this.chunk[..read]),
Poll::Ready(Err(error)) => return Poll::Ready(Some(Err(error))),
Poll::Pending => return Poll::Pending,
}
}
}
}
fn string_from_line(line: Vec<u8>) -> io::Result<String> {
String::from_utf8(line).map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))
}
#[cfg(test)]
mod tests {
use super::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, copy, copy_bidirectional};
use crate::io::StreamExt;
use crate::{run, spawn};
use core::pin::Pin;
use core::task::{Context, Poll};
use std::cell::RefCell;
use std::io;
use std::rc::Rc;
struct MemoryIo {
read: Vec<u8>,
read_pos: usize,
written: Rc<RefCell<Vec<u8>>>,
max_write: usize,
flushed: Rc<RefCell<bool>>,
closed: Rc<RefCell<bool>>,
}
impl MemoryIo {
fn new(read: &[u8], max_write: usize) -> Self {
Self {
read: read.to_vec(),
read_pos: 0,
written: Rc::new(RefCell::new(Vec::new())),
max_write,
flushed: Rc::new(RefCell::new(false)),
closed: Rc::new(RefCell::new(false)),
}
}
fn written(&self) -> Rc<RefCell<Vec<u8>>> {
Rc::clone(&self.written)
}
fn closed(&self) -> Rc<RefCell<bool>> {
Rc::clone(&self.closed)
}
}
impl AsyncRead for MemoryIo {
fn poll_read(
mut self: Pin<&mut Self>,
_cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
let available = self.read.len() - self.read_pos;
let read = buf.len().min(available).min(3);
let start = self.read_pos;
let end = start + read;
buf[..read].copy_from_slice(&self.read[start..end]);
self.read_pos = end;
Poll::Ready(Ok(read))
}
}
impl AsyncWrite for MemoryIo {
fn poll_write(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
if *self.closed.borrow() {
return Poll::Ready(Err(io::Error::new(
io::ErrorKind::BrokenPipe,
"writer is closed",
)));
}
let written = buf.len().min(self.max_write);
self.written.borrow_mut().extend_from_slice(&buf[..written]);
Poll::Ready(Ok(written))
}
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
*self.flushed.borrow_mut() = true;
Poll::Ready(Ok(()))
}
fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
*self.closed.borrow_mut() = true;
Poll::Ready(Ok(()))
}
}
#[test]
fn copy_writes_all_bytes_and_flushes() {
let written = Rc::new(RefCell::new(Vec::new()));
let flushed = Rc::new(RefCell::new(false));
{
let written = Rc::clone(&written);
let flushed = Rc::clone(&flushed);
spawn(async move {
let mut reader = MemoryIo::new(b"abcdefghi", 2);
let mut writer = MemoryIo::new(b"", 2);
writer.written = written;
writer.flushed = flushed;
let copied = copy(&mut reader, &mut writer).await.unwrap();
assert_eq!(copied, 9);
});
}
run();
assert_eq!(&*written.borrow(), b"abcdefghi");
assert!(*flushed.borrow());
}
#[test]
fn copy_bidirectional_copies_both_ways_and_closes_write_halves() {
let mut left = MemoryIo::new(b"left-to-right", 4);
let mut right = MemoryIo::new(b"right-to-left", 5);
let left_written = left.written();
let right_written = right.written();
let left_closed = left.closed();
let right_closed = right.closed();
spawn(async move {
let copied = copy_bidirectional(&mut left, &mut right).await.unwrap();
assert_eq!(copied, (13, 13));
});
run();
assert_eq!(&*right_written.borrow(), b"left-to-right");
assert_eq!(&*left_written.borrow(), b"right-to-left");
assert!(*left_closed.borrow());
assert!(*right_closed.borrow());
}
#[test]
fn read_exact_reports_unexpected_eof_on_short_read() {
let observed = Rc::new(RefCell::new(None::<io::ErrorKind>));
spawn({
let observed = Rc::clone(&observed);
async move {
let mut reader = MemoryIo::new(b"abc", 1);
let mut buf = [0u8; 4];
let error = reader
.read_exact(&mut buf)
.await
.expect_err("short read should fail");
*observed.borrow_mut() = Some(error.kind());
}
});
run();
assert_eq!(*observed.borrow(), Some(io::ErrorKind::UnexpectedEof));
}
#[test]
fn write_all_reports_write_zero() {
let observed = Rc::new(RefCell::new(None::<io::ErrorKind>));
spawn({
let observed = Rc::clone(&observed);
async move {
let mut writer = MemoryIo::new(b"", 0);
let error = writer
.write_all(b"cannot progress")
.await
.expect_err("zero write should fail");
*observed.borrow_mut() = Some(error.kind());
}
});
run();
assert_eq!(*observed.borrow(), Some(io::ErrorKind::WriteZero));
}
#[test]
fn lines_strip_lf_but_retain_cr_from_crlf() {
let observed = Rc::new(RefCell::new(Vec::new()));
spawn({
let observed = Rc::clone(&observed);
async move {
let reader = MemoryIo::new(b"alpha\r\nbeta\nunterminated", 1);
let lines = reader
.lines()
.collect::<Vec<_>>()
.await
.into_iter()
.collect::<io::Result<Vec<_>>>()
.unwrap();
*observed.borrow_mut() = lines;
}
});
run();
assert_eq!(
*observed.borrow(),
vec![
"alpha\r".to_string(),
"beta".to_string(),
"unterminated".to_string()
]
);
}
#[test]
fn copy_empty_reader_flushes_writer_without_writing() {
let flushed = Rc::new(RefCell::new(false));
let written = Rc::new(RefCell::new(Vec::new()));
spawn({
let flushed = Rc::clone(&flushed);
let written = Rc::clone(&written);
async move {
let mut reader = MemoryIo::new(b"", 1);
let mut writer = MemoryIo::new(b"", 1);
writer.flushed = flushed;
writer.written = written;
let copied = copy(&mut reader, &mut writer).await.unwrap();
assert_eq!(copied, 0);
}
});
run();
assert!(*flushed.borrow());
assert!(written.borrow().is_empty());
}
struct ParkingIo {
chunks: std::collections::VecDeque<Vec<u8>>,
park_before_next_read: bool,
written: Rc<RefCell<Vec<u8>>>,
closed: Rc<RefCell<bool>>,
}
impl ParkingIo {
fn new(chunks: &[&[u8]]) -> Self {
Self {
chunks: chunks.iter().map(|chunk| chunk.to_vec()).collect(),
park_before_next_read: true,
written: Rc::new(RefCell::new(Vec::new())),
closed: Rc::new(RefCell::new(false)),
}
}
fn written(&self) -> Rc<RefCell<Vec<u8>>> {
Rc::clone(&self.written)
}
}
impl AsyncRead for ParkingIo {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
if self.park_before_next_read {
self.park_before_next_read = false;
cx.waker().wake_by_ref();
return Poll::Pending;
}
self.park_before_next_read = true;
match self.chunks.pop_front() {
Some(chunk) => {
let read = chunk.len().min(buf.len());
buf[..read].copy_from_slice(&chunk[..read]);
Poll::Ready(Ok(read))
}
None => Poll::Ready(Ok(0)),
}
}
}
impl AsyncWrite for ParkingIo {
fn poll_write(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
self.written.borrow_mut().extend_from_slice(buf);
Poll::Ready(Ok(buf.len()))
}
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
*self.closed.borrow_mut() = true;
Poll::Ready(Ok(()))
}
}
#[test]
fn copy_does_not_re_emit_chunks_when_reads_park() {
let observed = Rc::new(RefCell::new(0u64));
let written = Rc::new(RefCell::new(Vec::new()));
spawn({
let observed = Rc::clone(&observed);
let written = Rc::clone(&written);
async move {
let mut reader = ParkingIo::new(&[b"hello"]);
let mut writer = ParkingIo::new(&[]);
writer.written = Rc::clone(&written);
*observed.borrow_mut() = copy(&mut reader, &mut writer).await.unwrap();
}
});
run();
assert_eq!(*observed.borrow(), 5);
assert_eq!(&*written.borrow(), b"hello");
}
#[test]
fn copy_bidirectional_does_not_re_emit_chunks_when_reads_park() {
let observed = Rc::new(RefCell::new((0u64, 0u64)));
let mut left = ParkingIo::new(&[b"to-right"]);
let mut right = ParkingIo::new(&[b"to-left"]);
let left_written = left.written();
let right_written = right.written();
spawn({
let observed = Rc::clone(&observed);
async move {
*observed.borrow_mut() = copy_bidirectional(&mut left, &mut right).await.unwrap();
}
});
run();
assert_eq!(*observed.borrow(), (8, 7));
assert_eq!(&*right_written.borrow(), b"to-right");
assert_eq!(&*left_written.borrow(), b"to-left");
}
}