1 //! Extra streaming decompression functionality.
2 //!
3 //! As of now this is mainly inteded for use to build a higher-level wrapper.
4 use crate::alloc::boxed::Box;
5 use core::{cmp, mem};
6
7 use crate::inflate::core::{decompress, inflate_flags, DecompressorOxide, TINFL_LZ_DICT_SIZE};
8 use crate::inflate::TINFLStatus;
9 use crate::{DataFormat, MZError, MZFlush, MZResult, MZStatus, StreamResult};
10
11 /// Tag that determines reset policy of [InflateState](struct.InflateState.html)
12 pub trait ResetPolicy {
13 /// Performs reset
reset(&self, state: &mut InflateState)14 fn reset(&self, state: &mut InflateState);
15 }
16
17 /// Resets state, without performing expensive ops (e.g. zeroing buffer)
18 ///
19 /// Note that not zeroing buffer can lead to security issues when dealing with untrusted input.
20 pub struct MinReset;
21
22 impl ResetPolicy for MinReset {
reset(&self, state: &mut InflateState)23 fn reset(&self, state: &mut InflateState) {
24 state.decompressor().init();
25 state.dict_ofs = 0;
26 state.dict_avail = 0;
27 state.first_call = true;
28 state.has_flushed = false;
29 state.last_status = TINFLStatus::NeedsMoreInput;
30 }
31 }
32
33 /// Resets state and zero memory, continuing to use the same data format.
34 pub struct ZeroReset;
35
36 impl ResetPolicy for ZeroReset {
37 #[inline]
reset(&self, state: &mut InflateState)38 fn reset(&self, state: &mut InflateState) {
39 MinReset.reset(state);
40 state.dict = [0; TINFL_LZ_DICT_SIZE];
41 }
42 }
43
44 /// Full reset of the state, including zeroing memory.
45 ///
46 /// Requires to provide new data format.
47 pub struct FullReset(pub DataFormat);
48
49 impl ResetPolicy for FullReset {
50 #[inline]
reset(&self, state: &mut InflateState)51 fn reset(&self, state: &mut InflateState) {
52 ZeroReset.reset(state);
53 state.data_format = self.0;
54 }
55 }
56
57 /// A struct that compbines a decompressor with extra data for streaming decompression.
58 ///
59 pub struct InflateState {
60 /// Inner decompressor struct
61 decomp: DecompressorOxide,
62
63 /// Buffer of input bytes for matches.
64 /// TODO: Could probably do this a bit cleaner with some
65 /// Cursor-like class.
66 /// We may also look into whether we need to keep a buffer here, or just one in the
67 /// decompressor struct.
68 dict: [u8; TINFL_LZ_DICT_SIZE],
69 /// Where in the buffer are we currently at?
70 dict_ofs: usize,
71 /// How many bytes of data to be flushed is there currently in the buffer?
72 dict_avail: usize,
73
74 first_call: bool,
75 has_flushed: bool,
76
77 /// Whether the input data is wrapped in a zlib header and checksum.
78 /// TODO: This should be stored in the decompressor.
79 data_format: DataFormat,
80 last_status: TINFLStatus,
81 }
82
83 impl Default for InflateState {
default() -> Self84 fn default() -> Self {
85 InflateState {
86 decomp: DecompressorOxide::default(),
87 dict: [0; TINFL_LZ_DICT_SIZE],
88 dict_ofs: 0,
89 dict_avail: 0,
90 first_call: true,
91 has_flushed: false,
92 data_format: DataFormat::Raw,
93 last_status: TINFLStatus::NeedsMoreInput,
94 }
95 }
96 }
97 impl InflateState {
98 /// Create a new state.
99 ///
100 /// Note that this struct is quite large due to internal buffers, and as such storing it on
101 /// the stack is not recommended.
102 ///
103 /// # Parameters
104 /// `data_format`: Determines whether the compressed data is assumed to wrapped with zlib
105 /// metadata.
new(data_format: DataFormat) -> InflateState106 pub fn new(data_format: DataFormat) -> InflateState {
107 let mut b = InflateState::default();
108 b.data_format = data_format;
109 b
110 }
111
112 /// Create a new state on the heap.
113 ///
114 /// # Parameters
115 /// `data_format`: Determines whether the compressed data is assumed to wrapped with zlib
116 /// metadata.
new_boxed(data_format: DataFormat) -> Box<InflateState>117 pub fn new_boxed(data_format: DataFormat) -> Box<InflateState> {
118 let mut b: Box<InflateState> = Box::default();
119 b.data_format = data_format;
120 b
121 }
122
123 /// Access the innner decompressor.
decompressor(&mut self) -> &mut DecompressorOxide124 pub fn decompressor(&mut self) -> &mut DecompressorOxide {
125 &mut self.decomp
126 }
127
128 /// Return the status of the last call to `inflate` with this `InflateState`.
last_status(&self) -> TINFLStatus129 pub const fn last_status(&self) -> TINFLStatus {
130 self.last_status
131 }
132
133 /// Create a new state using miniz/zlib style window bits parameter.
134 ///
135 /// The decompressor does not support different window sizes. As such,
136 /// any positive (>0) value will set the zlib header flag, while a negative one
137 /// will not.
new_boxed_with_window_bits(window_bits: i32) -> Box<InflateState>138 pub fn new_boxed_with_window_bits(window_bits: i32) -> Box<InflateState> {
139 let mut b: Box<InflateState> = Box::default();
140 b.data_format = DataFormat::from_window_bits(window_bits);
141 b
142 }
143
144 #[inline]
145 /// Reset the decompressor without re-allocating memory, using the given
146 /// data format.
reset(&mut self, data_format: DataFormat)147 pub fn reset(&mut self, data_format: DataFormat) {
148 self.reset_as(FullReset(data_format));
149 }
150
151 #[inline]
152 /// Resets the state according to specified policy.
reset_as<T: ResetPolicy>(&mut self, policy: T)153 pub fn reset_as<T: ResetPolicy>(&mut self, policy: T) {
154 policy.reset(self)
155 }
156 }
157
158 /// Try to decompress from `input` to `output` with the given `InflateState`
159 ///
160 /// # Errors
161 ///
162 /// Returns `MZError::Buf` If the size of the `output` slice is empty or no progress was made due to
163 /// lack of expected input data or called after the decompression was
164 /// finished without MZFlush::Finish.
165 ///
166 /// Returns `MZError::Param` if the compressor parameters are set wrong.
inflate( state: &mut InflateState, input: &[u8], output: &mut [u8], flush: MZFlush, ) -> StreamResult167 pub fn inflate(
168 state: &mut InflateState,
169 input: &[u8],
170 output: &mut [u8],
171 flush: MZFlush,
172 ) -> StreamResult {
173 let mut bytes_consumed = 0;
174 let mut bytes_written = 0;
175 let mut next_in = input;
176 let mut next_out = output;
177
178 if flush == MZFlush::Full {
179 return StreamResult::error(MZError::Stream);
180 }
181
182 let mut decomp_flags = inflate_flags::TINFL_FLAG_COMPUTE_ADLER32;
183 if state.data_format == DataFormat::Zlib {
184 decomp_flags |= inflate_flags::TINFL_FLAG_PARSE_ZLIB_HEADER;
185 }
186
187 let first_call = state.first_call;
188 state.first_call = false;
189 if (state.last_status as i32) < 0 {
190 return StreamResult::error(MZError::Data);
191 }
192
193 if state.has_flushed && (flush != MZFlush::Finish) {
194 return StreamResult::error(MZError::Stream);
195 }
196 state.has_flushed |= flush == MZFlush::Finish;
197
198 if (flush == MZFlush::Finish) && first_call {
199 decomp_flags |= inflate_flags::TINFL_FLAG_USING_NON_WRAPPING_OUTPUT_BUF;
200
201 let status = decompress(&mut state.decomp, next_in, next_out, 0, decomp_flags);
202 let in_bytes = status.1;
203 let out_bytes = status.2;
204 let status = status.0;
205
206 state.last_status = status;
207
208 bytes_consumed += in_bytes;
209 bytes_written += out_bytes;
210
211 let ret_status = {
212 if (status as i32) < 0 {
213 Err(MZError::Data)
214 } else if status != TINFLStatus::Done {
215 state.last_status = TINFLStatus::Failed;
216 Err(MZError::Buf)
217 } else {
218 Ok(MZStatus::StreamEnd)
219 }
220 };
221 return StreamResult {
222 bytes_consumed,
223 bytes_written,
224 status: ret_status,
225 };
226 }
227
228 if flush != MZFlush::Finish {
229 decomp_flags |= inflate_flags::TINFL_FLAG_HAS_MORE_INPUT;
230 }
231
232 if state.dict_avail != 0 {
233 bytes_written += push_dict_out(state, &mut next_out);
234 return StreamResult {
235 bytes_consumed,
236 bytes_written,
237 status: Ok(
238 if (state.last_status == TINFLStatus::Done) && (state.dict_avail == 0) {
239 MZStatus::StreamEnd
240 } else {
241 MZStatus::Ok
242 },
243 ),
244 };
245 }
246
247 let status = inflate_loop(
248 state,
249 &mut next_in,
250 &mut next_out,
251 &mut bytes_consumed,
252 &mut bytes_written,
253 decomp_flags,
254 flush,
255 );
256 StreamResult {
257 bytes_consumed,
258 bytes_written,
259 status,
260 }
261 }
262
inflate_loop( state: &mut InflateState, next_in: &mut &[u8], next_out: &mut &mut [u8], total_in: &mut usize, total_out: &mut usize, decomp_flags: u32, flush: MZFlush, ) -> MZResult263 fn inflate_loop(
264 state: &mut InflateState,
265 next_in: &mut &[u8],
266 next_out: &mut &mut [u8],
267 total_in: &mut usize,
268 total_out: &mut usize,
269 decomp_flags: u32,
270 flush: MZFlush,
271 ) -> MZResult {
272 let orig_in_len = next_in.len();
273 loop {
274 let status = decompress(
275 &mut state.decomp,
276 *next_in,
277 &mut state.dict,
278 state.dict_ofs,
279 decomp_flags,
280 );
281
282 let in_bytes = status.1;
283 let out_bytes = status.2;
284 let status = status.0;
285
286 state.last_status = status;
287
288 *next_in = &next_in[in_bytes..];
289 *total_in += in_bytes;
290
291 state.dict_avail = out_bytes;
292 *total_out += push_dict_out(state, next_out);
293
294 // The stream was corrupted, and decompression failed.
295 if (status as i32) < 0 {
296 return Err(MZError::Data);
297 }
298
299 // The decompressor has flushed all it's data and is waiting for more input, but
300 // there was no more input provided.
301 if (status == TINFLStatus::NeedsMoreInput) && orig_in_len == 0 {
302 return Err(MZError::Buf);
303 }
304
305 if flush == MZFlush::Finish {
306 if status == TINFLStatus::Done {
307 // There is not enough space in the output buffer to flush the remaining
308 // decompressed data in the internal buffer.
309 return if state.dict_avail != 0 {
310 Err(MZError::Buf)
311 } else {
312 Ok(MZStatus::StreamEnd)
313 };
314 // No more space in the output buffer, but we're not done.
315 } else if next_out.is_empty() {
316 return Err(MZError::Buf);
317 }
318 } else {
319 // We're not expected to finish, so it's fine if we can't flush everything yet.
320 let empty_buf = next_in.is_empty() || next_out.is_empty();
321 if (status == TINFLStatus::Done) || empty_buf || (state.dict_avail != 0) {
322 return if (status == TINFLStatus::Done) && (state.dict_avail == 0) {
323 // No more data left, we're done.
324 Ok(MZStatus::StreamEnd)
325 } else {
326 // Ok for now, still waiting for more input data or output space.
327 Ok(MZStatus::Ok)
328 };
329 }
330 }
331 }
332 }
333
push_dict_out(state: &mut InflateState, next_out: &mut &mut [u8]) -> usize334 fn push_dict_out(state: &mut InflateState, next_out: &mut &mut [u8]) -> usize {
335 let n = cmp::min(state.dict_avail as usize, next_out.len());
336 (next_out[..n]).copy_from_slice(&state.dict[state.dict_ofs..state.dict_ofs + n]);
337 *next_out = &mut mem::replace(next_out, &mut [])[n..];
338 state.dict_avail -= n;
339 state.dict_ofs = (state.dict_ofs + (n)) & (TINFL_LZ_DICT_SIZE - 1);
340 n
341 }
342
343 #[cfg(test)]
344 mod test {
345 use super::{inflate, InflateState};
346 use crate::{DataFormat, MZFlush, MZStatus};
347 use std::vec;
348
349 #[test]
test_state()350 fn test_state() {
351 let encoded = [
352 120u8, 156, 243, 72, 205, 201, 201, 215, 81, 168, 202, 201, 76, 82, 4, 0, 27, 101, 4,
353 19,
354 ];
355 let mut out = vec![0; 50];
356 let mut state = InflateState::new_boxed(DataFormat::Zlib);
357 let res = inflate(&mut state, &encoded, &mut out, MZFlush::Finish);
358 let status = res.status.expect("Failed to decompress!");
359 assert_eq!(status, MZStatus::StreamEnd);
360 assert_eq!(out[..res.bytes_written as usize], b"Hello, zlib!"[..]);
361 assert_eq!(res.bytes_consumed, encoded.len());
362
363 state.reset_as(super::ZeroReset);
364 out.iter_mut().map(|x| *x = 0).count();
365 let res = inflate(&mut state, &encoded, &mut out, MZFlush::Finish);
366 let status = res.status.expect("Failed to decompress!");
367 assert_eq!(status, MZStatus::StreamEnd);
368 assert_eq!(out[..res.bytes_written as usize], b"Hello, zlib!"[..]);
369 assert_eq!(res.bytes_consumed, encoded.len());
370
371 state.reset_as(super::MinReset);
372 out.iter_mut().map(|x| *x = 0).count();
373 let res = inflate(&mut state, &encoded, &mut out, MZFlush::Finish);
374 let status = res.status.expect("Failed to decompress!");
375 assert_eq!(status, MZStatus::StreamEnd);
376 assert_eq!(out[..res.bytes_written as usize], b"Hello, zlib!"[..]);
377 assert_eq!(res.bytes_consumed, encoded.len());
378 }
379 }
380