miniz_oxide/inflate/
stream.rs

Help
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///
60#[derive(Clone)]
61pub struct InflateState {
62    /// Inner decompressor struct
63    decomp: DecompressorOxide,
64
65    /// Buffer of input bytes for matches.
66    /// TODO: Could probably do this a bit cleaner with some
67    /// Cursor-like class.
68    /// We may also look into whether we need to keep a buffer here, or just one in the
69    /// decompressor struct.
70    dict: [u8; TINFL_LZ_DICT_SIZE],
71    /// Where in the buffer are we currently at?
72    dict_ofs: usize,
73    /// How many bytes of data to be flushed is there currently in the buffer?
74    dict_avail: usize,
75
76    first_call: bool,
77    has_flushed: bool,
78
79    /// Whether the input data is wrapped in a zlib header and checksum.
80    /// TODO: This should be stored in the decompressor.
81    data_format: DataFormat,
82    last_status: TINFLStatus,
83}
84
85impl Default for InflateState {
86    fn default() -> Self {
87        InflateState {
88            decomp: DecompressorOxide::default(),
89            dict: [0; TINFL_LZ_DICT_SIZE],
90            dict_ofs: 0,
91            dict_avail: 0,
92            first_call: true,
93            has_flushed: false,
94            data_format: DataFormat::Raw,
95            last_status: TINFLStatus::NeedsMoreInput,
96        }
97    }
98}
99impl InflateState {
100    /// Create a new state.
101    ///
102    /// Note that this struct is quite large due to internal buffers, and as such storing it on
103    /// the stack is not recommended.
104    ///
105    /// # Parameters
106    /// `data_format`: Determines whether the compressed data is assumed to wrapped with zlib
107    /// metadata.
108    pub fn new(data_format: DataFormat) -> InflateState {
109        InflateState {
110            data_format,
111            ..Default::default()
112        }
113    }
114
115    /// Create a new state on the heap.
116    ///
117    /// # Parameters
118    /// `data_format`: Determines whether the compressed data is assumed to wrapped with zlib
119    /// metadata.
120    #[cfg(feature = "with-alloc")]
121    pub fn new_boxed(data_format: DataFormat) -> Box<InflateState> {
122        let mut b: Box<InflateState> = Box::default();
123        b.data_format = data_format;
124        b
125    }
126
127    /// Access the innner decompressor.
128    pub fn decompressor(&mut self) -> &mut DecompressorOxide {
129        &mut self.decomp
130    }
131
132    /// Return the status of the last call to `inflate` with this `InflateState`.
133    pub const fn last_status(&self) -> TINFLStatus {
134        self.last_status
135    }
136
137    /// Create a new state using miniz/zlib style window bits parameter.
138    ///
139    /// The decompressor does not support different window sizes. As such,
140    /// any positive (>0) value will set the zlib header flag, while a negative one
141    /// will not.
142    #[cfg(feature = "with-alloc")]
143    pub fn new_boxed_with_window_bits(window_bits: i32) -> Box<InflateState> {
144        let mut b: Box<InflateState> = Box::default();
145        b.data_format = DataFormat::from_window_bits(window_bits);
146        b
147    }
148
149    #[inline]
150    /// Reset the decompressor without re-allocating memory, using the given
151    /// data format.
152    pub fn reset(&mut self, data_format: DataFormat) {
153        self.reset_as(FullReset(data_format));
154    }
155
156    #[inline]
157    /// Resets the state according to specified policy.
158    pub fn reset_as<T: ResetPolicy>(&mut self, policy: T) {
159        policy.reset(self)
160    }
161}
162
163/// Try to decompress from `input` to `output` with the given [`InflateState`]
164///
165/// # `flush`
166///
167/// Generally, the various [`MZFlush`] flags have meaning only on the compression side.  They can be
168/// supplied here, but the only one that has any semantic meaning is [`MZFlush::Finish`], which is a
169/// signal that the stream is expected to finish, and failing to do so is an error.  It isn't
170/// necessary to specify it when the stream ends; you'll still get returned a
171/// [`MZStatus::StreamEnd`] anyway.  Other values either have no effect or cause errors.  It's
172/// likely that you'll almost always just want to use [`MZFlush::None`].
173///
174/// # Errors
175///
176/// Returns [`MZError::Buf`] if the size of the `output` slice is empty or no progress was made due
177/// to lack of expected input data, or if called with [`MZFlush::Finish`] and input wasn't all
178/// consumed.
179///
180/// Returns [`MZError::Data`] if this or a a previous call failed with an error return from
181/// [`TINFLStatus`]; probably indicates corrupted data.
182///
183/// Returns [`MZError::Stream`] when called with [`MZFlush::Full`] (meaningless on
184/// decompression), or when called without [`MZFlush::Finish`] after an earlier call with
185/// [`MZFlush::Finish`] has been made.
186pub fn inflate(
187    state: &mut InflateState,
188    input: &[u8],
189    output: &mut [u8],
190    flush: MZFlush,
191) -> StreamResult {
192    let mut bytes_consumed = 0;
193    let mut bytes_written = 0;
194    let mut next_in = input;
195    let mut next_out = output;
196
197    if flush == MZFlush::Full {
198        return StreamResult::error(MZError::Stream);
199    }
200
201    let mut decomp_flags = if state.data_format == DataFormat::Zlib {
202        inflate_flags::TINFL_FLAG_COMPUTE_ADLER32
203    } else {
204        inflate_flags::TINFL_FLAG_IGNORE_ADLER32
205    };
206
207    if (state.data_format == DataFormat::Zlib)
208        | (state.data_format == DataFormat::ZLibIgnoreChecksum)
209    {
210        decomp_flags |= inflate_flags::TINFL_FLAG_PARSE_ZLIB_HEADER;
211    }
212
213    let first_call = state.first_call;
214    state.first_call = false;
215    if state.last_status == TINFLStatus::FailedCannotMakeProgress {
216        return StreamResult::error(MZError::Buf);
217    }
218    if (state.last_status as i32) < 0 {
219        return StreamResult::error(MZError::Data);
220    }
221
222    if state.has_flushed && (flush != MZFlush::Finish) {
223        return StreamResult::error(MZError::Stream);
224    }
225    state.has_flushed |= flush == MZFlush::Finish;
226
227    if (flush == MZFlush::Finish) && first_call {
228        decomp_flags |= inflate_flags::TINFL_FLAG_USING_NON_WRAPPING_OUTPUT_BUF;
229
230        // The caller is indicating that they want to finish the compression and this is the first call with the current stream
231        // so we can simply write directly to the output buffer.
232        // If there is not enough space for all of the decompressed data we will end up with a failure regardless.
233        let status = decompress(&mut state.decomp, next_in, next_out, 0, decomp_flags);
234        let in_bytes = status.1;
235        let out_bytes = status.2;
236        let status = status.0;
237
238        state.last_status = status;
239
240        bytes_consumed += in_bytes;
241        bytes_written += out_bytes;
242
243        let ret_status = {
244            if status == TINFLStatus::FailedCannotMakeProgress {
245                Err(MZError::Buf)
246            } else if (status as i32) < 0 {
247                Err(MZError::Data)
248            } else if status != TINFLStatus::Done {
249                state.last_status = TINFLStatus::Failed;
250                Err(MZError::Buf)
251            } else {
252                Ok(MZStatus::StreamEnd)
253            }
254        };
255        return StreamResult {
256            bytes_consumed,
257            bytes_written,
258            status: ret_status,
259        };
260    }
261
262    if flush != MZFlush::Finish {
263        decomp_flags |= inflate_flags::TINFL_FLAG_HAS_MORE_INPUT;
264    }
265
266    if state.dict_avail != 0 {
267        bytes_written += push_dict_out(state, &mut next_out);
268        return StreamResult {
269            bytes_consumed,
270            bytes_written,
271            status: Ok(
272                if (state.last_status == TINFLStatus::Done) && (state.dict_avail == 0) {
273                    MZStatus::StreamEnd
274                } else {
275                    MZStatus::Ok
276                },
277            ),
278        };
279    }
280
281    let status = inflate_loop(
282        state,
283        &mut next_in,
284        &mut next_out,
285        &mut bytes_consumed,
286        &mut bytes_written,
287        decomp_flags,
288        flush,
289    );
290    StreamResult {
291        bytes_consumed,
292        bytes_written,
293        status,
294    }
295}
296
297fn inflate_loop(
298    state: &mut InflateState,
299    next_in: &mut &[u8],
300    next_out: &mut &mut [u8],
301    total_in: &mut usize,
302    total_out: &mut usize,
303    decomp_flags: u32,
304    flush: MZFlush,
305) -> MZResult {
306    let orig_in_len = next_in.len();
307    loop {
308        let status = decompress(
309            &mut state.decomp,
310            next_in,
311            &mut state.dict,
312            state.dict_ofs,
313            decomp_flags,
314        );
315
316        let in_bytes = status.1;
317        let out_bytes = status.2;
318        let status = status.0;
319
320        state.last_status = status;
321
322        *next_in = &next_in[in_bytes..];
323        *total_in += in_bytes;
324
325        state.dict_avail = out_bytes;
326        *total_out += push_dict_out(state, next_out);
327
328        // The stream was corrupted, and decompression failed.
329        if (status as i32) < 0 {
330            return Err(MZError::Data);
331        }
332
333        // The decompressor has flushed all it's data and is waiting for more input, but
334        // there was no more input provided.
335        if (status == TINFLStatus::NeedsMoreInput) && orig_in_len == 0 {
336            return Err(MZError::Buf);
337        }
338
339        if flush == MZFlush::Finish {
340            if status == TINFLStatus::Done {
341                // There is not enough space in the output buffer to flush the remaining
342                // decompressed data in the internal buffer.
343                return if state.dict_avail != 0 {
344                    Err(MZError::Buf)
345                } else {
346                    Ok(MZStatus::StreamEnd)
347                };
348            // No more space in the output buffer, but we're not done.
349            } else if next_out.is_empty() {
350                return Err(MZError::Buf);
351            }
352        } else {
353            // We're not expected to finish, so it's fine if we can't flush everything yet.
354            let empty_buf = next_in.is_empty() || next_out.is_empty();
355            if (status == TINFLStatus::Done) || empty_buf || (state.dict_avail != 0) {
356                return if (status == TINFLStatus::Done) && (state.dict_avail == 0) {
357                    // No more data left, we're done.
358                    Ok(MZStatus::StreamEnd)
359                } else {
360                    // Ok for now, still waiting for more input data or output space.
361                    Ok(MZStatus::Ok)
362                };
363            }
364        }
365    }
366}
367
368fn push_dict_out(state: &mut InflateState, next_out: &mut &mut [u8]) -> usize {
369    let n = cmp::min(state.dict_avail, next_out.len());
370    (next_out[..n]).copy_from_slice(&state.dict[state.dict_ofs..state.dict_ofs + n]);
371    *next_out = &mut mem::take(next_out)[n..];
372    state.dict_avail -= n;
373    state.dict_ofs = (state.dict_ofs + (n)) & (TINFL_LZ_DICT_SIZE - 1);
374    n
375}
376
377#[cfg(all(test, feature = "with-alloc"))]
378mod test {
379    use super::{inflate, InflateState};
380    use crate::{DataFormat, MZFlush, MZStatus};
381    use alloc::vec;
382
383    #[test]
384    fn test_state() {
385        let encoded = [
386            120u8, 156, 243, 72, 205, 201, 201, 215, 81, 168, 202, 201, 76, 82, 4, 0, 27, 101, 4,
387            19,
388        ];
389        let mut out = vec![0; 50];
390        let mut state = InflateState::new_boxed(DataFormat::Zlib);
391        let res = inflate(&mut state, &encoded, &mut out, MZFlush::Finish);
392        let status = res.status.expect("Failed to decompress!");
393        assert_eq!(status, MZStatus::StreamEnd);
394        assert_eq!(out[..res.bytes_written as usize], b"Hello, zlib!"[..]);
395        assert_eq!(res.bytes_consumed, encoded.len());
396
397        state.reset_as(super::ZeroReset);
398        out.iter_mut().map(|x| *x = 0).count();
399        let res = inflate(&mut state, &encoded, &mut out, MZFlush::Finish);
400        let status = res.status.expect("Failed to decompress!");
401        assert_eq!(status, MZStatus::StreamEnd);
402        assert_eq!(out[..res.bytes_written as usize], b"Hello, zlib!"[..]);
403        assert_eq!(res.bytes_consumed, encoded.len());
404
405        state.reset_as(super::MinReset);
406        out.iter_mut().map(|x| *x = 0).count();
407        let res = inflate(&mut state, &encoded, &mut out, MZFlush::Finish);
408        let status = res.status.expect("Failed to decompress!");
409        assert_eq!(status, MZStatus::StreamEnd);
410        assert_eq!(out[..res.bytes_written as usize], b"Hello, zlib!"[..]);
411        assert_eq!(res.bytes_consumed, encoded.len());
412        assert_eq!(state.decompressor().adler32(), Some(459605011));
413
414        // Test state when not computing adler.
415        state = InflateState::new_boxed(DataFormat::ZLibIgnoreChecksum);
416        out.iter_mut().map(|x| *x = 0).count();
417        let res = inflate(&mut state, &encoded, &mut out, MZFlush::Finish);
418        let status = res.status.expect("Failed to decompress!");
419        assert_eq!(status, MZStatus::StreamEnd);
420        assert_eq!(out[..res.bytes_written as usize], b"Hello, zlib!"[..]);
421        assert_eq!(res.bytes_consumed, encoded.len());
422        // Not computed, so should be Some(1)
423        assert_eq!(state.decompressor().adler32(), Some(1));
424        // Should still have the checksum read from the header file.
425        assert_eq!(state.decompressor().adler32_header(), Some(459605011))
426    }
427
428    #[test]
429    fn test_partial_continue() {
430        let encoded = [
431            120u8, 156, 243, 72, 205, 201, 201, 215, 81, 168, 202, 201, 76, 82, 4, 0, 27, 101, 4,
432            19,
433        ];
434
435        // Feed input bytes one at a time to the decompressor
436        let mut out = vec![0; 50];
437        let mut state = InflateState::new_boxed(DataFormat::Zlib);
438        let mut part_in = 0;
439        let mut part_out = 0;
440        for i in 1..=encoded.len() {
441            let res = inflate(
442                &mut state,
443                &encoded[part_in..i],
444                &mut out[part_out..],
445                MZFlush::None,
446            );
447            let status = res.status.expect("Failed to decompress!");
448            if i == encoded.len() {
449                assert_eq!(status, MZStatus::StreamEnd);
450            } else {
451                assert_eq!(status, MZStatus::Ok);
452            }
453            part_out += res.bytes_written as usize;
454            part_in += res.bytes_consumed;
455        }
456
457        assert_eq!(out[..part_out as usize], b"Hello, zlib!"[..]);
458        assert_eq!(part_in, encoded.len());
459        assert_eq!(state.decompressor().adler32(), Some(459605011));
460    }
461
462    // Inflate part of a stream and clone the inflate state.
463    // Discard the original state and resume the stream from the clone.
464    #[test]
465    fn test_rewind_and_resume() {
466        let encoded = [
467            120u8, 156, 243, 72, 205, 201, 201, 215, 81, 168, 202, 201, 76, 82, 4, 0, 27, 101, 4,
468            19,
469        ];
470        let decoded = b"Hello, zlib!";
471
472        // Feed partial input bytes to the decompressor
473        let mut out = vec![0; 50];
474        let mut state = InflateState::new_boxed(DataFormat::Zlib);
475        let res1 = inflate(&mut state, &encoded[..10], &mut out, MZFlush::None);
476        let status = res1.status.expect("Failed to decompress!");
477        assert_eq!(status, MZStatus::Ok);
478
479        // Clone the state and discard the original
480        let mut resume = state.clone();
481        drop(state);
482
483        // Resume the stream using the cloned state
484        let res2 = inflate(
485            &mut resume,
486            &encoded[res1.bytes_consumed..],
487            &mut out[res1.bytes_written..],
488            MZFlush::Finish,
489        );
490        let status = res2.status.expect("Failed to decompress!");
491        assert_eq!(status, MZStatus::StreamEnd);
492
493        assert_eq!(res1.bytes_consumed + res2.bytes_consumed, encoded.len());
494        assert_eq!(res1.bytes_written + res2.bytes_written, decoded.len());
495        assert_eq!(
496            &out[..res1.bytes_written + res2.bytes_written as usize],
497            decoded
498        );
499        assert_eq!(resume.decompressor().adler32(), Some(459605011));
500    }
501}
miniz_oxide/inflate/stream.rs

miniz_oxide/inflate/
stream.rs
