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