brotli_decompressor/
writer.rs

Help
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 get_ref(&self) -> &W {
57        &self.0.get_ref().0
58    }
59    pub fn get_mut(&mut self) -> &mut W {
60        &mut self.0.get_mut().0
61    }
62    pub fn into_inner(self) -> Result<W, W> {
63        match self.0.into_inner() {
64            Ok(w) => Ok(w.0),
65            Err(w) => Err(w.0),
66        }
67    }
68}
69#[cfg(feature="std")]
70impl<W: Write,
71     BufferType : SliceWrapperMut<u8>,
72     AllocU8 : Allocator<u8>,
73     AllocU32 : Allocator<u32>,
74     AllocHC : Allocator<HuffmanCode> > Write for DecompressorWriterCustomAlloc<W,
75                                                                         BufferType,
76                                                                         AllocU8,
77                                                                         AllocU32,
78                                                                         AllocHC> {
79  	fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
80       self.0.write(buf)
81    }
82  	fn flush(&mut self) -> Result<(), Error> {
83       self.0.flush()
84    }
85}
86
87#[cfg(feature="std")]
88impl<W: Write,
89     BufferType : SliceWrapperMut<u8>,
90     AllocU8 : Allocator<u8>,
91     AllocU32 : Allocator<u32>,
92     AllocHC : Allocator<HuffmanCode> > DecompressorWriterCustomAlloc<W,
93                                                                         BufferType,
94                                                                         AllocU8,
95                                                                         AllocU32,
96                                                                         AllocHC> {
97    pub fn close(&mut self) -> Result<(), Error>{
98        self.0.close()
99    }
100}
101
102
103#[cfg(not(any(feature="unsafe", not(feature="std"))))]
104pub struct DecompressorWriter<W: Write>(DecompressorWriterCustomAlloc<W,
105                                                         <StandardAlloc
106                                                          as Allocator<u8>>::AllocatedMemory,
107                                                         StandardAlloc,
108                                                         StandardAlloc,
109                                                         StandardAlloc>);
110
111
112#[cfg(not(any(feature="unsafe", not(feature="std"))))]
113impl<W: Write> DecompressorWriter<W> {
114  pub fn new(w: W, buffer_size: usize) -> Self {
115      Self::new_with_custom_dictionary(w, buffer_size, <StandardAlloc as Allocator<u8>>::AllocatedMemory::default())
116  }
117  pub fn new_with_custom_dictionary(w: W, buffer_size: usize, dict: <StandardAlloc as Allocator<u8>>::AllocatedMemory) -> Self {
118    let mut alloc = StandardAlloc::default();
119    let buffer = <StandardAlloc as Allocator<u8>>::alloc_cell(&mut alloc, if buffer_size == 0 {4096} else {buffer_size});
120    DecompressorWriter::<W>(DecompressorWriterCustomAlloc::<W,
121                                                <StandardAlloc
122                                                 as Allocator<u8>>::AllocatedMemory,
123                                                StandardAlloc,
124                                                StandardAlloc,
125                                                StandardAlloc>::new_with_custom_dictionary(w,
126                                                                              buffer,
127                                                                              alloc,
128                                                                              StandardAlloc::default(),
129                                                                              StandardAlloc::default(),
130                                                                              dict))
131  }
132
133  pub fn get_ref(&self) -> &W {
134      self.0.get_ref()
135  }
136  pub fn get_mut(&mut self) -> &mut W {
137      self.0.get_mut()
138  }
139  pub fn into_inner(self) -> Result<W, W> {
140    self.0.into_inner()
141  }
142}
143
144
145#[cfg(all(feature="unsafe", feature="std"))]
146pub struct DecompressorWriter<W: Write>(DecompressorWriterCustomAlloc<W,
147                                                         <HeapAlloc<u8>
148                                                          as Allocator<u8>>::AllocatedMemory,
149                                                         HeapAlloc<u8>,
150                                                         HeapAlloc<u32>,
151                                                         HeapAlloc<HuffmanCode> >);
152
153
154#[cfg(all(feature="unsafe", feature="std"))]
155impl<W: Write> DecompressorWriter<W> {
156  pub fn new(w: W, buffer_size: usize) -> Self {
157    let dict = <HeapAlloc<u8> as Allocator<u8>>::AllocatedMemory::default();
158    Self::new_with_custom_dictionary(w, buffer_size, dict)
159  }
160  pub fn new_with_custom_dictionary(w: W, buffer_size: usize, dict: <HeapAlloc<u8> as Allocator<u8>>::AllocatedMemory) -> Self {
161    let mut alloc_u8 = HeapAlloc::<u8>::new(0);
162    let buffer = alloc_u8.alloc_cell(buffer_size);
163    let alloc_u32 = HeapAlloc::<u32>::new(0);
164    let alloc_hc = HeapAlloc::<HuffmanCode>::new(HuffmanCode{bits:2, value: 1});
165    DecompressorWriter::<W>(DecompressorWriterCustomAlloc::<W,
166                                                <HeapAlloc<u8>
167                                                 as Allocator<u8>>::AllocatedMemory,
168                                                HeapAlloc<u8>,
169                                                HeapAlloc<u32>,
170                                                HeapAlloc<HuffmanCode> >
171      ::new_with_custom_dictionary(w, buffer, alloc_u8, alloc_u32, alloc_hc, dict))
172  }
173
174  pub fn get_ref(&self) -> &W {
175      self.0.get_ref()
176  }
177  pub fn get_mut(&mut self) -> &mut W {
178    &mut (self.0).0.get_mut().0
179  }
180  pub fn into_inner(self) -> Result<W, W> {
181    self.0.into_inner()
182  }
183}
184
185#[cfg(feature="std")]
186impl<W: Write> DecompressorWriter<W> {
187    pub fn close(&mut self) -> Result<(), Error>{
188        self.0.close()
189    }
190}
191#[cfg(feature="std")]
192impl<W: Write> Write for DecompressorWriter<W> {
193  	fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
194       self.0.write(buf)
195    }
196  	fn flush(&mut self) -> Result<(), Error> {
197       self.0.flush()
198    }
199}
200
201pub struct DecompressorWriterCustomIo<ErrType,
202                                W: CustomWrite<ErrType>,
203                                BufferType: SliceWrapperMut<u8>,
204                                AllocU8: Allocator<u8>,
205                                AllocU32: Allocator<u32>,
206                                AllocHC: Allocator<HuffmanCode>>
207{
208  output_buffer: BufferType,
209  total_out: usize,
210  output: Option<W>,
211  error_if_invalid_data: Option<ErrType>,
212  state: BrotliState<AllocU8, AllocU32, AllocHC>,
213}
214
215
216pub fn write_all<ErrType, W: CustomWrite<ErrType>>(writer: &mut W, mut buf : &[u8]) -> Result<(), ErrType> {
217    while buf.len() != 0 {
218          match writer.write(buf) {
219                Ok(bytes_written) => buf = &buf[bytes_written..],
220                Err(e) => return Err(e),
221          }
222    }
223    Ok(())
224}
225
226
227impl<ErrType,
228     W: CustomWrite<ErrType>,
229     BufferType : SliceWrapperMut<u8>,
230     AllocU8,
231     AllocU32,
232     AllocHC> DecompressorWriterCustomIo<ErrType, W, BufferType, AllocU8, AllocU32, AllocHC>
233 where AllocU8 : Allocator<u8>, AllocU32 : Allocator<u32>, AllocHC : Allocator<HuffmanCode>
234{
235
236    pub fn new(w: W, buffer : BufferType,
237               alloc_u8 : AllocU8, alloc_u32 : AllocU32, alloc_hc : AllocHC,
238               invalid_data_error_type : ErrType) -> Self {
239           let dict = AllocU8::AllocatedMemory::default();
240           Self::new_with_custom_dictionary(w, buffer, alloc_u8, alloc_u32, alloc_hc, dict, invalid_data_error_type)
241    }
242    pub fn new_with_custom_dictionary(w: W, buffer : BufferType,
243               alloc_u8 : AllocU8, alloc_u32 : AllocU32, alloc_hc : AllocHC,
244               dict: AllocU8::AllocatedMemory,
245               invalid_data_error_type : ErrType) -> Self {
246        DecompressorWriterCustomIo::<ErrType, W, BufferType, AllocU8, AllocU32, AllocHC>{
247            output_buffer : buffer,
248            total_out : 0,
249            output: Some(w),
250            state : BrotliState::new_with_custom_dictionary(alloc_u8,
251                                                                 alloc_u32,
252                                                                 alloc_hc,
253                                                                 dict),
254            error_if_invalid_data : Some(invalid_data_error_type),
255        }
256    }
257    pub fn close(&mut self) -> Result<(), ErrType>{
258        loop {
259            let mut avail_in : usize = 0;
260            let mut input_offset : usize = 0;
261            let mut avail_out : usize = self.output_buffer.slice_mut().len();
262            let mut output_offset : usize = 0;
263            let ret = BrotliDecompressStream(
264                &mut avail_in,
265                &mut input_offset,
266                &[],
267                &mut avail_out,
268                &mut output_offset,
269                self.output_buffer.slice_mut(),                
270                &mut self.total_out,
271                &mut self.state);
272          // already closed.
273          if self.error_if_invalid_data.is_none() {
274              return Ok(());
275          }
276          match write_all(self.output.as_mut().unwrap(), &self.output_buffer.slice_mut()[..output_offset]) {
277            Ok(_) => {},
278            Err(e) => return Err(e),
279           }
280           match ret {
281           BrotliResult::NeedsMoreInput => return self.error_if_invalid_data.take().map(|e|Err(e)).unwrap_or(Ok(())),
282           BrotliResult::NeedsMoreOutput => {},
283           BrotliResult::ResultSuccess => {
284               return Ok(());
285           },
286           BrotliResult::ResultFailure => return self.error_if_invalid_data.take().map(|e|Err(e)).unwrap_or(Ok(()))
287           }
288        }
289    }
290
291    pub fn get_ref(&self) -> &W {
292        self.output.as_ref().unwrap()
293    }
294    pub fn get_mut(&mut self) -> &mut W {
295        self.output.as_mut().unwrap()
296    }
297    pub fn into_inner(mut self) -> Result<W, W> {
298        match self.close() {
299            Ok(_) => Ok((core::mem::replace(&mut self.output, None).unwrap())),
300            Err(_) => Err((core::mem::replace(&mut self.output, None).unwrap())),
301        }
302    }
303}
304
305impl<ErrType,
306     W: CustomWrite<ErrType>,
307     BufferType : SliceWrapperMut<u8>,
308     AllocU8 : Allocator<u8>,
309     AllocU32 : Allocator<u32>,
310     AllocHC : Allocator<HuffmanCode> > Drop for DecompressorWriterCustomIo<ErrType,
311                                                                                     W,
312                                                                                     BufferType,
313                                                                                     AllocU8,
314                                                                                     AllocU32,
315                                                                                     AllocHC> {
316    fn drop(&mut self) {
317        if self.output.is_some() {
318            match self.close() {
319                Ok(_) => {},
320                Err(_) => {},
321            }
322        }
323    }
324}
325
326impl<ErrType,
327     W: CustomWrite<ErrType>,
328     BufferType : SliceWrapperMut<u8>,
329     AllocU8 : Allocator<u8>,
330     AllocU32 : Allocator<u32>,
331     AllocHC : Allocator<HuffmanCode> > CustomWrite<ErrType> for DecompressorWriterCustomIo<ErrType,
332                                                                                     W,
333                                                                                     BufferType,
334                                                                                     AllocU8,
335                                                                                     AllocU32,
336                                                                                     AllocHC> {
337	fn write(&mut self, buf: &[u8]) -> Result<usize, ErrType > {
338        let mut avail_in = buf.len();
339        let mut input_offset : usize = 0;
340        loop {
341            let mut output_offset = 0;
342            let mut avail_out = self.output_buffer.slice_mut().len();
343            let op_result = BrotliDecompressStream(&mut avail_in,
344                                     &mut input_offset,
345                                     &buf[..],
346                                     &mut avail_out,
347                                     &mut output_offset,
348                                     self.output_buffer.slice_mut(),
349                                     &mut self.total_out,
350                                     &mut self.state);
351         match write_all(self.output.as_mut().unwrap(), &self.output_buffer.slice_mut()[..output_offset]) {
352          Ok(_) => {},
353          Err(e) => return Err(e),
354         }
355         match op_result {
356          BrotliResult::NeedsMoreInput => assert_eq!(avail_in, 0),
357          BrotliResult::NeedsMoreOutput => continue,
358          BrotliResult::ResultSuccess => {
359              return Ok(input_offset);
360          }
361          BrotliResult::ResultFailure => return self.error_if_invalid_data.take().map(|e|Err(e)).unwrap_or(Ok(0)),
362       }
363        if avail_in == 0 {
364           break
365        }
366      }
367      Ok(buf.len())
368    }
369    fn flush(&mut self) -> Result<(), ErrType> {
370       self.output.as_mut().unwrap().flush()
371    }
372}
373
374#[cfg(feature="std")]
375#[cfg(test)]
376mod test {
377    use super::DecompressorWriter;
378    use std::vec::Vec;
379    use std::io::Write;
380    // Brotli-compressed "hello\n" and 2 extra bytes
381
382
383    #[test]
384    fn write_extra() {
385        let contents = b"\x8f\x02\x80\x68\x65\x6c\x6c\x6f\x0a\x03\x67\x6f\x6f\x64\x62\x79\x65\x0a";
386        let mut decoder = DecompressorWriter::new(Vec::new(), 0);
387        let n = decoder.write(contents).unwrap();
388        assert_eq!(n, 10);
389        // Ensure that we can continue to not send data to the writer
390        // as it has consumed the entire file.
391        let n = decoder.write(contents).unwrap();
392        assert_eq!(n, 0);
393
394        let mut decoder = DecompressorWriter::new(Vec::new(), 0);
395        let e = decoder.write_all(contents).unwrap_err();
396        assert!(e.kind() == std::io::ErrorKind::WriteZero);
397        assert_eq!(decoder.get_ref().as_slice(), b"hello\n");
398    }
399}
brotli_decompressor/writer.rs

brotli_decompressor/
writer.rs
