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")]
5use crate::alloc::boxed::Box;
6use core::{cmp, mem};
7
8use crate::inflate::core::{decompress, inflate_flags, DecompressorOxide, TINFL_LZ_DICT_SIZE};
9use crate::inflate::TINFLStatus;
10use crate::{DataFormat, MZError, MZFlush, MZResult, MZStatus, StreamResult};
11
12/// Tag that determines reset policy of [InflateState](struct.InflateState.html)
13pub 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.
21pub struct MinReset;
22
23impl 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.
35pub struct ZeroReset;
36
37impl 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.
48pub struct FullReset(pub DataFormat);
49
50impl 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///
60pub 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
84impl 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}
98impl 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.
185pub 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 == TINFLStatus::FailedCannotMakeProgress {
215 return StreamResult::error(MZError::Buf);
216 }
217 if (state.last_status as i32) < 0 {
218 return StreamResult::error(MZError::Data);
219 }
220
221 if state.has_flushed && (flush != MZFlush::Finish) {
222 return StreamResult::error(MZError::Stream);
223 }
224 state.has_flushed |= flush == MZFlush::Finish;
225
226 if (flush == MZFlush::Finish) && first_call {
227 decomp_flags |= inflate_flags::TINFL_FLAG_USING_NON_WRAPPING_OUTPUT_BUF;
228
229 let status = decompress(&mut state.decomp, next_in, next_out, 0, decomp_flags);
230 let in_bytes = status.1;
231 let out_bytes = status.2;
232 let status = status.0;
233
234 state.last_status = status;
235
236 bytes_consumed += in_bytes;
237 bytes_written += out_bytes;
238
239 let ret_status = {
240 if status == TINFLStatus::FailedCannotMakeProgress {
241 Err(MZError::Buf)
242 } else if (status as i32) < 0 {
243 Err(MZError::Data)
244 } else if status != TINFLStatus::Done {
245 state.last_status = TINFLStatus::Failed;
246 Err(MZError::Buf)
247 } else {
248 Ok(MZStatus::StreamEnd)
249 }
250 };
251 return StreamResult {
252 bytes_consumed,
253 bytes_written,
254 status: ret_status,
255 };
256 }
257
258 if flush != MZFlush::Finish {
259 decomp_flags |= inflate_flags::TINFL_FLAG_HAS_MORE_INPUT;
260 }
261
262 if state.dict_avail != 0 {
263 bytes_written += push_dict_out(state, &mut next_out);
264 return StreamResult {
265 bytes_consumed,
266 bytes_written,
267 status: Ok(
268 if (state.last_status == TINFLStatus::Done) && (state.dict_avail == 0) {
269 MZStatus::StreamEnd
270 } else {
271 MZStatus::Ok
272 },
273 ),
274 };
275 }
276
277 let status = inflate_loop(
278 state,
279 &mut next_in,
280 &mut next_out,
281 &mut bytes_consumed,
282 &mut bytes_written,
283 decomp_flags,
284 flush,
285 );
286 StreamResult {
287 bytes_consumed,
288 bytes_written,
289 status,
290 }
291}
292
293fn inflate_loop(
294 state: &mut InflateState,
295 next_in: &mut &[u8],
296 next_out: &mut &mut [u8],
297 total_in: &mut usize,
298 total_out: &mut usize,
299 decomp_flags: u32,
300 flush: MZFlush,
301) -> MZResult {
302 let orig_in_len = next_in.len();
303 loop {
304 let status = decompress(
305 &mut state.decomp,
306 next_in,
307 &mut state.dict,
308 state.dict_ofs,
309 decomp_flags,
310 );
311
312 let in_bytes = status.1;
313 let out_bytes = status.2;
314 let status = status.0;
315
316 state.last_status = status;
317
318 *next_in = &next_in[in_bytes..];
319 *total_in += in_bytes;
320
321 state.dict_avail = out_bytes;
322 *total_out += push_dict_out(state, next_out);
323
324 // The stream was corrupted, and decompression failed.
325 if (status as i32) < 0 {
326 return Err(MZError::Data);
327 }
328
329 // The decompressor has flushed all it's data and is waiting for more input, but
330 // there was no more input provided.
331 if (status == TINFLStatus::NeedsMoreInput) && orig_in_len == 0 {
332 return Err(MZError::Buf);
333 }
334
335 if flush == MZFlush::Finish {
336 if status == TINFLStatus::Done {
337 // There is not enough space in the output buffer to flush the remaining
338 // decompressed data in the internal buffer.
339 return if state.dict_avail != 0 {
340 Err(MZError::Buf)
341 } else {
342 Ok(MZStatus::StreamEnd)
343 };
344 // No more space in the output buffer, but we're not done.
345 } else if next_out.is_empty() {
346 return Err(MZError::Buf);
347 }
348 } else {
349 // We're not expected to finish, so it's fine if we can't flush everything yet.
350 let empty_buf = next_in.is_empty() || next_out.is_empty();
351 if (status == TINFLStatus::Done) || empty_buf || (state.dict_avail != 0) {
352 return if (status == TINFLStatus::Done) && (state.dict_avail == 0) {
353 // No more data left, we're done.
354 Ok(MZStatus::StreamEnd)
355 } else {
356 // Ok for now, still waiting for more input data or output space.
357 Ok(MZStatus::Ok)
358 };
359 }
360 }
361 }
362}
363
364fn push_dict_out(state: &mut InflateState, next_out: &mut &mut [u8]) -> usize {
365 let n: usize = cmp::min(v1:state.dict_avail, v2:next_out.len());
366 (next_out[..n]).copy_from_slice(&state.dict[state.dict_ofs..state.dict_ofs + n]);
367 *next_out = &mut mem::take(dest:next_out)[n..];
368 state.dict_avail -= n;
369 state.dict_ofs = (state.dict_ofs + (n)) & (TINFL_LZ_DICT_SIZE - 1);
370 n
371}
372
373#[cfg(all(test, feature = "with-alloc"))]
374mod test {
375 use super::{inflate, InflateState};
376 use crate::{DataFormat, MZFlush, MZStatus};
377 use alloc::vec;
378
379 #[test]
380 fn test_state() {
381 let encoded = [
382 120u8, 156, 243, 72, 205, 201, 201, 215, 81, 168, 202, 201, 76, 82, 4, 0, 27, 101, 4,
383 19,
384 ];
385 let mut out = vec![0; 50];
386 let mut state = InflateState::new_boxed(DataFormat::Zlib);
387 let res = inflate(&mut state, &encoded, &mut out, MZFlush::Finish);
388 let status = res.status.expect("Failed to decompress!");
389 assert_eq!(status, MZStatus::StreamEnd);
390 assert_eq!(out[..res.bytes_written as usize], b"Hello, zlib!"[..]);
391 assert_eq!(res.bytes_consumed, encoded.len());
392
393 state.reset_as(super::ZeroReset);
394 out.iter_mut().map(|x| *x = 0).count();
395 let res = inflate(&mut state, &encoded, &mut out, MZFlush::Finish);
396 let status = res.status.expect("Failed to decompress!");
397 assert_eq!(status, MZStatus::StreamEnd);
398 assert_eq!(out[..res.bytes_written as usize], b"Hello, zlib!"[..]);
399 assert_eq!(res.bytes_consumed, encoded.len());
400
401 state.reset_as(super::MinReset);
402 out.iter_mut().map(|x| *x = 0).count();
403 let res = inflate(&mut state, &encoded, &mut out, MZFlush::Finish);
404 let status = res.status.expect("Failed to decompress!");
405 assert_eq!(status, MZStatus::StreamEnd);
406 assert_eq!(out[..res.bytes_written as usize], b"Hello, zlib!"[..]);
407 assert_eq!(res.bytes_consumed, encoded.len());
408 assert_eq!(state.decompressor().adler32(), Some(459605011));
409
410 // Test state when not computing adler.
411 state = InflateState::new_boxed(DataFormat::ZLibIgnoreChecksum);
412 out.iter_mut().map(|x| *x = 0).count();
413 let res = inflate(&mut state, &encoded, &mut out, MZFlush::Finish);
414 let status = res.status.expect("Failed to decompress!");
415 assert_eq!(status, MZStatus::StreamEnd);
416 assert_eq!(out[..res.bytes_written as usize], b"Hello, zlib!"[..]);
417 assert_eq!(res.bytes_consumed, encoded.len());
418 // Not computed, so should be Some(1)
419 assert_eq!(state.decompressor().adler32(), Some(1));
420 // Should still have the checksum read from the header file.
421 assert_eq!(state.decompressor().adler32_header(), Some(459605011))
422 }
423}
424