limit_reader/
lib.rs

1use std::io;
2use std::io::Read;
3
4struct LimitReader<T: Read> {
5    reader: T,
6    limit: usize,
7    reader_count: usize,
8}
9
10impl<T: Read> LimitReader<T> {
11    fn new(reader: T, limit: usize) -> Self {
12        LimitReader {
13            reader,
14            limit,
15            reader_count: 0,
16        }
17    }
18}
19
20impl<T: Read> Read for LimitReader<T> {
21    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
22        let max_read = self.limit.min(buf.len()); // min of limit and buf.len()
23        let bytes_read = self.reader.read(&mut buf[..max_read])?;
24        self.reader_count += 1;
25        Ok(bytes_read)
26    }
27}
28
29#[cfg(test)]
30mod tests {
31    use std::fs;
32    use std::io::Read;
33
34    use crate::LimitReader;
35
36    #[test]
37    fn test_read() {
38        let hello_txt = fs::File::open("testdata/hello.txt").unwrap();
39        let mut reader = LimitReader::new(hello_txt, 5);
40        let mut r1 = vec![0u8];
41        reader.read_to_end(&mut r1).unwrap();
42
43        let mut hello_txt = fs::File::open("testdata/hello.txt").unwrap();
44        let mut r2 = vec![0u8];
45        hello_txt.read_to_end(&mut r2).unwrap();
46        assert_eq!(r1, r2);
47    }
48}