vendor/miniz_oxide/src/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")]
   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 = cmp::min(state.dict_avail as usize, 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(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 }