brotli_decompressor/
writer.rs

1use core;
2#[cfg(feature="std")]
3use std::io::{self, Error, ErrorKind, Write};
4#[cfg(feature="std")]
5pub use alloc_stdlib::StandardAlloc;
6#[cfg(all(feature="unsafe",feature="std"))]
7pub use alloc_stdlib::HeapAlloc;
8pub use huffman::{HuffmanCode, HuffmanTreeGroup};
9pub use state::BrotliState;
10// use io_wrappers::write_all;
11pub use io_wrappers::{CustomWrite};
12#[cfg(feature="std")]
13pub use io_wrappers::{IntoIoWriter, IoWriterWrapper};
14pub use super::decode::{BrotliDecompressStream, BrotliResult};
15pub use alloc::{AllocatedStackMemory, Allocator, SliceWrapper, SliceWrapperMut, StackAllocator};
16
17#[cfg(feature="std")]
18pub struct DecompressorWriterCustomAlloc<W: Write,
19     BufferType : SliceWrapperMut<u8>,
20     AllocU8 : Allocator<u8>,
21     AllocU32 : Allocator<u32>,
22     AllocHC : Allocator<HuffmanCode> >(DecompressorWriterCustomIo<io::Error,
23                                                             IntoIoWriter<W>,
24                                                             BufferType,
25                                                             AllocU8, AllocU32, AllocHC>);
26
27
28#[cfg(feature="std")]
29impl<W: Write,
30     BufferType : SliceWrapperMut<u8>,
31     AllocU8,
32     AllocU32,
33     AllocHC> DecompressorWriterCustomAlloc<W, BufferType, AllocU8, AllocU32, AllocHC>
34 where AllocU8 : Allocator<u8>, AllocU32 : Allocator<u32>, AllocHC : Allocator<HuffmanCode>
35    {
36    pub fn new(w: W, buffer : BufferType,
37               alloc_u8 : AllocU8, alloc_u32 : AllocU32, alloc_hc : AllocHC) -> Self {
38     let dict = AllocU8::AllocatedMemory::default();
39     Self::new_with_custom_dictionary(w, buffer, alloc_u8, alloc_u32, alloc_hc, dict)
40
41    }
42    pub fn new_with_custom_dictionary(w: W, buffer : BufferType,
43               alloc_u8 : AllocU8, alloc_u32 : AllocU32, alloc_hc : AllocHC, dict: AllocU8::AllocatedMemory) -> Self {
44        DecompressorWriterCustomAlloc::<W, BufferType, AllocU8, AllocU32, AllocHC>(
45          DecompressorWriterCustomIo::<Error,
46                                 IntoIoWriter<W>,
47                                 BufferType,
48                                 AllocU8, AllocU32, AllocHC>::new_with_custom_dictionary(IntoIoWriter::<W>(w),
49                                                                  buffer,
50                                                                  alloc_u8, alloc_u32, alloc_hc,
51                                                                  dict,
52                                                                  Error::new(ErrorKind::InvalidData,
53                                                                             "Invalid Data")))
54    }
55
56    pub fn attach_dictionary(&mut self, dict: AllocU8::AllocatedMemory) -> bool {
57      self.0.attach_dictionary(dict)
58    }
59    pub fn attach_serialized_dictionary(&mut self, dict: AllocU8::AllocatedMemory) -> bool {
60      self.0.attach_serialized_dictionary(dict)
61    }
62
63    pub fn get_ref(&self) -> &W {
64        &self.0.get_ref().0
65    }
66    pub fn get_mut(&mut self) -> &mut W {
67        &mut self.0.get_mut().0
68    }
69    pub fn into_inner(self) -> Result<W, W> {
70        match self.0.into_inner() {
71            Ok(w) => Ok(w.0),
72            Err(w) => Err(w.0),
73        }
74    }
75}
76#[cfg(feature="std")]
77impl<W: Write,
78     BufferType : SliceWrapperMut<u8>,
79     AllocU8 : Allocator<u8>,
80     AllocU32 : Allocator<u32>,
81     AllocHC : Allocator<HuffmanCode> > Write for DecompressorWriterCustomAlloc<W,
82                                                                         BufferType,
83                                                                         AllocU8,
84                                                                         AllocU32,
85                                                                         AllocHC> {
86  	fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
87       self.0.write(buf)
88    }
89  	fn flush(&mut self) -> Result<(), Error> {
90       self.0.flush()
91    }
92}
93
94#[cfg(feature="std")]
95impl<W: Write,
96     BufferType : SliceWrapperMut<u8>,
97     AllocU8 : Allocator<u8>,
98     AllocU32 : Allocator<u32>,
99     AllocHC : Allocator<HuffmanCode> > DecompressorWriterCustomAlloc<W,
100                                                                         BufferType,
101                                                                         AllocU8,
102                                                                         AllocU32,
103                                                                         AllocHC> {
104    pub fn close(&mut self) -> Result<(), Error>{
105        self.0.close()
106    }
107}
108
109
110#[cfg(not(any(feature="unsafe", not(feature="std"))))]
111pub struct DecompressorWriter<W: Write>(DecompressorWriterCustomAlloc<W,
112                                                         <StandardAlloc
113                                                          as Allocator<u8>>::AllocatedMemory,
114                                                         StandardAlloc,
115                                                         StandardAlloc,
116                                                         StandardAlloc>);
117
118
119#[cfg(not(any(feature="unsafe", not(feature="std"))))]
120impl<W: Write> DecompressorWriter<W> {
121  pub fn new(w: W, buffer_size: usize) -> Self {
122      Self::new_with_custom_dictionary(w, buffer_size, <StandardAlloc as Allocator<u8>>::AllocatedMemory::default())
123  }
124  pub fn new_with_custom_dictionary(w: W, buffer_size: usize, dict: <StandardAlloc as Allocator<u8>>::AllocatedMemory) -> Self {
125    let mut alloc = StandardAlloc::default();
126    let buffer = <StandardAlloc as Allocator<u8>>::alloc_cell(&mut alloc, if buffer_size == 0 {4096} else {buffer_size});
127    DecompressorWriter::<W>(DecompressorWriterCustomAlloc::<W,
128                                                <StandardAlloc
129                                                 as Allocator<u8>>::AllocatedMemory,
130                                                StandardAlloc,
131                                                StandardAlloc,
132                                                StandardAlloc>::new_with_custom_dictionary(w,
133                                                                              buffer,
134                                                                              alloc,
135                                                                              StandardAlloc::default(),
136                                                                              StandardAlloc::default(),
137                                                                              dict))
138  }
139
140  // Attaches an additional raw LZ77 prefix dictionary; only allowed before
141  // the first write. Returns false if the dictionary could not be attached.
142  pub fn attach_dictionary(&mut self, dict: <StandardAlloc as Allocator<u8>>::AllocatedMemory) -> bool {
143    self.0.attach_dictionary(dict)
144  }
145
146  // Attaches a serialized shared dictionary (magic 0x91 0x00, may contain an
147  // LZ77 prefix dictionary and custom word/transform lists); only allowed
148  // before the first write.
149  pub fn attach_serialized_dictionary(&mut self, dict: <StandardAlloc as Allocator<u8>>::AllocatedMemory) -> bool {
150    self.0.attach_serialized_dictionary(dict)
151  }
152
153  pub fn get_ref(&self) -> &W {
154      self.0.get_ref()
155  }
156  pub fn get_mut(&mut self) -> &mut W {
157      self.0.get_mut()
158  }
159  pub fn into_inner(self) -> Result<W, W> {
160    self.0.into_inner()
161  }
162}
163
164
165#[cfg(all(feature="unsafe", feature="std"))]
166pub struct DecompressorWriter<W: Write>(DecompressorWriterCustomAlloc<W,
167                                                         <HeapAlloc<u8>
168                                                          as Allocator<u8>>::AllocatedMemory,
169                                                         HeapAlloc<u8>,
170                                                         HeapAlloc<u32>,
171                                                         HeapAlloc<HuffmanCode> >);
172
173
174#[cfg(all(feature="unsafe", feature="std"))]
175impl<W: Write> DecompressorWriter<W> {
176  pub fn new(w: W, buffer_size: usize) -> Self {
177    let dict = <HeapAlloc<u8> as Allocator<u8>>::AllocatedMemory::default();
178    Self::new_with_custom_dictionary(w, buffer_size, dict)
179  }
180  pub fn new_with_custom_dictionary(w: W, buffer_size: usize, dict: <HeapAlloc<u8> as Allocator<u8>>::AllocatedMemory) -> Self {
181    let mut alloc_u8 = HeapAlloc::<u8>::new(0);
182    let buffer = alloc_u8.alloc_cell(if buffer_size == 0 {4096} else {buffer_size});
183    let alloc_u32 = HeapAlloc::<u32>::new(0);
184    let alloc_hc = HeapAlloc::<HuffmanCode>::new(HuffmanCode{bits:2, value: 1});
185    DecompressorWriter::<W>(DecompressorWriterCustomAlloc::<W,
186                                                <HeapAlloc<u8>
187                                                 as Allocator<u8>>::AllocatedMemory,
188                                                HeapAlloc<u8>,
189                                                HeapAlloc<u32>,
190                                                HeapAlloc<HuffmanCode> >
191      ::new_with_custom_dictionary(w, buffer, alloc_u8, alloc_u32, alloc_hc, dict))
192  }
193
194  // Attaches an additional raw LZ77 prefix dictionary; only allowed before
195  // the first write. Returns false if the dictionary could not be attached.
196  pub fn attach_dictionary(&mut self, dict: <HeapAlloc<u8> as Allocator<u8>>::AllocatedMemory) -> bool {
197    self.0.attach_dictionary(dict)
198  }
199
200  // Attaches a serialized shared dictionary (magic 0x91 0x00, may contain an
201  // LZ77 prefix dictionary and custom word/transform lists); only allowed
202  // before the first write.
203  pub fn attach_serialized_dictionary(&mut self, dict: <HeapAlloc<u8> as Allocator<u8>>::AllocatedMemory) -> bool {
204    self.0.attach_serialized_dictionary(dict)
205  }
206
207  pub fn get_ref(&self) -> &W {
208      self.0.get_ref()
209  }
210  pub fn get_mut(&mut self) -> &mut W {
211    &mut (self.0).0.get_mut().0
212  }
213  pub fn into_inner(self) -> Result<W, W> {
214    self.0.into_inner()
215  }
216}
217
218#[cfg(feature="std")]
219impl<W: Write> DecompressorWriter<W> {
220    pub fn close(&mut self) -> Result<(), Error>{
221        self.0.close()
222    }
223}
224#[cfg(feature="std")]
225impl<W: Write> Write for DecompressorWriter<W> {
226  	fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
227       self.0.write(buf)
228    }
229  	fn flush(&mut self) -> Result<(), Error> {
230       self.0.flush()
231    }
232}
233
234pub struct DecompressorWriterCustomIo<ErrType,
235                                W: CustomWrite<ErrType>,
236                                BufferType: SliceWrapperMut<u8>,
237                                AllocU8: Allocator<u8>,
238                                AllocU32: Allocator<u32>,
239                                AllocHC: Allocator<HuffmanCode>>
240{
241  output_buffer: BufferType,
242  total_out: usize,
243  output: Option<W>,
244  error_if_invalid_data: Option<ErrType>,
245  state: BrotliState<AllocU8, AllocU32, AllocHC>,
246}
247
248
249pub fn write_all<ErrType, W: CustomWrite<ErrType>>(writer: &mut W, mut buf : &[u8]) -> Result<(), ErrType> {
250    while buf.len() != 0 {
251          match writer.write(buf) {
252                Ok(bytes_written) => buf = &buf[bytes_written..],
253                Err(e) => return Err(e),
254          }
255    }
256    Ok(())
257}
258
259
260impl<ErrType,
261     W: CustomWrite<ErrType>,
262     BufferType : SliceWrapperMut<u8>,
263     AllocU8,
264     AllocU32,
265     AllocHC> DecompressorWriterCustomIo<ErrType, W, BufferType, AllocU8, AllocU32, AllocHC>
266 where AllocU8 : Allocator<u8>, AllocU32 : Allocator<u32>, AllocHC : Allocator<HuffmanCode>
267{
268
269    pub fn new(w: W, buffer : BufferType,
270               alloc_u8 : AllocU8, alloc_u32 : AllocU32, alloc_hc : AllocHC,
271               invalid_data_error_type : ErrType) -> Self {
272           let dict = AllocU8::AllocatedMemory::default();
273           Self::new_with_custom_dictionary(w, buffer, alloc_u8, alloc_u32, alloc_hc, dict, invalid_data_error_type)
274    }
275    pub fn new_with_custom_dictionary(w: W, buffer : BufferType,
276               alloc_u8 : AllocU8, alloc_u32 : AllocU32, alloc_hc : AllocHC,
277               dict: AllocU8::AllocatedMemory,
278               invalid_data_error_type : ErrType) -> Self {
279        DecompressorWriterCustomIo::<ErrType, W, BufferType, AllocU8, AllocU32, AllocHC>{
280            output_buffer : buffer,
281            total_out : 0,
282            output: Some(w),
283            state : BrotliState::new_with_custom_dictionary(alloc_u8,
284                                                                 alloc_u32,
285                                                                 alloc_hc,
286                                                                 dict),
287            error_if_invalid_data : Some(invalid_data_error_type),
288        }
289    }
290    pub fn close(&mut self) -> Result<(), ErrType>{
291        loop {
292            let mut avail_in : usize = 0;
293            let mut input_offset : usize = 0;
294            let mut avail_out : usize = self.output_buffer.slice_mut().len();
295            let mut output_offset : usize = 0;
296            let ret = BrotliDecompressStream(
297                &mut avail_in,
298                &mut input_offset,
299                &[],
300                &mut avail_out,
301                &mut output_offset,
302                self.output_buffer.slice_mut(),                
303                &mut self.total_out,
304                &mut self.state);
305          // already closed.
306          if self.error_if_invalid_data.is_none() {
307              return Ok(());
308          }
309          match write_all(self.output.as_mut().unwrap(), &self.output_buffer.slice_mut()[..output_offset]) {
310            Ok(_) => {},
311            Err(e) => return Err(e),
312           }
313           match ret {
314           BrotliResult::NeedsMoreInput => return self.error_if_invalid_data.take().map(|e|Err(e)).unwrap_or(Ok(())),
315           BrotliResult::NeedsMoreOutput => {},
316           BrotliResult::ResultSuccess => {
317               return Ok(());
318           },
319           BrotliResult::ResultFailure => return self.error_if_invalid_data.take().map(|e|Err(e)).unwrap_or(Ok(()))
320           }
321        }
322    }
323
324    // Attaches an additional raw LZ77 prefix dictionary; only allowed before
325    // the first write. Returns false if the dictionary could not be attached.
326    pub fn attach_dictionary(&mut self, dict: AllocU8::AllocatedMemory) -> bool {
327      self.state.attach_dictionary(dict)
328    }
329
330    // Attaches a serialized shared dictionary (magic 0x91 0x00, may contain
331    // an LZ77 prefix dictionary and custom word/transform lists); only
332    // allowed before the first write.
333    pub fn attach_serialized_dictionary(&mut self, dict: AllocU8::AllocatedMemory) -> bool {
334      self.state.attach_serialized_dictionary(dict)
335    }
336
337    pub fn get_ref(&self) -> &W {
338        self.output.as_ref().unwrap()
339    }
340    pub fn get_mut(&mut self) -> &mut W {
341        self.output.as_mut().unwrap()
342    }
343    pub fn into_inner(mut self) -> Result<W, W> {
344        match self.close() {
345            Ok(_) => Ok((core::mem::replace(&mut self.output, None).unwrap())),
346            Err(_) => Err((core::mem::replace(&mut self.output, None).unwrap())),
347        }
348    }
349}
350
351impl<ErrType,
352     W: CustomWrite<ErrType>,
353     BufferType : SliceWrapperMut<u8>,
354     AllocU8 : Allocator<u8>,
355     AllocU32 : Allocator<u32>,
356     AllocHC : Allocator<HuffmanCode> > Drop for DecompressorWriterCustomIo<ErrType,
357                                                                                     W,
358                                                                                     BufferType,
359                                                                                     AllocU8,
360                                                                                     AllocU32,
361                                                                                     AllocHC> {
362    fn drop(&mut self) {
363        if self.output.is_some() {
364            match self.close() {
365                Ok(_) => {},
366                Err(_) => {},
367            }
368        }
369    }
370}
371
372impl<ErrType,
373     W: CustomWrite<ErrType>,
374     BufferType : SliceWrapperMut<u8>,
375     AllocU8 : Allocator<u8>,
376     AllocU32 : Allocator<u32>,
377     AllocHC : Allocator<HuffmanCode> > CustomWrite<ErrType> for DecompressorWriterCustomIo<ErrType,
378                                                                                     W,
379                                                                                     BufferType,
380                                                                                     AllocU8,
381                                                                                     AllocU32,
382                                                                                     AllocHC> {
383	fn write(&mut self, buf: &[u8]) -> Result<usize, ErrType > {
384        let mut avail_in = buf.len();
385        let mut input_offset : usize = 0;
386        loop {
387            let mut output_offset = 0;
388            let mut avail_out = self.output_buffer.slice_mut().len();
389            let op_result = BrotliDecompressStream(&mut avail_in,
390                                     &mut input_offset,
391                                     &buf[..],
392                                     &mut avail_out,
393                                     &mut output_offset,
394                                     self.output_buffer.slice_mut(),
395                                     &mut self.total_out,
396                                     &mut self.state);
397         match write_all(self.output.as_mut().unwrap(), &self.output_buffer.slice_mut()[..output_offset]) {
398          Ok(_) => {},
399          Err(e) => return Err(e),
400         }
401         match op_result {
402          BrotliResult::NeedsMoreInput => assert_eq!(avail_in, 0),
403          BrotliResult::NeedsMoreOutput => continue,
404          BrotliResult::ResultSuccess => {
405              return Ok(input_offset);
406          }
407          BrotliResult::ResultFailure => return self.error_if_invalid_data.take().map(|e|Err(e)).unwrap_or(Ok(0)),
408       }
409        if avail_in == 0 {
410           break
411        }
412      }
413      Ok(buf.len())
414    }
415    fn flush(&mut self) -> Result<(), ErrType> {
416       self.output.as_mut().unwrap().flush()
417    }
418}
419
420#[cfg(feature="std")]
421#[cfg(test)]
422mod test {
423    use super::DecompressorWriter;
424    use std::vec::Vec;
425    use std::io::Write;
426    // Brotli-compressed "hello\n" and 2 extra bytes
427
428
429    #[test]
430    fn write_extra() {
431        let contents = b"\x8f\x02\x80\x68\x65\x6c\x6c\x6f\x0a\x03\x67\x6f\x6f\x64\x62\x79\x65\x0a";
432        let mut decoder = DecompressorWriter::new(Vec::new(), 0);
433        let n = decoder.write(contents).unwrap();
434        assert_eq!(n, 10);
435        // Ensure that we can continue to not send data to the writer
436        // as it has consumed the entire file.
437        let n = decoder.write(contents).unwrap();
438        assert_eq!(n, 0);
439
440        let mut decoder = DecompressorWriter::new(Vec::new(), 0);
441        let e = decoder.write_all(contents).unwrap_err();
442        assert!(e.kind() == std::io::ErrorKind::WriteZero);
443        assert_eq!(decoder.get_ref().as_slice(), b"hello\n");
444    }
445}
brotli_decompressor/writer.rs

brotli_decompressor/
writer.rs
