miniz_oxide/inflate/
stream.rs

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 = cmp::min(state.dict_avail, 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(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}