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