1 use core;
2 #[cfg(feature="std")]
3 use std::io::{self, Error, ErrorKind, Write};
4 #[cfg(feature="std")]
5 pub use alloc_stdlib::StandardAlloc;
6 #[cfg(all(feature="unsafe",feature="std"))]
7 pub use alloc_stdlib::HeapAllocUninitialized;
8 pub use huffman::{HuffmanCode, HuffmanTreeGroup};
9 pub use state::BrotliState;
10 // use io_wrappers::write_all;
11 pub use io_wrappers::{CustomWrite};
12 #[cfg(feature="std")]
13 pub use io_wrappers::{IntoIoWriter, IoWriterWrapper};
14 pub use super::decode::{BrotliDecompressStream, BrotliResult};
15 pub use alloc::{AllocatedStackMemory, Allocator, SliceWrapper, SliceWrapperMut, StackAllocator};
16 
17 #[cfg(feature="std")]
18 pub 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")]
29 impl<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     {
new(w: W, buffer : BufferType, alloc_u8 : AllocU8, alloc_u32 : AllocU32, alloc_hc : AllocHC) -> Self36     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     }
new_with_custom_dictionary(w: W, buffer : BufferType, alloc_u8 : AllocU8, alloc_u32 : AllocU32, alloc_hc : AllocHC, dict: AllocU8::AllocatedMemory) -> Self42     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 
get_ref(&self) -> &W56     pub fn get_ref(&self) -> &W {
57         &self.0.get_ref().0
58     }
get_mut(&mut self) -> &mut W59     pub fn get_mut(&mut self) -> &mut W {
60         &mut self.0.get_mut().0
61     }
into_inner(self) -> Result<W, W>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")]
70 impl<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> {
write(&mut self, buf: &[u8]) -> Result<usize, Error>79   	fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
80        self.0.write(buf)
81     }
flush(&mut self) -> Result<(), Error>82   	fn flush(&mut self) -> Result<(), Error> {
83        self.0.flush()
84     }
85 }
86 
87 
88 #[cfg(not(any(feature="unsafe", not(feature="std"))))]
89 pub struct DecompressorWriter<W: Write>(DecompressorWriterCustomAlloc<W,
90                                                          <StandardAlloc
91                                                           as Allocator<u8>>::AllocatedMemory,
92                                                          StandardAlloc,
93                                                          StandardAlloc,
94                                                          StandardAlloc>);
95 
96 
97 #[cfg(not(any(feature="unsafe", not(feature="std"))))]
98 impl<W: Write> DecompressorWriter<W> {
new(w: W, buffer_size: usize) -> Self99   pub fn new(w: W, buffer_size: usize) -> Self {
100       Self::new_with_custom_dictionary(w, buffer_size, <StandardAlloc as Allocator<u8>>::AllocatedMemory::default())
101   }
new_with_custom_dictionary(w: W, buffer_size: usize, dict: <StandardAlloc as Allocator<u8>>::AllocatedMemory) -> Self102   pub fn new_with_custom_dictionary(w: W, buffer_size: usize, dict: <StandardAlloc as Allocator<u8>>::AllocatedMemory) -> Self {
103     let mut alloc = StandardAlloc::default();
104     let buffer = <StandardAlloc as Allocator<u8>>::alloc_cell(&mut alloc, if buffer_size == 0 {4096} else {buffer_size});
105     DecompressorWriter::<W>(DecompressorWriterCustomAlloc::<W,
106                                                 <StandardAlloc
107                                                  as Allocator<u8>>::AllocatedMemory,
108                                                 StandardAlloc,
109                                                 StandardAlloc,
110                                                 StandardAlloc>::new_with_custom_dictionary(w,
111                                                                               buffer,
112                                                                               alloc,
113                                                                               StandardAlloc::default(),
114                                                                               StandardAlloc::default(),
115                                                                               dict))
116   }
117 
get_ref(&self) -> &W118   pub fn get_ref(&self) -> &W {
119       self.0.get_ref()
120   }
get_mut(&mut self) -> &mut W121   pub fn get_mut(&mut self) -> &mut W {
122       self.0.get_mut()
123   }
into_inner(self) -> Result<W, W>124   pub fn into_inner(self) -> Result<W, W> {
125     self.0.into_inner()
126   }
127 }
128 
129 
130 #[cfg(all(feature="unsafe", feature="std"))]
131 pub struct DecompressorWriter<W: Write>(DecompressorWriterCustomAlloc<W,
132                                                          <HeapAllocUninitialized<u8>
133                                                           as Allocator<u8>>::AllocatedMemory,
134                                                          HeapAllocUninitialized<u8>,
135                                                          HeapAllocUninitialized<u32>,
136                                                          HeapAllocUninitialized<HuffmanCode> >);
137 
138 
139 #[cfg(all(feature="unsafe", feature="std"))]
140 impl<W: Write> DecompressorWriter<W> {
new(w: W, buffer_size: usize) -> Self141   pub fn new(w: W, buffer_size: usize) -> Self {
142     let dict = <HeapAllocUninitialized<u8> as Allocator<u8>>::AllocatedMemory::default();
143     Self::new_with_custom_dictionary(w, buffer_size, dict)
144   }
new_with_custom_dictionary(w: W, buffer_size: usize, dict: <HeapAllocUninitialized<u8> as Allocator<u8>>::AllocatedMemory) -> Self145   pub fn new_with_custom_dictionary(w: W, buffer_size: usize, dict: <HeapAllocUninitialized<u8> as Allocator<u8>>::AllocatedMemory) -> Self {
146     let mut alloc_u8 = unsafe { HeapAllocUninitialized::<u8>::new() };
147     let buffer = alloc_u8.alloc_cell(buffer_size);
148     let alloc_u32 = unsafe { HeapAllocUninitialized::<u32>::new() };
149     let alloc_hc = unsafe { HeapAllocUninitialized::<HuffmanCode>::new() };
150     DecompressorWriter::<W>(DecompressorWriterCustomAlloc::<W,
151                                                 <HeapAllocUninitialized<u8>
152                                                  as Allocator<u8>>::AllocatedMemory,
153                                                 HeapAllocUninitialized<u8>,
154                                                 HeapAllocUninitialized<u32>,
155                                                 HeapAllocUninitialized<HuffmanCode> >
156       ::new_with_custom_dictionary(w, buffer, alloc_u8, alloc_u32, alloc_hc, dict))
157   }
158 
get_ref(&self) -> &W159   pub fn get_ref(&self) -> &W {
160       self.0.get_ref()
161   }
get_mut(&mut self) -> &mut W162   pub fn get_mut(&mut self) -> &mut W {
163     &mut (self.0).0.get_mut().0
164   }
into_inner(self) -> Result<W, W>165   pub fn into_inner(self) -> Result<W, W> {
166     self.0.into_inner()
167   }
168 }
169 
170 
171 #[cfg(feature="std")]
172 impl<W: Write> Write for DecompressorWriter<W> {
write(&mut self, buf: &[u8]) -> Result<usize, Error>173   	fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
174        self.0.write(buf)
175     }
flush(&mut self) -> Result<(), Error>176   	fn flush(&mut self) -> Result<(), Error> {
177        self.0.flush()
178     }
179 }
180 
181 pub struct DecompressorWriterCustomIo<ErrType,
182                                 W: CustomWrite<ErrType>,
183                                 BufferType: SliceWrapperMut<u8>,
184                                 AllocU8: Allocator<u8>,
185                                 AllocU32: Allocator<u32>,
186                                 AllocHC: Allocator<HuffmanCode>>
187 {
188   output_buffer: BufferType,
189   total_out: usize,
190   output: Option<W>,
191   error_if_invalid_data: Option<ErrType>,
192   state: BrotliState<AllocU8, AllocU32, AllocHC>,
193 }
194 
195 
write_all<ErrType, W: CustomWrite<ErrType>>(writer: &mut W, mut buf : &[u8]) -> Result<(), ErrType>196 pub fn write_all<ErrType, W: CustomWrite<ErrType>>(writer: &mut W, mut buf : &[u8]) -> Result<(), ErrType> {
197     while buf.len() != 0 {
198           match writer.write(buf) {
199                 Ok(bytes_written) => buf = &buf[bytes_written..],
200                 Err(e) => return Err(e),
201           }
202     }
203     Ok(())
204 }
205 
206 
207 impl<ErrType,
208      W: CustomWrite<ErrType>,
209      BufferType : SliceWrapperMut<u8>,
210      AllocU8,
211      AllocU32,
212      AllocHC> DecompressorWriterCustomIo<ErrType, W, BufferType, AllocU8, AllocU32, AllocHC>
213  where AllocU8 : Allocator<u8>, AllocU32 : Allocator<u32>, AllocHC : Allocator<HuffmanCode>
214 {
215 
new(w: W, buffer : BufferType, alloc_u8 : AllocU8, alloc_u32 : AllocU32, alloc_hc : AllocHC, invalid_data_error_type : ErrType) -> Self216     pub fn new(w: W, buffer : BufferType,
217                alloc_u8 : AllocU8, alloc_u32 : AllocU32, alloc_hc : AllocHC,
218                invalid_data_error_type : ErrType) -> Self {
219            let dict = AllocU8::AllocatedMemory::default();
220            Self::new_with_custom_dictionary(w, buffer, alloc_u8, alloc_u32, alloc_hc, dict, invalid_data_error_type)
221     }
new_with_custom_dictionary(w: W, buffer : BufferType, alloc_u8 : AllocU8, alloc_u32 : AllocU32, alloc_hc : AllocHC, dict: AllocU8::AllocatedMemory, invalid_data_error_type : ErrType) -> Self222     pub fn new_with_custom_dictionary(w: W, buffer : BufferType,
223                alloc_u8 : AllocU8, alloc_u32 : AllocU32, alloc_hc : AllocHC,
224                dict: AllocU8::AllocatedMemory,
225                invalid_data_error_type : ErrType) -> Self {
226         DecompressorWriterCustomIo::<ErrType, W, BufferType, AllocU8, AllocU32, AllocHC>{
227             output_buffer : buffer,
228             total_out : 0,
229             output: Some(w),
230             state : BrotliState::new_with_custom_dictionary(alloc_u8,
231                                                                  alloc_u32,
232                                                                  alloc_hc,
233                                                                  dict),
234             error_if_invalid_data : Some(invalid_data_error_type),
235         }
236     }
close(&mut self) -> Result<(), ErrType>237     fn close(&mut self) -> Result<(), ErrType>{
238         loop {
239             let mut avail_in : usize = 0;
240             let mut input_offset : usize = 0;
241             let mut avail_out : usize = self.output_buffer.slice_mut().len();
242             let mut output_offset : usize = 0;
243             let ret = BrotliDecompressStream(
244                 &mut avail_in,
245                 &mut input_offset,
246                 &[],
247                 &mut avail_out,
248                 &mut output_offset,
249                 self.output_buffer.slice_mut(),
250                 &mut self.total_out,
251                 &mut self.state);
252           match write_all(self.output.as_mut().unwrap(), &self.output_buffer.slice_mut()[..output_offset]) {
253             Ok(_) => {},
254             Err(e) => return Err(e),
255            }
256            match ret {
257            BrotliResult::NeedsMoreInput => return Err(self.error_if_invalid_data.take().unwrap()),
258            BrotliResult::NeedsMoreOutput => {},
259            BrotliResult::ResultSuccess => return Ok(()),
260            BrotliResult::ResultFailure => return Err(self.error_if_invalid_data.take().unwrap()),
261            }
262         }
263     }
264 
get_ref(&self) -> &W265     pub fn get_ref(&self) -> &W {
266         self.output.as_ref().unwrap()
267     }
get_mut(&mut self) -> &mut W268     pub fn get_mut(&mut self) -> &mut W {
269         self.output.as_mut().unwrap()
270     }
into_inner(mut self) -> Result<W, W>271     pub fn into_inner(mut self) -> Result<W, W> {
272         match self.close() {
273             Ok(_) => Ok((core::mem::replace(&mut self.output, None).unwrap())),
274             Err(_) => Err((core::mem::replace(&mut self.output, None).unwrap())),
275         }
276     }
277 }
278 
279 impl<ErrType,
280      W: CustomWrite<ErrType>,
281      BufferType : SliceWrapperMut<u8>,
282      AllocU8 : Allocator<u8>,
283      AllocU32 : Allocator<u32>,
284      AllocHC : Allocator<HuffmanCode> > Drop for DecompressorWriterCustomIo<ErrType,
285                                                                                      W,
286                                                                                      BufferType,
287                                                                                      AllocU8,
288                                                                                      AllocU32,
289                                                                                      AllocHC> {
drop(&mut self)290     fn drop(&mut self) {
291         if self.output.is_some() {
292             match self.close() {
293                 Ok(_) => {},
294                 Err(_) => {},
295             }
296         }
297     }
298 }
299 
300 impl<ErrType,
301      W: CustomWrite<ErrType>,
302      BufferType : SliceWrapperMut<u8>,
303      AllocU8 : Allocator<u8>,
304      AllocU32 : Allocator<u32>,
305      AllocHC : Allocator<HuffmanCode> > CustomWrite<ErrType> for DecompressorWriterCustomIo<ErrType,
306                                                                                      W,
307                                                                                      BufferType,
308                                                                                      AllocU8,
309                                                                                      AllocU32,
310                                                                                      AllocHC> {
write(&mut self, buf: &[u8]) -> Result<usize, ErrType >311 	fn write(&mut self, buf: &[u8]) -> Result<usize, ErrType > {
312         let mut avail_in = buf.len();
313         let mut input_offset : usize = 0;
314         loop {
315             let mut output_offset = 0;
316             let mut avail_out = self.output_buffer.slice_mut().len();
317             let op_result = BrotliDecompressStream(&mut avail_in,
318                                      &mut input_offset,
319                                      &buf[..],
320                                      &mut avail_out,
321                                      &mut output_offset,
322                                      self.output_buffer.slice_mut(),
323                                      &mut self.total_out,
324                                      &mut self.state);
325          match write_all(self.output.as_mut().unwrap(), &self.output_buffer.slice_mut()[..output_offset]) {
326           Ok(_) => {},
327           Err(e) => return Err(e),
328          }
329          match op_result {
330           BrotliResult::NeedsMoreInput => assert_eq!(avail_in, 0),
331           BrotliResult::NeedsMoreOutput => continue,
332           BrotliResult::ResultSuccess => return Ok((buf.len())),
333           BrotliResult::ResultFailure => return Err(self.error_if_invalid_data.take().unwrap()),
334         }
335         if avail_in == 0 {
336            break
337         }
338       }
339       Ok(buf.len())
340     }
flush(&mut self) -> Result<(), ErrType>341     fn flush(&mut self) -> Result<(), ErrType> {
342        self.output.as_mut().unwrap().flush()
343     }
344 }
345 
346