#
cargo/util/
io.rs

    1use std::io::{self, Read, Take};
    2
    3#[derive(Debug)]
    4pub struct LimitErrorReader<R> {
    5    inner: Take<R>,
    6}
    7
    8impl<R: Read> LimitErrorReader<R> {
    9    pub fn new(r: R, limit: u64) -> LimitErrorReader<R> {
    10        LimitErrorReader {
    11            inner: r.take(limit),
    12        }
    13    }
    14}
    15
    16impl<R: Read> Read for LimitErrorReader<R> {
    17    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
    18        match self.inner.read(buf) {
    19            Ok(0) if self.inner.limit() == 0 => Err(io::Error::new(
    20                io::ErrorKind::Other,
    21                "maximum limit reached when reading",
    22            )),
    23            e => e,
    24        }
    25    }
    26}
    27
    28#[cfg(test)]
    29mod tests {
    30    use super::LimitErrorReader;
    31
    32    use std::io::Read;
    33
    34    #[test]
    35    fn under_the_limit() {
    36        let buf = &[1; 7][..];
    37        let mut r = LimitErrorReader::new(buf, 8);
    38        let mut out = Vec::new();
    39        assert!(matches!(r.read_to_end(&mut out), Ok(7)));
    40        assert_eq!(buf, out.as_slice());
    41    }
    42
    43    #[test]
    44    #[should_panic = "maximum limit reached when reading"]
    45    fn over_the_limit() {
    46        let buf = &[1; 8][..];
    47        let mut r = LimitErrorReader::new(buf, 8);
    48        let mut out = Vec::new();
    49        r.read_to_end(&mut out).unwrap();
    50    }
    51}