Skip to main content

brotli_decompressor/
reader.rs

1#[cfg(feature="std")]
2use std::io::{self, Error, ErrorKind, Read};
3#[cfg(feature="std")]
4pub use alloc_stdlib::StandardAlloc;
5#[cfg(all(feature="unsafe",feature="std"))]
6pub use alloc_stdlib::HeapAlloc;
7pub use huffman::{HuffmanCode, HuffmanTreeGroup};
8pub use state::BrotliState;
9// use io_wrappers::write_all;
10pub use io_wrappers::{CustomRead, CustomWrite};
11#[cfg(feature="std")]
12pub use io_wrappers::{IntoIoReader, IoReaderWrapper, IoWriterWrapper};
13pub use super::decode::{BrotliDecompressStream, BrotliResult};
14pub use alloc::{AllocatedStackMemory, Allocator, SliceWrapper, SliceWrapperMut, StackAllocator};
15
16#[cfg(feature="std")]
17pub struct DecompressorCustomAlloc<R: Read,
18     BufferType : SliceWrapperMut<u8>,
19     AllocU8 : Allocator<u8>,
20     AllocU32 : Allocator<u32>,
21     AllocHC : Allocator<HuffmanCode> >(DecompressorCustomIo<io::Error,
22                                                             IntoIoReader<R>,
23                                                             BufferType,
24                                                             AllocU8, AllocU32, AllocHC>);
25
26
27#[cfg(feature="std")]
28impl<R: Read,
29     BufferType : SliceWrapperMut<u8>,
30     AllocU8,
31     AllocU32,
32     AllocHC> DecompressorCustomAlloc<R, BufferType, AllocU8, AllocU32, AllocHC>
33 where AllocU8 : Allocator<u8>, AllocU32 : Allocator<u32>, AllocHC : Allocator<HuffmanCode>
34    {
35
36    pub fn new(r: R, buffer : BufferType,
37               alloc_u8 : AllocU8, alloc_u32 : AllocU32, alloc_hc : AllocHC) -> Self {
38        DecompressorCustomAlloc::<R, BufferType, AllocU8, AllocU32, AllocHC>(
39          DecompressorCustomIo::<Error,
40                                 IntoIoReader<R>,
41                                 BufferType,
42                                 AllocU8, AllocU32, AllocHC>::new(IntoIoReader::<R>(r),
43                                                                  buffer,
44                                                                  alloc_u8, alloc_u32, alloc_hc,
45                                                                  Error::new(ErrorKind::InvalidData,
46                                                                             "Invalid Data")))
47    }
48
49    pub fn new_with_custom_dictionary(r: R, buffer : BufferType,
50               alloc_u8 : AllocU8, alloc_u32 : AllocU32, alloc_hc : AllocHC,
51               dict: AllocU8::AllocatedMemory) -> Self {
52        DecompressorCustomAlloc::<R, BufferType, AllocU8, AllocU32, AllocHC>(
53          DecompressorCustomIo::<Error,
54                                 IntoIoReader<R>,
55                                 BufferType,
56                                 AllocU8, AllocU32, AllocHC>::new_with_custom_dictionary(IntoIoReader::<R>(r),
57                                                                                   buffer,
58                                                                                   alloc_u8, alloc_u32, alloc_hc,
59                                                                                   dict,
60                                                                                   Error::new(ErrorKind::InvalidData,
61                                                                                              "Invalid Data")))
62    }
63
64    pub fn attach_dictionary(&mut self, dict: AllocU8::AllocatedMemory) -> bool {
65      self.0.attach_dictionary(dict)
66    }
67    pub fn attach_serialized_dictionary(&mut self, dict: AllocU8::AllocatedMemory) -> bool {
68      self.0.attach_serialized_dictionary(dict)
69    }
70
71    pub fn get_ref(&self) -> &R {
72      &self.0.get_ref().0
73    }
74    pub fn get_mut(&mut self) -> &mut R {
75      &mut self.0.get_mut().0
76    }
77    pub fn into_inner(self) -> R {
78      self.0.into_inner().0
79    }
80}
81#[cfg(feature="std")]
82impl<R: Read,
83     BufferType : SliceWrapperMut<u8>,
84     AllocU8 : Allocator<u8>,
85     AllocU32 : Allocator<u32>,
86     AllocHC : Allocator<HuffmanCode> > Read for DecompressorCustomAlloc<R,
87                                                                         BufferType,
88                                                                         AllocU8,
89                                                                         AllocU32,
90                                                                         AllocHC> {
91  	fn read(&mut self, buf: &mut [u8]) -> Result<usize, Error> {
92       self.0.read(buf)
93    }
94}
95
96
97#[cfg(not(any(feature="unsafe", not(feature="std"))))]
98pub struct Decompressor<R: Read>(DecompressorCustomAlloc<R,
99                                                         <StandardAlloc
100                                                          as Allocator<u8>>::AllocatedMemory,
101                                                         StandardAlloc,
102                                                         StandardAlloc,
103                                                         StandardAlloc>);
104
105
106#[cfg(not(any(feature="unsafe", not(feature="std"))))]
107impl<R: Read> Decompressor<R> {
108  pub fn new(r: R, buffer_size: usize) -> Self {
109     let dict = <StandardAlloc as Allocator<u8>>::AllocatedMemory::default();
110     Self::new_with_custom_dict(r, buffer_size, dict)
111  }
112  pub fn new_with_custom_dict(r: R, buffer_size: usize, dict: <StandardAlloc as Allocator<u8>>::AllocatedMemory) -> Self {
113    let mut alloc = StandardAlloc::default();
114    let buffer = <StandardAlloc as Allocator<u8>>::alloc_cell(&mut alloc, if buffer_size == 0 {4096} else {buffer_size});
115    Decompressor::<R>(DecompressorCustomAlloc::<R,
116                                                <StandardAlloc
117                                                 as Allocator<u8>>::AllocatedMemory,
118                                                StandardAlloc,
119                                                StandardAlloc,
120                                                StandardAlloc>::new_with_custom_dictionary(r,
121                                                                              buffer,
122                                                                              alloc,
123                                                                              StandardAlloc::default(),
124                                                                              StandardAlloc::default(),
125                                                                              dict))
126  }
127
128  // Attaches an additional raw LZ77 prefix dictionary; only allowed before
129  // the first read. Returns false if the dictionary could not be attached.
130  pub fn attach_dictionary(&mut self, dict: <StandardAlloc as Allocator<u8>>::AllocatedMemory) -> bool {
131    self.0.attach_dictionary(dict)
132  }
133
134  // Attaches a serialized shared dictionary (magic 0x91 0x00, may contain an
135  // LZ77 prefix dictionary and custom word/transform lists); only allowed
136  // before the first read.
137  pub fn attach_serialized_dictionary(&mut self, dict: <StandardAlloc as Allocator<u8>>::AllocatedMemory) -> bool {
138    self.0.attach_serialized_dictionary(dict)
139  }
140
141  pub fn get_ref(&self) -> &R {
142    &self.0.get_ref()
143  }
144  pub fn get_mut(&mut self) -> &mut R {
145    &mut ((self.0).0).get_mut().0
146  }
147  pub fn into_inner(self) -> R {
148    self.0.into_inner()
149  }
150}
151
152
153#[cfg(all(feature="unsafe", feature="std"))]
154pub struct Decompressor<R: Read>(DecompressorCustomAlloc<R,
155                                                         <HeapAlloc<u8>
156                                                          as Allocator<u8>>::AllocatedMemory,
157                                                         HeapAlloc<u8>,
158                                                         HeapAlloc<u32>,
159                                                         HeapAlloc<HuffmanCode> >);
160
161
162#[cfg(all(feature="unsafe", feature="std"))]
163impl<R: Read> Decompressor<R> {
164  pub fn new(r: R, buffer_size: usize) -> Self {
165     let dict = <HeapAlloc<u8> as Allocator<u8>>::AllocatedMemory::default();
166     Self::new_with_custom_dictionary(r, buffer_size, dict)
167  }
168  pub fn new_with_custom_dictionary(r: R, buffer_size: usize, dict: <HeapAlloc<u8>
169                                                 as Allocator<u8>>::AllocatedMemory) -> Self {
170    let mut alloc_u8 = HeapAlloc::<u8>::new(0);
171    let buffer = alloc_u8.alloc_cell(if buffer_size == 0 {4096} else {buffer_size});
172    let alloc_u32 = HeapAlloc::<u32>::new(0);
173    let alloc_hc = HeapAlloc::<HuffmanCode>::new(HuffmanCode{
174        bits:0, value: 0,
175    });
176    Decompressor::<R>(DecompressorCustomAlloc::<R,
177                                                <HeapAlloc<u8>
178                                                 as Allocator<u8>>::AllocatedMemory,
179                                                HeapAlloc<u8>,
180                                                HeapAlloc<u32>,
181                                                HeapAlloc<HuffmanCode> >
182      ::new_with_custom_dictionary(r, buffer, alloc_u8, alloc_u32, alloc_hc, dict))
183  }
184
185  // Attaches an additional raw LZ77 prefix dictionary; only allowed before
186  // the first read. Returns false if the dictionary could not be attached.
187  pub fn attach_dictionary(&mut self, dict: <HeapAlloc<u8> as Allocator<u8>>::AllocatedMemory) -> bool {
188    self.0.attach_dictionary(dict)
189  }
190
191  // Attaches a serialized shared dictionary (magic 0x91 0x00, may contain an
192  // LZ77 prefix dictionary and custom word/transform lists); only allowed
193  // before the first read.
194  pub fn attach_serialized_dictionary(&mut self, dict: <HeapAlloc<u8> as Allocator<u8>>::AllocatedMemory) -> bool {
195    self.0.attach_serialized_dictionary(dict)
196  }
197
198  pub fn get_ref(&self) -> &R {
199    self.0.get_ref()
200  }
201  pub fn get_mut(&mut self) -> &mut R {
202    &mut (self.0).0.get_mut().0
203  }
204  pub fn into_inner(self) -> R {
205    self.0.into_inner()
206  }
207}
208
209
210#[cfg(feature="std")]
211impl<R: Read> Read for Decompressor<R> {
212  fn read(&mut self, buf: &mut [u8]) -> Result<usize, Error> {
213    self.0.read(buf)
214  }
215}
216
217pub struct DecompressorCustomIo<ErrType,
218                                R: CustomRead<ErrType>,
219                                BufferType: SliceWrapperMut<u8>,
220                                AllocU8: Allocator<u8>,
221                                AllocU32: Allocator<u32>,
222                                AllocHC: Allocator<HuffmanCode>>
223{
224  input_buffer: BufferType,
225  total_out: usize,
226  input_offset: usize,
227  input_len: usize,
228  input: R,
229  error_if_invalid_data: Option<ErrType>,
230  state: BrotliState<AllocU8, AllocU32, AllocHC>,
231  done: bool,
232}
233
234impl<ErrType,
235     R: CustomRead<ErrType>,
236     BufferType : SliceWrapperMut<u8>,
237     AllocU8,
238     AllocU32,
239     AllocHC> DecompressorCustomIo<ErrType, R, BufferType, AllocU8, AllocU32, AllocHC>
240 where AllocU8 : Allocator<u8>, AllocU32 : Allocator<u32>, AllocHC : Allocator<HuffmanCode>
241{
242
243    pub fn new(r: R, buffer : BufferType,
244               alloc_u8 : AllocU8, alloc_u32 : AllocU32, alloc_hc : AllocHC,
245               invalid_data_error_type : ErrType) -> Self {
246     let dict = AllocU8::AllocatedMemory::default();
247     Self::new_with_custom_dictionary(r, buffer, alloc_u8, alloc_u32, alloc_hc, dict, invalid_data_error_type)
248    }
249    pub fn new_with_custom_dictionary(r: R, buffer : BufferType,
250               alloc_u8 : AllocU8, alloc_u32 : AllocU32, alloc_hc : AllocHC,
251               dict: AllocU8::AllocatedMemory,
252               invalid_data_error_type : ErrType) -> Self {
253        DecompressorCustomIo::<ErrType, R, BufferType, AllocU8, AllocU32, AllocHC>{
254            input_buffer : buffer,
255            total_out : 0,
256            input_offset : 0,
257            input_len : 0,
258            input: r,
259            state : BrotliState::new_with_custom_dictionary(alloc_u8,
260                                     alloc_u32,
261                                     alloc_hc,
262                                     dict),
263            error_if_invalid_data : Some(invalid_data_error_type),
264            done: false,
265        }
266    }
267
268    // Attaches an additional raw LZ77 prefix dictionary; only allowed before
269    // the first read. Returns false if the dictionary could not be attached.
270    pub fn attach_dictionary(&mut self, dict: AllocU8::AllocatedMemory) -> bool {
271      self.state.attach_dictionary(dict)
272    }
273
274    // Attaches a serialized shared dictionary (magic 0x91 0x00, may contain
275    // an LZ77 prefix dictionary and custom word/transform lists); only
276    // allowed before the first read.
277    pub fn attach_serialized_dictionary(&mut self, dict: AllocU8::AllocatedMemory) -> bool {
278      self.state.attach_serialized_dictionary(dict)
279    }
280
281    pub fn get_ref(&self) -> &R {
282      &self.input
283    }
284    pub fn get_mut(&mut self) -> &mut R {
285      &mut self.input
286    }
287    pub fn into_inner(self) -> R {
288      match self {
289        DecompressorCustomIo {
290          input_buffer: _ib,
291          total_out: _to,
292          state: _state,
293          input_offset: _io,
294          input_len: _il,
295          error_if_invalid_data:_eiid,
296          input,
297          done: _done,
298        } =>{
299          input
300        }
301    }
302    }
303
304    pub fn copy_to_front(&mut self) {
305        let avail_in = self.input_len - self.input_offset;
306        if self.input_offset == self.input_buffer.slice_mut().len() {
307            self.input_offset = 0;
308            self.input_len = 0;
309        } else if self.input_offset + 256 > self.input_buffer.slice_mut().len() && avail_in < self.input_offset {
310            let (first, second) = self.input_buffer.slice_mut().split_at_mut(self.input_offset);
311            self.input_len -= self.input_offset;
312            first[0..avail_in].clone_from_slice(&second[0..avail_in]);
313            self.input_offset = 0;
314        }
315    }
316}
317
318impl<ErrType,
319     R: CustomRead<ErrType>,
320     BufferType : SliceWrapperMut<u8>,
321     AllocU8 : Allocator<u8>,
322     AllocU32 : Allocator<u32>,
323     AllocHC : Allocator<HuffmanCode> > CustomRead<ErrType> for DecompressorCustomIo<ErrType,
324                                                                                     R,
325                                                                                     BufferType,
326                                                                                     AllocU8,
327                                                                                     AllocU32,
328                                                                                     AllocHC> {
329  /// This variant of read will return Ok(number of bytes read) until the file
330  /// Is completed at which point it will return Ok(0).
331  /// However if there are additional unconsumed bytes in the buffer, it will
332  /// return Err(InvalidData) at that point. Otherwise it will keep returning
333  /// Ok(0).
334  ///
335  /// # Arguments
336  ///
337  /// * `buf` - The buffer to read into
338  ///
339  /// # Errors
340  ///
341  /// Returns Ok(0) if the file has been fully decompressed.
342  /// If the file has been fully decompressed but there are additional
343  /// non-brotli bytes in the buffer, then return an InvalidData error.
344  /// Also upstream errors from the reader are returned.
345  fn read(&mut self, buf: &mut [u8]) -> Result<usize, ErrType > {
346    let mut output_offset : usize = 0;
347    let mut avail_out = buf.len() - output_offset;
348    let mut avail_in = self.input_len - self.input_offset;
349    while avail_out == buf.len() {
350      match BrotliDecompressStream(&mut avail_in,
351                                   &mut self.input_offset,
352                                   &self.input_buffer.slice_mut()[..],
353                                   &mut avail_out,
354                                   &mut output_offset,
355                                   buf,
356                                   &mut self.total_out,
357                                   &mut self.state) {
358        BrotliResult::NeedsMoreInput => {
359          self.copy_to_front();
360          if output_offset != 0 {
361            // The decompressor successfully decoded some bytes, but still requires more
362            // we do not wish to risk self.input.read returning an error, so instead we
363            // opt to return what we have and do not invoke the read trait method
364            return Ok(output_offset);
365          }
366          match self.input.read(&mut self.input_buffer.slice_mut()[self.input_len..]) {
367            Err(e) => {
368              return Err(e);
369            },
370            Ok(size) => if size == 0 {
371              return self.error_if_invalid_data.take().map(|e| Err(e)).unwrap_or(Ok(0));
372            }else {
373              self.input_len += size;
374              avail_in = self.input_len - self.input_offset;
375            },
376          }
377        },
378        BrotliResult::NeedsMoreOutput => {
379          break;
380        },
381        BrotliResult::ResultSuccess => {
382            if output_offset == 0 {
383                if !self.done {
384                    self.done = true;
385                } else if self.input_len != self.input_offset {
386                    // Did not consume entire input; report error.
387                    return self.error_if_invalid_data.take().map(|e| Err(e)).unwrap_or(Ok(output_offset));
388                }
389            }
390            return Ok(output_offset);
391        }
392        BrotliResult::ResultFailure => return self.error_if_invalid_data.take().map(|e| Err(e)).unwrap_or(Ok(0)),
393      }
394    }
395    Ok(output_offset)
396  }
397}
398
399#[cfg(feature="std")]
400#[test]
401fn test_no_vanishing_bytes() {
402    use std::string::ToString;
403
404    // Output from this command:
405    let compressed_with_extra = b"\x8f\x02\x80\x68\x65\x6c\x6c\x6f\x0a\x03\x67\x6f\x6f\x64\x62\x79\x65\x0a";
406    // Make sure that read_to_string returns the data.
407    let cursor = std::io::Cursor::new(compressed_with_extra);
408    let mut reader = super::Decompressor::new(cursor, 8000);
409    assert_eq!(std::io::read_to_string(&mut reader).unwrap(), "hello\n");
410
411    // However you can call read extra times to make sure there's no data.
412    let cursor = std::io::Cursor::new(compressed_with_extra);
413    let mut reader = super::Decompressor::new(cursor, 8000);
414    let mut data = std::vec::Vec::<u8>::default();
415    loop {
416        let mut buf = [0u8;5];
417        let offset = reader.read(&mut buf).unwrap();
418        if offset == 0 {
419            break;
420        }
421        data.extend_from_slice(&buf[..offset]);
422    }
423    assert_eq!(
424        &data,
425        &['h' as u8, 'e' as u8, 'l' as u8, 'l' as u8, 'o' as u8, '\n' as u8]);
426
427    // But calling read, one last time, results in an error because there
428    // were leftover bytes in the buffer.
429    let mut buf = [0u8;5];
430    assert_eq!(reader.read(&mut buf).unwrap_err().kind(),
431               io::ErrorKind::InvalidData);
432    data.clear();
433
434
435}
436
437#[cfg(feature="std")]
438#[test]
439fn test_repeated_read_returns_zero() {
440    use std::string::ToString;
441
442    // Output from this command:
443    let compressed_without_extra = b"\x8f\x02\x80\x68\x65\x6c\x6c\x6f\x0a\x03";
444    // Make sure that read_to_string returns the data.
445    let cursor = std::io::Cursor::new(compressed_without_extra);
446    let mut reader = super::Decompressor::new(cursor, 8000);
447    assert_eq!(std::io::read_to_string(&mut reader).unwrap(), "hello\n");
448
449    // However you can call read extra times to make sure there's no data.
450    let cursor = std::io::Cursor::new(compressed_without_extra);
451    let mut reader = super::Decompressor::new(cursor, 8000);
452    let mut data = std::vec::Vec::<u8>::default();
453    loop {
454        let mut buf = [0u8;5];
455        let offset = reader.read(&mut buf).unwrap();
456        if offset == 0 {
457            break;
458        }
459        data.extend_from_slice(&buf[..offset]);
460    }
461    assert_eq!(&data, &['h' as u8, 'e' as u8, 'l' as u8, 'l' as u8, 'o' as u8, '\n' as u8]);
462    let mut buf = [0u8;5];
463    assert_eq!(reader.read(&mut buf).unwrap(), 0);
464    data.clear();
465
466
467}
468