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 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
288fn 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
359fn 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)]
369mod 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