inflate: add opt-in limit on Huffman table rebuilds (fixes #198)

miniz_oxide/src/inflate/core.rs

@@ -283,6 +283,20 @@ pub struct DecompressorOxide {
     counter: u32,
     /// Number of extra bits for the last length or distance code.
     num_extra: u8,
+    /// Number of times the Huffman decode tables have been (re)built so far while decoding the
+    /// current stream. Reset to 0 whenever decoding (re)starts from [`State::Start`].
+    ///
+    /// See [`DecompressorOxide::set_max_huffman_table_rebuilds`] for why this is tracked.
+    huffman_table_rebuilds: u32,
+    /// Maximum value [`Self::huffman_table_rebuilds`] is allowed to reach before decompression
+    /// is aborted with [`TINFLStatus::HuffmanTableRebuildLimitExceeded`][crate::inflate::TINFLStatus::HuffmanTableRebuildLimitExceeded].
+    ///
+    /// Unlike most other fields on this struct, this is caller-configured state: it defaults to
+    /// `u32::MAX` (no limit, preserving prior behavior) and is *not* reset by [`Self::init`] or
+    /// by the `Start` state, so it only needs to be set once per `DecompressorOxide` even if the
+    /// same instance is reused (e.g. via [`crate::inflate::stream::InflateState`] resets) to
+    /// decode multiple streams.
+    max_huffman_table_rebuilds: u32,
     /// Number of entries in each huffman table.
     table_sizes: [u16; MAX_HUFF_TABLES],
     /// Buffer of input data.
@@ -316,6 +330,58 @@ impl DecompressorOxide {
         self.state = core::State::Start;
     }
 
+    /// Sets the maximum number of times decompression is allowed to (re)build the internal
+    /// Huffman decode tables while decoding a single deflate stream, and returns an early
+    /// failure ([`TINFLStatus::HuffmanTableRebuildLimitExceeded`]) if that budget is exceeded.
+    ///
+    /// # Background
+    ///
+    /// Every `BTYPE=1` (static Huffman) or `BTYPE=2` (dynamic Huffman) deflate block header
+    /// forces a full rebuild of the decoder's Huffman lookup tables, at a fixed cost that does
+    /// not depend on how many literal/match symbols (including zero) the block actually
+    /// encodes before its end-of-block marker. Since the DEFLATE format allows an arbitrarily
+    /// long chain of minimal, near-empty blocks (each one legally as short as 10 bits for a
+    /// static block containing nothing but an end-of-block code), a crafted input of a few
+    /// hundred KiB to a few MiB can force many millions of these rebuilds while producing
+    /// little or no decompressed output at all. This decouples decompression cost from output
+    /// size, which means the existing output-size limit (see [`decompress_with_limit`]) does
+    /// **not** bound this cost: the limit is only ever compared against bytes written, and such
+    /// an input may write close to nothing.
+    ///
+    /// This is a legal (if highly unusual) deflate stream, not a malformed one, so it cannot be
+    /// rejected outright without risking rejecting some legitimate, if degenerate, streams.
+    /// Setting a limit here lets a caller who cares about this attack vector bound the cost
+    /// explicitly, independently of the output-size limit.
+    ///
+    /// # Default
+    ///
+    /// Defaults to `u32::MAX`, i.e. no limit, which preserves prior behavior for callers that do
+    /// not opt in.
+    ///
+    /// # Picking a value
+    ///
+    /// Each deflate block triggers 2 rebuilds if static, or 3 if dynamic (one each for the
+    /// huffman-length, literal/length, and distance tables), so this is roughly `2-3x` the
+    /// number of blocks a stream is allowed to contain. What value is "reasonable" depends
+    /// entirely on the caller's own workload (in particular, how many legitimate blocks a
+    /// normal input for that workload might contain) traded off against how much worst-case CPU
+    /// time is acceptable; there is no value that is safe for every caller, so none is applied
+    /// by default.
+    ///
+    /// # Persistence across resets
+    ///
+    /// Unlike most of this struct's fields, this value is caller-configured, not decode state:
+    /// it is left untouched by [`Self::init`] and by the `Start` state, so it only needs to be
+    /// set once per instance even if reused (e.g. via [`crate::inflate::stream::InflateState`]
+    /// resets) to decode multiple streams. Note that the (feature-gated)
+    /// `from_block_boundary_state` constructor builds a fresh instance via [`Default`] and does
+    /// *not* carry over a previously configured limit; call this again afterwards if that
+    /// matters for your use case.
+    #[inline]
+    pub fn set_max_huffman_table_rebuilds(&mut self, max_huffman_table_rebuilds: u32) {
+        self.max_huffman_table_rebuilds = max_huffman_table_rebuilds;
+    }
+
     /// Returns the adler32 checksum of the currently decompressed data.
     /// Note: Will return Some(1) if decompressing zlib but ignoring adler32.
     #[inline]
@@ -413,6 +479,8 @@ impl Default for DecompressorOxide {
             dist: 0,
             counter: 0,
             num_extra: 0,
+            huffman_table_rebuilds: 0,
+            max_huffman_table_rebuilds: u32::MAX,
             table_sizes: [0; MAX_HUFF_TABLES],
             bit_buf: 0,
             // TODO:(oyvindln) Check that copies here are optimized out in release mode.
@@ -472,6 +540,7 @@ enum State {
     BadCodeSizeDistPrevLookup,
     InvalidLitlen,
     InvalidDist,
+    HuffmanTableRebuildLimitExceeded,
 }
 
 impl State {
@@ -489,6 +558,7 @@ impl State {
                 | BadCodeSizeDistPrevLookup
                 | InvalidLitlen
                 | InvalidDist
+                | HuffmanTableRebuildLimitExceeded
         )
     }
 
@@ -871,6 +941,17 @@ fn init_tree(r: &mut DecompressorOxide, l: &mut LocalVars) -> Option<Action> {
     loop {
         let bt = r.block_type as usize;
 
+        // This is the expensive part of decoding a block header: rebuilding a table below is an
+        // O(FAST_LOOKUP_SIZE + MAX_HUFF_TREE_SIZE) operation regardless of how many symbols the
+        // table actually encodes, and a single degenerate deflate stream can force an
+        // unbounded number of these rebuilds while producing little or no output (see
+        // `DecompressorOxide::set_max_huffman_table_rebuilds`). Check the caller's budget, if
+        // any, before doing any of that work.
+        r.huffman_table_rebuilds = r.huffman_table_rebuilds.saturating_add(1);
+        if r.huffman_table_rebuilds > r.max_huffman_table_rebuilds {
+            return Some(Action::Jump(HuffmanTableRebuildLimitExceeded));
+        }
+
         let code_sizes = match bt {
             LITLEN_TABLE => &mut r.code_size_literal[..],
             DIST_TABLE => &mut r.code_size_dist,
@@ -1469,6 +1550,7 @@ pub fn decompress_with_limit(
                 r.z_header1 = 0;
                 r.z_adler32 = 1;
                 r.check_adler32 = 1;
+                r.huffman_table_rebuilds = 0;
                 if flags & TINFL_FLAG_PARSE_ZLIB_HEADER != 0 {
                     Action::Jump(State::ReadZlibCmf)
                 } else {
@@ -2017,6 +2099,12 @@ pub fn decompress_with_limit(
             // We are done.
             DoneForever => break TINFLStatus::Done,
 
+            // The caller-configured Huffman table rebuild budget (see
+            // `DecompressorOxide::set_max_huffman_table_rebuilds`) was exceeded.
+            HuffmanTableRebuildLimitExceeded => {
+                break TINFLStatus::HuffmanTableRebuildLimitExceeded
+            }
+
             // Anything else indicates failure.
             // BadZlibHeader | BadRawLength | BadDistOrLiteralTableLength | BlockTypeUnexpected |
             // DistanceOutOfBounds |

miniz_oxide/src/inflate/mod.rs

@@ -21,6 +21,7 @@ const TINFL_STATUS_NEEDS_MORE_INPUT: i32 = 1;
 const TINFL_STATUS_HAS_MORE_OUTPUT: i32 = 2;
 #[cfg(feature = "block-boundary")]
 const TINFL_STATUS_BLOCK_BOUNDARY: i32 = 3;
+const TINFL_STATUS_HUFFMAN_REBUILD_LIMIT_EXCEEDED: i32 = -5;
 
 /// Return status codes.
 #[repr(i8)]
@@ -66,6 +67,19 @@ pub enum TINFLStatus {
     /// There is still pending data that didn't fit in the output buffer.
     HasMoreOutput = TINFL_STATUS_HAS_MORE_OUTPUT as i8,
 
+    /// The number of Huffman decode table (re)builds allowed by a limit configured via
+    /// [`DecompressorOxide::set_max_huffman_table_rebuilds`][core::DecompressorOxide::set_max_huffman_table_rebuilds]
+    /// was exceeded, and decompression was stopped early.
+    ///
+    /// Every `BTYPE=1`/`BTYPE=2` (static/dynamic Huffman) deflate block header forces a full
+    /// rebuild of the decoder's Huffman tables, independently of how many symbols (including
+    /// zero) the block actually encodes. Since DEFLATE allows chaining an arbitrarily long
+    /// sequence of minimal, near-empty blocks, decompression cost can be decoupled from output
+    /// size this way; unlike [`HasMoreOutput`][Self::HasMoreOutput], an output-size limit alone
+    /// does not bound this cost, since such inputs can produce little or no output at all. This
+    /// status is only ever returned if the caller has opted in to a rebuild limit.
+    HuffmanTableRebuildLimitExceeded = TINFL_STATUS_HUFFMAN_REBUILD_LIMIT_EXCEEDED as i8,
+
     /// Reached the end of a deflate block, and the start of the next block.
     ///
     /// At this point, you can suspend decompression and later resume with a new `DecompressorOxide`.
@@ -90,6 +104,7 @@ impl TINFLStatus {
             TINFL_STATUS_DONE => Some(Done),
             TINFL_STATUS_NEEDS_MORE_INPUT => Some(NeedsMoreInput),
             TINFL_STATUS_HAS_MORE_OUTPUT => Some(HasMoreOutput),
+            TINFL_STATUS_HUFFMAN_REBUILD_LIMIT_EXCEEDED => Some(HuffmanTableRebuildLimitExceeded),
             #[cfg(feature = "block-boundary")]
             TINFL_STATUS_BLOCK_BOUNDARY => Some(BlockBoundary),
             _ => None,
@@ -119,6 +134,9 @@ impl alloc::fmt::Display for DecompressError {
             TINFLStatus::Done => "", // Unreachable
             TINFLStatus::NeedsMoreInput => "Truncated input stream",
             TINFLStatus::HasMoreOutput => "Output size exceeded the specified limit",
+            TINFLStatus::HuffmanTableRebuildLimitExceeded => {
+                "Huffman table rebuild limit exceeded"
+            }
             #[cfg(feature = "block-boundary")]
             TINFLStatus::BlockBoundary => "Reached end of a deflate block",
         })
@@ -144,7 +162,7 @@ fn decompress_error(status: TINFLStatus, output: Vec<u8>) -> Result<Vec<u8>, Dec
 #[inline]
 #[cfg(feature = "with-alloc")]
 pub fn decompress_to_vec(input: &[u8]) -> Result<Vec<u8>, DecompressError> {
-    decompress_to_vec_inner(input, 0, usize::MAX)
+    decompress_to_vec_inner(input, 0, usize::MAX, u32::MAX)
 }
 
 /// Decompress the deflate-encoded data (with a zlib wrapper) in `input` to a vector.
@@ -161,6 +179,7 @@ pub fn decompress_to_vec_zlib(input: &[u8]) -> Result<Vec<u8>, DecompressError>
         input,
         inflate_flags::TINFL_FLAG_PARSE_ZLIB_HEADER,
         usize::MAX,
+        u32::MAX,
     )
 }
 
@@ -179,7 +198,35 @@ pub fn decompress_to_vec_with_limit(
     input: &[u8],
     max_size: usize,
 ) -> Result<Vec<u8>, DecompressError> {
-    decompress_to_vec_inner(input, 0, max_size)
+    decompress_to_vec_inner(input, 0, max_size, u32::MAX)
+}
+
+/// Decompress the deflate-encoded data in `input` to a vector, bounding both the output size
+/// and the number of times the decompressor is allowed to rebuild its internal Huffman decode
+/// tables.
+///
+/// The vector is grown to at most `max_size` bytes, exactly like
+/// [`decompress_to_vec_with_limit`]. In addition, decompression stops early with
+/// [`TINFLStatus::HuffmanTableRebuildLimitExceeded`] if more than `max_huffman_table_rebuilds`
+/// deflate block headers are processed.
+///
+/// This second limit exists because a deflate stream can legally consist of an arbitrarily
+/// long chain of minimal, near-empty blocks that each force an expensive Huffman table rebuild
+/// while producing little or no output. Such an input can be very cheap to store (well under a
+/// couple MiB) yet expensive to decompress (multiple CPU seconds or more), all while never
+/// coming close to triggering `max_size` since so little (or nothing) is ever written to the
+/// output. See [`DecompressorOxide::set_max_huffman_table_rebuilds`][core::DecompressorOxide::set_max_huffman_table_rebuilds]
+/// for more details, and for guidance on picking a value for `max_huffman_table_rebuilds`.
+///
+/// Returns a [`Result`] containing the [`Vec`] of decompressed data on success, and a [struct][DecompressError] on failure.
+#[inline]
+#[cfg(feature = "with-alloc")]
+pub fn decompress_to_vec_with_limits(
+    input: &[u8],
+    max_size: usize,
+    max_huffman_table_rebuilds: u32,
+) -> Result<Vec<u8>, DecompressError> {
+    decompress_to_vec_inner(input, 0, max_size, max_huffman_table_rebuilds)
 }
 
 /// Decompress the deflate-encoded data (with a zlib wrapper) in `input` to a vector.
@@ -196,7 +243,35 @@ pub fn decompress_to_vec_zlib_with_limit(
     input: &[u8],
     max_size: usize,
 ) -> Result<Vec<u8>, DecompressError> {
-    decompress_to_vec_inner(input, inflate_flags::TINFL_FLAG_PARSE_ZLIB_HEADER, max_size)
+    decompress_to_vec_inner(
+        input,
+        inflate_flags::TINFL_FLAG_PARSE_ZLIB_HEADER,
+        max_size,
+        u32::MAX,
+    )
+}
+
+/// Decompress the deflate-encoded data (with a zlib wrapper) in `input` to a vector, bounding
+/// both the output size and the number of times the decompressor is allowed to rebuild its
+/// internal Huffman decode tables.
+///
+/// See [`decompress_to_vec_with_limits`] for details on `max_huffman_table_rebuilds` and why it
+/// exists.
+///
+/// Returns a [`Result`] containing the [`Vec`] of decompressed data on success, and a [struct][DecompressError] on failure.
+#[inline]
+#[cfg(feature = "with-alloc")]
+pub fn decompress_to_vec_zlib_with_limits(
+    input: &[u8],
+    max_size: usize,
+    max_huffman_table_rebuilds: u32,
+) -> Result<Vec<u8>, DecompressError> {
+    decompress_to_vec_inner(
+        input,
+        inflate_flags::TINFL_FLAG_PARSE_ZLIB_HEADER,
+        max_size,
+        max_huffman_table_rebuilds,
+    )
 }
 
 /// Backend of various to-[`Vec`] decompressions.
@@ -207,11 +282,13 @@ fn decompress_to_vec_inner(
     mut input: &[u8],
     flags: u32,
     max_output_size: usize,
+    max_huffman_table_rebuilds: u32,
 ) -> Result<Vec<u8>, DecompressError> {
     let flags = flags | inflate_flags::TINFL_FLAG_USING_NON_WRAPPING_OUTPUT_BUF;
     let mut ret: Vec<u8> = vec![0; input.len().saturating_mul(2).min(max_output_size)];
 
     let mut decomp = Box::<DecompressorOxide>::default();
+    decomp.set_max_huffman_table_rebuilds(max_huffman_table_rebuilds);
 
     let mut out_pos = 0;
     loop {