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