#
cargo/util/
io.rs

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51

    use std::io::{self, Read, Take};

    #[derive(Debug)]
    pub struct LimitErrorReader<R> {
        inner: Take<R>,
    }

    impl<R: Read> LimitErrorReader<R> {
        pub fn new(r: R, limit: u64) -> LimitErrorReader<R> {
            LimitErrorReader {
                inner: r.take(limit),
            }
        }
    }

    impl<R: Read> Read for LimitErrorReader<R> {
        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
            match self.inner.read(buf) {
                Ok(0) if self.inner.limit() == 0 => Err(io::Error::new(
                    io::ErrorKind::Other,
                    "maximum limit reached when reading",
                )),
                e => e,
            }
        }
    }

    #[cfg(test)]
    mod tests {
        use super::LimitErrorReader;

        use std::io::Read;

        #[test]
        fn under_the_limit() {
            let buf = &[1; 7][..];
            let mut r = LimitErrorReader::new(buf, 8);
            let mut out = Vec::new();
            assert!(matches!(r.read_to_end(&mut out), Ok(7)));
            assert_eq!(buf, out.as_slice());
        }

        #[test]
        #[should_panic = "maximum limit reached when reading"]
        fn over_the_limit() {
            let buf = &[1; 8][..];
            let mut r = LimitErrorReader::new(buf, 8);
            let mut out = Vec::new();
            r.read_to_end(&mut out).unwrap();
        }
    }