compio_io/util/
take.rs

1use compio_buf::{BufResult, IntoInner, IoBufMut, buf_try};
2
3use crate::{AsyncBufRead, AsyncRead, IoResult};
4
5/// Read up to a limit number of bytes from reader.
6#[derive(Debug)]
7pub struct Take<R> {
8    reader: R,
9    limit: u64,
10}
11
12impl<T> Take<T> {
13    pub(crate) fn new(reader: T, limit: u64) -> Self {
14        Self { reader, limit }
15    }
16
17    /// Returns the number of bytes that can be read before this instance will
18    /// return EOF.
19    ///
20    /// # Note
21    ///
22    /// This instance may reach `EOF` after reading fewer bytes than indicated
23    /// by this method if the underlying [`AsyncRead`] instance reaches EOF.
24    pub fn limit(&self) -> u64 {
25        self.limit
26    }
27
28    /// Sets the number of bytes that can be read before this instance will
29    /// return EOF. This is the same as constructing a new `Take` instance, so
30    /// the amount of bytes read and the previous limit value don't matter when
31    /// calling this method.
32    pub fn set_limit(&mut self, limit: u64) {
33        self.limit = limit;
34    }
35
36    /// Consumes the `Take`, returning the wrapped reader.
37    pub fn into_inner(self) -> T {
38        self.reader
39    }
40
41    /// Gets a reference to the underlying reader.
42    pub fn get_ref(&self) -> &T {
43        &self.reader
44    }
45
46    /// Gets a mutable reference to the underlying reader.
47    ///
48    /// Care should be taken to avoid modifying the internal I/O state of the
49    /// underlying reader as doing so may corrupt the internal limit of this
50    /// `Take`.
51    pub fn get_mut(&mut self) -> &mut T {
52        &mut self.reader
53    }
54}
55
56impl<R: AsyncRead> AsyncRead for Take<R> {
57    async fn read<B: IoBufMut>(&mut self, buf: B) -> BufResult<usize, B> {
58        if self.limit == 0 {
59            return BufResult(Ok(0), buf);
60        }
61
62        let max = self.limit.min(buf.buf_capacity() as u64) as usize;
63        let buf = buf.slice(..max);
64
65        let (n, buf) = buf_try!(self.reader.read(buf).await.into_inner());
66        assert!(n as u64 <= self.limit, "number of read bytes exceeds limit");
67        self.limit -= n as u64;
68
69        BufResult(Ok(n), buf)
70    }
71}
72
73impl<R: AsyncBufRead> AsyncBufRead for Take<R> {
74    async fn fill_buf(&mut self) -> IoResult<&'_ [u8]> {
75        if self.limit == 0 {
76            return Ok(&[]);
77        }
78
79        let buf = self.reader.fill_buf().await?;
80        let cap = self.limit.min(buf.len() as u64) as usize;
81        Ok(&buf[..cap])
82    }
83
84    fn consume(&mut self, amount: usize) {
85        // Don't let callers reset the limit by passing an overlarge value
86        let amount = self.limit.min(amount as u64) as usize;
87        self.limit -= amount as u64;
88        self.reader.consume(amount);
89    }
90}
compio_io/util/take.rs

compio_io/util/
take.rs
