src/util/io.rs

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();
    }
}