nuclei/io/util/
read_ext.rs

1//! [`AsyncReadExt`]: the reads every [`AsyncRead`] gets, as futures.
2
3use core::future::Future;
4use core::pin::Pin;
5use core::task::{Context, Poll, ready};
6use std::io;
7
8use futures_io::AsyncRead;
9
10use super::ReadBuf;
11
12/// One typed-read method per line: name, type, width, conversion, what it
13/// reads.
14macro_rules! read_ints {
15    ($($name:ident, $ty:ty, $n:literal, $conv:ident, $what:literal;)+) => {$(
16        #[doc = concat!("Reads ", $what, " (", stringify!($n), " bytes, exactly).")]
17        fn $name(&mut self) -> ReadNumber<'_, Self, $ty, $n>
18        where
19            Self: Unpin,
20        {
21            ReadNumber {
22                reader: self,
23                bytes: [0; $n],
24                done: 0,
25                convert: <$ty>::$conv,
26            }
27        }
28    )+};
29}
30
31/// Reads on any [`AsyncRead`]: into a slice, a [`ReadBuf`], to the end, as
32/// typed numbers in big or little endian.
33pub trait AsyncReadExt: AsyncRead {
34    /// Reads some bytes into `buf`, answering how many (0 at the end).
35    fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> Read<'a, Self>
36    where
37        Self: Unpin,
38    {
39        Read { reader: self, buf }
40    }
41
42    /// Reads some bytes into the unfilled part of `buf` and counts them
43    /// filled, answering how many (0 at the end, or when `buf` is full).
44    fn read_buf<'a, 'b>(&'a mut self, buf: &'a mut ReadBuf<'b>) -> ReadBufFuture<'a, 'b, Self>
45    where
46        Self: Unpin,
47    {
48        ReadBufFuture { reader: self, buf }
49    }
50
51    /// Fills `buf` exactly, answering its length; an end of the stream
52    /// first is an [`UnexpectedEof`](io::ErrorKind::UnexpectedEof) error.
53    fn read_exact<'a>(&'a mut self, buf: &'a mut [u8]) -> ReadExact<'a, Self>
54    where
55        Self: Unpin,
56    {
57        ReadExact {
58            reader: self,
59            buf,
60            done: 0,
61        }
62    }
63
64    /// Reads to the end of the stream, appending to `buf`; answers how many
65    /// bytes it appended.
66    fn read_to_end<'a>(&'a mut self, buf: &'a mut Vec<u8>) -> ReadToEnd<'a, Self>
67    where
68        Self: Unpin,
69    {
70        let start = buf.len();
71        ReadToEnd {
72            reader: self,
73            buf,
74            start,
75            zeroed: 0,
76        }
77    }
78
79    /// Reads to the end of the stream, appending to `buf` if what it read is
80    /// UTF-8 (an [`InvalidData`](io::ErrorKind::InvalidData) error and `buf`
81    /// as it was otherwise); answers how many bytes it appended.
82    fn read_to_string<'a>(&'a mut self, buf: &'a mut String) -> ReadToString<'a, Self>
83    where
84        Self: Unpin,
85    {
86        ReadToString {
87            reader: self,
88            buf,
89            bytes: Vec::new(),
90            zeroed: 0,
91        }
92    }
93
94    /// A reader that ends after `limit` bytes of this one.
95    fn take(self, limit: u64) -> Take<Self>
96    where
97        Self: Sized,
98    {
99        Take { inner: self, limit }
100    }
101
102    /// A reader that reads this one to its end, then `next`.
103    fn chain<R: AsyncRead>(self, next: R) -> Chain<Self, R>
104    where
105        Self: Sized,
106    {
107        Chain {
108            first: self,
109            second: next,
110            first_done: false,
111        }
112    }
113
114    read_ints! {
115        read_u8, u8, 1, from_be_bytes, "an unsigned byte";
116        read_i8, i8, 1, from_be_bytes, "a signed byte";
117        read_u16, u16, 2, from_be_bytes, "a big-endian `u16`";
118        read_u16_le, u16, 2, from_le_bytes, "a little-endian `u16`";
119        read_i16, i16, 2, from_be_bytes, "a big-endian `i16`";
120        read_i16_le, i16, 2, from_le_bytes, "a little-endian `i16`";
121        read_u32, u32, 4, from_be_bytes, "a big-endian `u32`";
122        read_u32_le, u32, 4, from_le_bytes, "a little-endian `u32`";
123        read_i32, i32, 4, from_be_bytes, "a big-endian `i32`";
124        read_i32_le, i32, 4, from_le_bytes, "a little-endian `i32`";
125        read_u64, u64, 8, from_be_bytes, "a big-endian `u64`";
126        read_u64_le, u64, 8, from_le_bytes, "a little-endian `u64`";
127        read_i64, i64, 8, from_be_bytes, "a big-endian `i64`";
128        read_i64_le, i64, 8, from_le_bytes, "a little-endian `i64`";
129        read_u128, u128, 16, from_be_bytes, "a big-endian `u128`";
130        read_u128_le, u128, 16, from_le_bytes, "a little-endian `u128`";
131        read_i128, i128, 16, from_be_bytes, "a big-endian `i128`";
132        read_i128_le, i128, 16, from_le_bytes, "a little-endian `i128`";
133        read_f32, f32, 4, from_be_bytes, "a big-endian `f32`";
134        read_f32_le, f32, 4, from_le_bytes, "a little-endian `f32`";
135        read_f64, f64, 8, from_be_bytes, "a big-endian `f64`";
136        read_f64_le, f64, 8, from_le_bytes, "a little-endian `f64`";
137    }
138}
139
140impl<R: AsyncRead + ?Sized> AsyncReadExt for R {}
141
142/// The future of [`read`](AsyncReadExt::read).
143#[derive(Debug)]
144#[must_use = "futures do nothing unless polled"]
145pub struct Read<'a, R: ?Sized> {
146    reader: &'a mut R,
147    buf: &'a mut [u8],
148}
149
150impl<R: AsyncRead + Unpin + ?Sized> Future for Read<'_, R> {
151    type Output = io::Result<usize>;
152
153    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> {
154        let this = self.get_mut();
155        Pin::new(&mut *this.reader).poll_read(cx, this.buf)
156    }
157}
158
159/// The future of [`read_buf`](AsyncReadExt::read_buf).
160#[derive(Debug)]
161#[must_use = "futures do nothing unless polled"]
162pub struct ReadBufFuture<'a, 'b, R: ?Sized> {
163    reader: &'a mut R,
164    buf: &'a mut ReadBuf<'b>,
165}
166
167impl<R: AsyncRead + Unpin + ?Sized> Future for ReadBufFuture<'_, '_, R> {
168    type Output = io::Result<usize>;
169
170    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> {
171        let this = self.get_mut();
172        poll_read_buf(Pin::new(&mut *this.reader), cx, this.buf)
173    }
174}
175
176/// Reads into the unfilled part of `buf` and counts what came as filled.
177pub(crate) fn poll_read_buf<R: AsyncRead + ?Sized>(
178    reader: Pin<&mut R>,
179    cx: &mut Context<'_>,
180    buf: &mut ReadBuf<'_>,
181) -> Poll<io::Result<usize>> {
182    if buf.remaining() == 0 {
183        return Poll::Ready(Ok(0));
184    }
185    let n = ready!(reader.poll_read(cx, buf.initialize_unfilled()))?;
186    buf.advance(n);
187    Poll::Ready(Ok(n))
188}
189
190/// The future of [`read_exact`](AsyncReadExt::read_exact).
191#[derive(Debug)]
192#[must_use = "futures do nothing unless polled"]
193pub struct ReadExact<'a, R: ?Sized> {
194    reader: &'a mut R,
195    buf: &'a mut [u8],
196    done: usize,
197}
198
199impl<R: AsyncRead + Unpin + ?Sized> Future for ReadExact<'_, R> {
200    type Output = io::Result<usize>;
201
202    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> {
203        let this = self.get_mut();
204        ready!(fill(
205            Pin::new(&mut *this.reader),
206            cx,
207            this.buf,
208            &mut this.done
209        ))?;
210        Poll::Ready(Ok(this.buf.len()))
211    }
212}
213
214/// Reads until `buf[*done..]` is full; an end first is an error.
215fn fill<R: AsyncRead + ?Sized>(
216    mut reader: Pin<&mut R>,
217    cx: &mut Context<'_>,
218    buf: &mut [u8],
219    done: &mut usize,
220) -> Poll<io::Result<()>> {
221    while *done < buf.len() {
222        let n = ready!(reader.as_mut().poll_read(cx, &mut buf[*done..]))?;
223        if n == 0 {
224            return Poll::Ready(Err(io::ErrorKind::UnexpectedEof.into()));
225        }
226        *done += n;
227    }
228    Poll::Ready(Ok(()))
229}
230
231/// The future of a typed read ([`read_u32`](AsyncReadExt::read_u32) and the
232/// like).
233#[must_use = "futures do nothing unless polled"]
234pub struct ReadNumber<'a, R: ?Sized, T, const N: usize> {
235    reader: &'a mut R,
236    bytes: [u8; N],
237    done: usize,
238    convert: fn([u8; N]) -> T,
239}
240
241impl<R: AsyncRead + Unpin + ?Sized, T, const N: usize> Future for ReadNumber<'_, R, T, N> {
242    type Output = io::Result<T>;
243
244    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<T>> {
245        let this = self.get_mut();
246        ready!(fill(
247            Pin::new(&mut *this.reader),
248            cx,
249            &mut this.bytes,
250            &mut this.done
251        ))?;
252        Poll::Ready(Ok((this.convert)(this.bytes)))
253    }
254}
255
256impl<R: ?Sized, T, const N: usize> core::fmt::Debug for ReadNumber<'_, R, T, N> {
257    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
258        f.debug_struct("ReadNumber")
259            .field("width", &N)
260            .field("done", &self.done)
261            .finish_non_exhaustive()
262    }
263}
264
265/// The future of [`read_to_end`](AsyncReadExt::read_to_end).
266#[derive(Debug)]
267#[must_use = "futures do nothing unless polled"]
268pub struct ReadToEnd<'a, R: ?Sized> {
269    reader: &'a mut R,
270    buf: &'a mut Vec<u8>,
271    start: usize,
272    /// Bytes of `buf`'s spare capacity already zeroed.
273    zeroed: usize,
274}
275
276impl<R: AsyncRead + Unpin + ?Sized> Future for ReadToEnd<'_, R> {
277    type Output = io::Result<usize>;
278
279    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> {
280        let this = self.get_mut();
281        ready!(read_to_end(
282            Pin::new(&mut *this.reader),
283            cx,
284            this.buf,
285            &mut this.zeroed
286        ))?;
287        Poll::Ready(Ok(this.buf.len() - this.start))
288    }
289}
290
291/// The least room a read to the end asks of the vector, and the most one
292/// read takes: the vector grows by doubling in between.
293const MIN_READ: usize = 32;
294const MAX_READ: usize = 64 * 1024;
295
296/// Reads to the end, appending to `buf`. Reads land in `buf`'s spare
297/// capacity, at most [`MAX_READ`] bytes a read; when less than [`MIN_READ`]
298/// is spare, the vector grows by doubling (by [`MAX_READ`] at most). Spare
299/// memory is zeroed once, when it is new (`zeroed` counts what already is,
300/// across polls), not before each read.
301fn read_to_end<R: AsyncRead + ?Sized>(
302    mut reader: Pin<&mut R>,
303    cx: &mut Context<'_>,
304    buf: &mut Vec<u8>,
305    zeroed: &mut usize,
306) -> Poll<io::Result<()>> {
307    loop {
308        if buf.capacity() - buf.len() < MIN_READ {
309            let before = buf.capacity();
310            buf.reserve(buf.len().clamp(MIN_READ, MAX_READ));
311            if buf.capacity() != before {
312                // A new allocation keeps only the bytes in use.
313                *zeroed = 0;
314            }
315        }
316        let spare = buf.spare_capacity_mut();
317        let room = spare.len().min(MAX_READ);
318        if *zeroed < room {
319            for byte in &mut spare[*zeroed..room] {
320                byte.write(0);
321            }
322            *zeroed = room;
323        }
324        // SAFETY: the first `room` spare bytes are initialized (zeroed now
325        // or before, and only ever written with bytes since).
326        let dst = unsafe { core::slice::from_raw_parts_mut(spare.as_mut_ptr().cast::<u8>(), room) };
327        let n = ready!(reader.as_mut().poll_read(cx, dst))?;
328        if n == 0 {
329            return Poll::Ready(Ok(()));
330        }
331        if n > room {
332            return Poll::Ready(Err(io::Error::new(
333                io::ErrorKind::InvalidData,
334                "the reader reported more bytes than it was given room for",
335            )));
336        }
337        // SAFETY: the first `n <= room` spare bytes are initialized and hold
338        // what the reader read.
339        unsafe { buf.set_len(buf.len() + n) };
340        *zeroed -= n;
341    }
342}
343
344/// The future of [`read_to_string`](AsyncReadExt::read_to_string).
345#[derive(Debug)]
346#[must_use = "futures do nothing unless polled"]
347pub struct ReadToString<'a, R: ?Sized> {
348    reader: &'a mut R,
349    buf: &'a mut String,
350    bytes: Vec<u8>,
351    zeroed: usize,
352}
353
354impl<R: AsyncRead + Unpin + ?Sized> Future for ReadToString<'_, R> {
355    type Output = io::Result<usize>;
356
357    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> {
358        let this = self.get_mut();
359        ready!(read_to_end(
360            Pin::new(&mut *this.reader),
361            cx,
362            &mut this.bytes,
363            &mut this.zeroed
364        ))?;
365        let bytes = core::mem::take(&mut this.bytes);
366        match String::from_utf8(bytes) {
367            Ok(text) => {
368                this.buf.push_str(&text);
369                Poll::Ready(Ok(text.len()))
370            }
371            Err(_) => Poll::Ready(Err(io::Error::new(
372                io::ErrorKind::InvalidData,
373                "the stream did not hold UTF-8",
374            ))),
375        }
376    }
377}
378
379/// A reader that ends after a number of bytes of another; made by
380/// [`take`](AsyncReadExt::take).
381#[derive(Debug)]
382pub struct Take<R> {
383    inner: R,
384    limit: u64,
385}
386
387impl<R> Take<R> {
388    /// The bytes left before it ends.
389    pub fn limit(&self) -> u64 {
390        self.limit
391    }
392
393    /// Sets the bytes left before it ends.
394    pub fn set_limit(&mut self, limit: u64) {
395        self.limit = limit;
396    }
397
398    /// The reader it reads from.
399    pub fn get_ref(&self) -> &R {
400        &self.inner
401    }
402
403    /// The reader it reads from, to change.
404    pub fn get_mut(&mut self) -> &mut R {
405        &mut self.inner
406    }
407
408    /// The reader it reads from, taken back.
409    pub fn into_inner(self) -> R {
410        self.inner
411    }
412}
413
414impl<R: AsyncRead + Unpin> AsyncRead for Take<R> {
415    fn poll_read(
416        self: Pin<&mut Self>,
417        cx: &mut Context<'_>,
418        buf: &mut [u8],
419    ) -> Poll<io::Result<usize>> {
420        let this = self.get_mut();
421        if this.limit == 0 {
422            return Poll::Ready(Ok(0));
423        }
424        let max = usize::try_from(this.limit).map_or(buf.len(), |l| l.min(buf.len()));
425        let n = ready!(Pin::new(&mut this.inner).poll_read(cx, &mut buf[..max]))?;
426        this.limit -= n as u64;
427        Poll::Ready(Ok(n))
428    }
429}
430
431/// A reader that reads one to its end, then another; made by
432/// [`chain`](AsyncReadExt::chain).
433#[derive(Debug)]
434pub struct Chain<A, B> {
435    first: A,
436    second: B,
437    first_done: bool,
438}
439
440impl<A, B> Chain<A, B> {
441    /// The two readers.
442    pub fn get_ref(&self) -> (&A, &B) {
443        (&self.first, &self.second)
444    }
445
446    /// The two readers, to change.
447    pub fn get_mut(&mut self) -> (&mut A, &mut B) {
448        (&mut self.first, &mut self.second)
449    }
450
451    /// The two readers, taken back.
452    pub fn into_inner(self) -> (A, B) {
453        (self.first, self.second)
454    }
455}
456
457impl<A: AsyncRead + Unpin, B: AsyncRead + Unpin> AsyncRead for Chain<A, B> {
458    fn poll_read(
459        self: Pin<&mut Self>,
460        cx: &mut Context<'_>,
461        buf: &mut [u8],
462    ) -> Poll<io::Result<usize>> {
463        let this = self.get_mut();
464        if !this.first_done {
465            let n = ready!(Pin::new(&mut this.first).poll_read(cx, buf))?;
466            if n > 0 || buf.is_empty() {
467                return Poll::Ready(Ok(n));
468            }
469            this.first_done = true;
470        }
471        Pin::new(&mut this.second).poll_read(cx, buf)
472    }
473}
nuclei/io/util/read_ext.rs

nuclei/io/util/
read_ext.rs
