atm_io_utils/
limited_reader.rs

1//! A wrapper around a reader that limits how many bytes can be read from it.
2
3use std::cmp::min;
4
5use futures_core::Poll;
6use futures_core::task::Context;
7use futures_io::{AsyncRead, Error as FutIoErr};
8
9/// Wraps a reader and limits the number of bytes that can be read from it. Once the limit has been
10/// reached, further calls to poll_read will return `Ok(Ready(0))`.
11pub struct LimitedReader<R> {
12    inner: R,
13    remaining: usize,
14}
15
16impl<R> LimitedReader<R> {
17    /// Create a new `LimitedReader`, wrapping the given reader.
18    pub fn new(inner: R, limit: usize) -> LimitedReader<R> {
19        LimitedReader {
20            inner: inner,
21            remaining: limit,
22        }
23    }
24}
25
26impl<R: AsyncRead> AsyncRead for LimitedReader<R> {
27    fn poll_read(&mut self, cx: &mut Context, buf: &mut [u8]) -> Poll<usize, FutIoErr> {
28        let upper = min(self.remaining, buf.len());
29        self.inner.poll_read(cx, &mut buf[..upper])
30    }
31}
atm_io_utils/limited_reader.rs

atm_io_utils/
limited_reader.rs
