nuclei/io/util/
copy.rs

1//! Copying from a reader to a writer: [`copy`], [`copy_buf`] and, both ways
2//! at once in one task, [`copy_bidirectional`].
3
4use core::future::poll_fn;
5use core::pin::Pin;
6use core::task::{Context, Poll, ready};
7use std::io;
8
9use futures_io::{AsyncBufRead, AsyncRead, AsyncWrite};
10
11/// The buffer [`copy`] and [`copy_bidirectional`] use for each direction.
12const DEFAULT_BUF_SIZE: usize = 8 * 1024;
13
14/// One direction of a copy: a buffer, what is in it, and the count copied.
15#[derive(Debug)]
16struct CopyBuffer {
17    read_done: bool,
18    need_flush: bool,
19    pos: usize,
20    cap: usize,
21    amount: u64,
22    buf: Box<[u8]>,
23}
24
25impl CopyBuffer {
26    fn new(size: usize) -> Self {
27        Self {
28            read_done: false,
29            need_flush: false,
30            pos: 0,
31            cap: 0,
32            amount: 0,
33            buf: vec![0; size.max(1)].into_boxed_slice(),
34        }
35    }
36
37    /// Copies until the reader ends and the writer is flushed; answers the
38    /// bytes copied. Before it waits on the reader it flushes what it wrote,
39    /// so nothing waits in the writer meanwhile.
40    fn poll_copy<R, W>(
41        &mut self,
42        cx: &mut Context<'_>,
43        mut reader: Pin<&mut R>,
44        mut writer: Pin<&mut W>,
45    ) -> Poll<io::Result<u64>>
46    where
47        R: AsyncRead + ?Sized,
48        W: AsyncWrite + ?Sized,
49    {
50        loop {
51            if self.pos == self.cap && !self.read_done {
52                match reader.as_mut().poll_read(cx, &mut self.buf) {
53                    Poll::Ready(Ok(0)) => self.read_done = true,
54                    Poll::Ready(Ok(n)) => {
55                        self.pos = 0;
56                        self.cap = n;
57                    }
58                    Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
59                    Poll::Pending => {
60                        if self.need_flush {
61                            ready!(writer.as_mut().poll_flush(cx))?;
62                            self.need_flush = false;
63                        }
64                        return Poll::Pending;
65                    }
66                }
67            }
68            while self.pos < self.cap {
69                let n = ready!(
70                    writer
71                        .as_mut()
72                        .poll_write(cx, &self.buf[self.pos..self.cap])
73                )?;
74                if n == 0 {
75                    return Poll::Ready(Err(io::Error::new(
76                        io::ErrorKind::WriteZero,
77                        "the writer took no more bytes",
78                    )));
79                }
80                self.pos += n;
81                self.amount += n as u64;
82                self.need_flush = true;
83            }
84            if self.pos == self.cap && self.read_done {
85                ready!(writer.as_mut().poll_flush(cx))?;
86                return Poll::Ready(Ok(self.amount));
87            }
88        }
89    }
90}
91
92/// Copies everything `reader` gives into `writer`, then flushes it; answers
93/// the bytes copied.
94///
95/// ```
96/// # nuclei::block_on(async {
97/// let mut reader: &[u8] = b"hello";
98/// let mut writer = Vec::new();
99/// let copied = nuclei::io::copy(&mut reader, &mut writer).await.unwrap();
100/// assert_eq!(copied, 5);
101/// assert_eq!(writer, b"hello");
102/// # });
103/// ```
104pub async fn copy<R, W>(reader: &mut R, writer: &mut W) -> io::Result<u64>
105where
106    R: AsyncRead + Unpin + ?Sized,
107    W: AsyncWrite + Unpin + ?Sized,
108{
109    let mut buf = CopyBuffer::new(DEFAULT_BUF_SIZE);
110    poll_fn(|cx| buf.poll_copy(cx, Pin::new(&mut *reader), Pin::new(&mut *writer))).await
111}
112
113/// Copies everything a buffered `reader` gives into `writer` from the
114/// reader's own buffer, then flushes it; answers the bytes copied.
115pub async fn copy_buf<R, W>(reader: &mut R, writer: &mut W) -> io::Result<u64>
116where
117    R: AsyncBufRead + Unpin + ?Sized,
118    W: AsyncWrite + Unpin + ?Sized,
119{
120    let mut amount = 0u64;
121    poll_fn(|cx| {
122        loop {
123            let mut reader = Pin::new(&mut *reader);
124            let writer = Pin::new(&mut *writer);
125            let available = ready!(reader.as_mut().poll_fill_buf(cx))?;
126            if available.is_empty() {
127                ready!(writer.poll_flush(cx))?;
128                return Poll::Ready(Ok(amount));
129            }
130            let n = ready!(writer.poll_write(cx, available))?;
131            if n == 0 {
132                return Poll::Ready(Err(io::Error::new(
133                    io::ErrorKind::WriteZero,
134                    "the writer took no more bytes",
135                )));
136            }
137            reader.consume(n);
138            amount += n as u64;
139        }
140    })
141    .await
142}
143
144/// One direction of [`copy_bidirectional`].
145#[derive(Debug)]
146enum Transfer {
147    Running(CopyBuffer),
148    ShuttingDown(u64),
149    Done(u64),
150}
151
152impl Transfer {
153    /// Copies, then closes the writer once the reader ended: the peer of the
154    /// writer sees the end the reader saw.
155    fn poll<R, W>(
156        &mut self,
157        cx: &mut Context<'_>,
158        reader: Pin<&mut R>,
159        mut writer: Pin<&mut W>,
160    ) -> Poll<io::Result<u64>>
161    where
162        R: AsyncRead + ?Sized,
163        W: AsyncWrite + ?Sized,
164    {
165        if let Self::Running(buf) = self {
166            let amount = ready!(buf.poll_copy(cx, reader, writer.as_mut()))?;
167            *self = Self::ShuttingDown(amount);
168        }
169        self.poll_after_copy(cx, writer)
170    }
171
172    fn poll_after_copy<W: AsyncWrite + ?Sized>(
173        &mut self,
174        cx: &mut Context<'_>,
175        writer: Pin<&mut W>,
176    ) -> Poll<io::Result<u64>> {
177        match *self {
178            Self::ShuttingDown(amount) => {
179                ready!(writer.poll_close(cx))?;
180                *self = Self::Done(amount);
181                Poll::Ready(Ok(amount))
182            }
183            Self::Done(amount) => Poll::Ready(Ok(amount)),
184            Self::Running(_) => Poll::Pending,
185        }
186    }
187}
188
189/// Copies `a` into `b` and `b` into `a` at once, in the calling task, each
190/// direction through a buffer of 8 KiB; a direction whose reader ends closes
191/// its writer. Answers the bytes copied each way, `(a to b, b to a)`, once
192/// both ended; the first error ends both.
193pub async fn copy_bidirectional<A, B>(a: &mut A, b: &mut B) -> io::Result<(u64, u64)>
194where
195    A: AsyncRead + AsyncWrite + Unpin + ?Sized,
196    B: AsyncRead + AsyncWrite + Unpin + ?Sized,
197{
198    copy_bidirectional_with_sizes(a, b, DEFAULT_BUF_SIZE, DEFAULT_BUF_SIZE).await
199}
200
201/// [`copy_bidirectional`] with a buffer of `a_to_b` bytes for one direction
202/// and of `b_to_a` bytes for the other.
203pub async fn copy_bidirectional_with_sizes<A, B>(
204    a: &mut A,
205    b: &mut B,
206    a_to_b: usize,
207    b_to_a: usize,
208) -> io::Result<(u64, u64)>
209where
210    A: AsyncRead + AsyncWrite + Unpin + ?Sized,
211    B: AsyncRead + AsyncWrite + Unpin + ?Sized,
212{
213    let mut forward = Transfer::Running(CopyBuffer::new(a_to_b));
214    let mut backward = Transfer::Running(CopyBuffer::new(b_to_a));
215    poll_fn(|cx| {
216        let ab = forward.poll(cx, Pin::new(&mut *a), Pin::new(&mut *b))?;
217        let ba = backward.poll(cx, Pin::new(&mut *b), Pin::new(&mut *a))?;
218        match (ab, ba) {
219            (Poll::Ready(ab), Poll::Ready(ba)) => Poll::Ready(Ok((ab, ba))),
220            _ => Poll::Pending,
221        }
222    })
223    .await
224}
nuclei/io/util/copy.rs

nuclei/io/util/
copy.rs
