1 use std::io;
2 use std::cmp;
3 
4 use super::*;
5 
6 /// Limits the amount of data that can be read from a
7 /// `BufferedReader`.
8 pub struct Limitor<T: BufferedReader<C>, C> {
9     reader: T,
10     limit: u64,
11 
12     cookie: C,
13 }
14 
15 impl<T: BufferedReader<C>, C> fmt::Display for Limitor<T, C> {
fmt(&self, f: &mut fmt::Formatter) -> fmt::Result16     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
17         write!(f, "Limitor ({} bytes)", self.limit)
18     }
19 }
20 
21 impl<T: BufferedReader<C>, C> fmt::Debug for Limitor<T, C> {
fmt(&self, f: &mut fmt::Formatter) -> fmt::Result22     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
23         f.debug_struct("Limitor")
24             .field("limit", &self.limit)
25             .field("reader", &self.reader)
26             .finish()
27     }
28 }
29 
30 impl<T: BufferedReader<()>> Limitor<T, ()> {
31     /// Instantiates a new limitor.
32     ///
33     /// `reader` is the source to wrap.  `limit` is the maximum number
34     /// of bytes that can be read from the source.
new(reader: T, limit: u64) -> Self35     pub fn new(reader: T, limit: u64) -> Self {
36         Self::with_cookie(reader, limit, ())
37     }
38 }
39 
40 impl<T: BufferedReader<C>, C> Limitor<T, C> {
41     /// Like `new()`, but sets a cookie.
42     ///
43     /// The cookie can be retrieved using the `cookie_ref` and
44     /// `cookie_mut` methods, and set using the `cookie_set` method.
with_cookie(reader: T, limit: u64, cookie: C) -> Limitor<T, C>45     pub fn with_cookie(reader: T, limit: u64, cookie: C)
46             -> Limitor<T, C> {
47         Limitor {
48             reader,
49             limit,
50             cookie,
51         }
52     }
53 }
54 
55 impl<T: BufferedReader<C>, C> io::Read for Limitor<T, C> {
read(&mut self, buf: &mut [u8]) -> Result<usize, io::Error>56     fn read(&mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
57         let len = cmp::min(self.limit, buf.len() as u64) as usize;
58         let result = self.reader.read(&mut buf[0..len]);
59         if let Ok(amount) = result {
60             self.limit -= amount as u64;
61         }
62         result
63     }
64 }
65 
66 impl<T: BufferedReader<C>, C> BufferedReader<C> for Limitor<T, C> {
buffer(&self) -> &[u8]67     fn buffer(&self) -> &[u8] {
68         let buf = self.reader.buffer();
69         &buf[..cmp::min(buf.len(),
70                         cmp::min(std::usize::MAX as u64,
71                                  self.limit) as usize)]
72     }
73 
74     /// Return the buffer.  Ensure that it contains at least `amount`
75     /// bytes.
data(&mut self, amount: usize) -> Result<&[u8], io::Error>76     fn data(&mut self, amount: usize) -> Result<&[u8], io::Error> {
77         let amount = cmp::min(amount as u64, self.limit) as usize;
78         let result = self.reader.data(amount);
79         match result {
80             Ok(ref buffer) =>
81                 if buffer.len() as u64 > self.limit {
82                     return Ok(&buffer[0..self.limit as usize]);
83                 } else {
84                     return Ok(buffer);
85                 },
86             Err(err) => return Err(err),
87         }
88     }
89 
consume(&mut self, amount: usize) -> &[u8]90     fn consume(&mut self, amount: usize) -> &[u8] {
91         assert!(amount as u64 <= self.limit);
92         self.limit -= amount as u64;
93         let data = self.reader.consume(amount);
94         return &data[..cmp::min(self.limit + amount as u64, data.len() as u64) as usize];
95     }
96 
data_consume(&mut self, amount: usize) -> Result<&[u8], io::Error>97     fn data_consume(&mut self, amount: usize) -> Result<&[u8], io::Error> {
98         let amount = cmp::min(amount as u64, self.limit) as usize;
99         let result = self.reader.data_consume(amount);
100         if let Ok(ref buffer) = result {
101             let amount = cmp::min(amount, buffer.len());
102             self.limit -= amount as u64;
103             return Ok(&buffer[
104                 ..cmp::min(buffer.len() as u64, self.limit + amount as u64) as usize]);
105         }
106         return result;
107     }
108 
data_consume_hard(&mut self, amount: usize) -> Result<&[u8], io::Error>109     fn data_consume_hard(&mut self, amount: usize) -> Result<&[u8], io::Error> {
110         if amount as u64 > self.limit {
111             return Err(Error::new(ErrorKind::UnexpectedEof, "EOF"));
112         }
113         let result = self.reader.data_consume_hard(amount);
114         if let Ok(ref buffer) = result {
115             let amount = cmp::min(amount, buffer.len());
116             self.limit -= amount as u64;
117             return Ok(&buffer[
118                 ..cmp::min(buffer.len() as u64, self.limit + amount as u64) as usize]);
119         }
120         return result;
121     }
122 
consummated(&mut self) -> bool123     fn consummated(&mut self) -> bool {
124         self.limit == 0
125     }
126 
get_mut(&mut self) -> Option<&mut dyn BufferedReader<C>>127     fn get_mut(&mut self) -> Option<&mut dyn BufferedReader<C>> {
128         Some(&mut self.reader)
129     }
130 
get_ref(&self) -> Option<&dyn BufferedReader<C>>131     fn get_ref(&self) -> Option<&dyn BufferedReader<C>> {
132         Some(&self.reader)
133     }
134 
into_inner<'b>(self: Box<Self>) -> Option<Box<dyn BufferedReader<C> + 'b>> where Self: 'b135     fn into_inner<'b>(self: Box<Self>) -> Option<Box<dyn BufferedReader<C> + 'b>>
136         where Self: 'b {
137         Some(self.reader.as_boxed())
138     }
139 
cookie_set(&mut self, cookie: C) -> C140     fn cookie_set(&mut self, cookie: C) -> C {
141         use std::mem;
142 
143         mem::replace(&mut self.cookie, cookie)
144     }
145 
cookie_ref(&self) -> &C146     fn cookie_ref(&self) -> &C {
147         &self.cookie
148     }
149 
cookie_mut(&mut self) -> &mut C150     fn cookie_mut(&mut self) -> &mut C {
151         &mut self.cookie
152     }
153 }
154 
155 #[cfg(test)]
156 mod test {
157     use super::*;
158 
159     #[test]
buffered_reader_limitor_test()160     fn buffered_reader_limitor_test() {
161         let data : &[u8] = b"01234567890123456789";
162 
163         /* Add a single limitor.  */
164         {
165             let mut bio : Box<dyn BufferedReader<()>>
166                 = Box::new(Memory::new(data));
167 
168             bio = {
169                 let mut bio2 = Box::new(Limitor::new(bio, 5));
170                 {
171                     let result = bio2.data(5).unwrap();
172                     assert_eq!(result.len(), 5);
173                     assert_eq!(result, &b"01234"[..]);
174                 }
175                 bio2.consume(5);
176                 {
177                     let result = bio2.data(1).unwrap();
178                     assert_eq!(result.len(), 0);
179                     assert_eq!(result, &b""[..]);
180                 }
181 
182                 bio2.into_inner().unwrap()
183             };
184 
185             {
186                 {
187                     let result = bio.data(15).unwrap();
188                     assert_eq!(result.len(), 15);
189                     assert_eq!(result, &b"567890123456789"[..]);
190                 }
191                 bio.consume(15);
192                 {
193                     let result = bio.data(1).unwrap();
194                     assert_eq!(result.len(), 0);
195                     assert_eq!(result, &b""[..]);
196                 }
197             }
198         }
199 
200         /* Try with two limitors where the first one imposes the real
201          * limit.  */
202         {
203             let mut bio : Box<dyn BufferedReader<()>>
204                 = Box::new(Memory::new(data));
205 
206             bio = {
207                 let bio2 : Box<dyn BufferedReader<()>>
208                     = Box::new(Limitor::new(bio, 5));
209                 // We limit to 15 bytes, but bio2 will still limit us to 5
210                 // bytes.
211                 let mut bio3 : Box<dyn BufferedReader<()>>
212                     = Box::new(Limitor::new(bio2, 15));
213                 {
214                     let result = bio3.data(100).unwrap();
215                     assert_eq!(result.len(), 5);
216                     assert_eq!(result, &b"01234"[..]);
217                 }
218                 bio3.consume(5);
219                 {
220                     let result = bio3.data(1).unwrap();
221                     assert_eq!(result.len(), 0);
222                     assert_eq!(result, &b""[..]);
223                 }
224 
225                 bio3.into_inner().unwrap().into_inner().unwrap()
226             };
227 
228             {
229                 {
230                     let result = bio.data(15).unwrap();
231                     assert_eq!(result.len(), 15);
232                     assert_eq!(result, &b"567890123456789"[..]);
233                 }
234                 bio.consume(15);
235                 {
236                     let result = bio.data(1).unwrap();
237                     assert_eq!(result.len(), 0);
238                     assert_eq!(result, &b""[..]);
239                 }
240             }
241         }
242     }
243 
244     // Test that buffer() returns the same data as data().
245     #[test]
buffer_test()246     fn buffer_test() {
247         // Test vector.
248         let size = 10 * DEFAULT_BUF_SIZE;
249         let mut input = Vec::with_capacity(size);
250         let mut v = 0u8;
251         for _ in 0..size {
252             input.push(v);
253             if v == std::u8::MAX {
254                 v = 0;
255             } else {
256                 v += 1;
257             }
258         }
259 
260         let reader = Generic::new(&input[..], None);
261         let size = size / 2;
262         let input = &input[..size];
263         let mut reader = Limitor::new(reader, input.len() as u64);
264 
265         // Gather some stats to make it easier to figure out whether
266         // this test is working.
267         let stats_count =  2 * DEFAULT_BUF_SIZE;
268         let mut stats = vec![0usize; stats_count];
269 
270         for i in 0..input.len() {
271             let data = reader.data(DEFAULT_BUF_SIZE + 1).unwrap().to_vec();
272             assert!(data.len() > 0);
273             assert_eq!(data, reader.buffer());
274             // And, we may as well check to make sure we read the
275             // right data.
276             assert_eq!(data, &input[i..i+data.len()]);
277 
278             stats[cmp::min(data.len(), stats_count - 1)] += 1;
279 
280             // Consume one byte and see what happens.
281             reader.consume(1);
282         }
283 
284         if false {
285             for i in 0..stats.len() {
286                 if stats[i] > 0 {
287                     if i == stats.len() - 1 {
288                         eprint!(">=");
289                     }
290                     eprintln!("{}: {}", i, stats[i]);
291                 }
292             }
293         }
294     }
295 
296     #[test]
consummated()297     fn consummated() {
298         let data = b"0123456789";
299 
300         let mut l = Limitor::new(Memory::new(data), 10);
301         l.drop_eof().unwrap();
302         assert!(l.consummated());
303 
304         let mut l = Limitor::new(Memory::new(data), 20);
305         l.drop_eof().unwrap();
306         eprintln!("{:?}", l);
307         assert!(! l.consummated());
308     }
309 }
310