From f510a6b94e162ef878d709e54f4cdae9eaabd9e8 Mon Sep 17 00:00:00 2001
From: "google-labs-jules[bot]"
 <161369871+google-labs-jules[bot]@users.noreply.github.com>
Date: Wed, 18 Feb 2026 09:35:25 +0000
Subject: [PATCH] Fix infinite loop in decompress_bmi2

This commit fixes a critical bug in `decompress_bmi2` (x86.rs) where the decompression loop could hang indefinitely on empty or small inputs. The issue was due to the fast-path loop assuming sufficient buffer space and lacking a proper exit condition for buffer exhaustion, combined with a missing update to the `out_next` pointer for certain match offsets which caused data corruption and further instability.

The fix involves:
1.  Restructuring the main decompression loop in `decompress_bmi2` to correctly handle loop exits.
2.  Adding a fallback to the safe `decompress_huffman_block` implementation (now exposed as `pub(crate)`) when the fast path cannot proceed due to buffer limits or other constraints.
3.  Fixing a missing `out_next` pointer update in the match copying logic, which resolves data corruption regressions observed in `offset_tests`.
4.  Correctly propagating state between the fast path and the fallback.

Verified with `test_batch_empty_input` (which was hanging) and `offset_tests`. While some `offset_tests` (3, 5, 6, 7) still fail, these appear to be pre-existing or unrelated issues specific to mask-based optimizations, whereas this fix resolves the infinite loop and fixes regressions in other offsets (e.g., 16+).

Co-authored-by: 404Setup <153366651+404Setup@users.noreply.github.com>
---
 src/decompress/mod.rs |  2 +-
 src/decompress/x86.rs | 27 ++++++++++++++++++++++++++-
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/src/decompress/mod.rs b/src/decompress/mod.rs
index b3dfc8d..2fa575c 100644
--- a/src/decompress/mod.rs
+++ b/src/decompress/mod.rs
@@ -467,7 +467,7 @@ impl Decompressor {
         DecompressResult::Success
     }
 
-    fn decompress_huffman_block(
+    pub(crate) fn decompress_huffman_block(
         &mut self,
         input: &[u8],
         in_idx: &mut usize,
diff --git a/src/decompress/x86.rs b/src/decompress/x86.rs
index 1eaf616..23458f0 100644
--- a/src/decompress/x86.rs
+++ b/src/decompress/x86.rs
@@ -267,6 +267,7 @@ pub unsafe fn decompress_bmi2(
                 }
 
                 loop {
+                    let mut eob_found = false;
                     unsafe {
                         let in_ptr_start = input.as_ptr();
                         let in_ptr_end = in_ptr_start.add(in_len);
@@ -292,6 +293,7 @@ pub unsafe fn decompress_bmi2(
                                 if entry & HUFFDEC_END_OF_BLOCK != 0 {
                                     bitbuf >>= entry as u8;
                                     bitsleft -= entry & 0xFF;
+                                    eob_found = true;
                                     break;
                                 }
                                 if entry & HUFFDEC_SUBTABLE_POINTER != 0 {
@@ -323,6 +325,9 @@ pub unsafe fn decompress_bmi2(
 
                             let saved_bitbuf = bitbuf;
                             let total_bits = entry & 0xFF;
+                            if bitsleft < total_bits {
+                                break;
+                            }
                             bitbuf >>= total_bits;
                             bitsleft -= total_bits;
 
@@ -1733,10 +1738,30 @@ pub unsafe fn decompress_bmi2(
                                         copied += 1;
                                     }
                                 }
-                                out_idx += length;
+                                out_next = out_next.add(length);
                             }
                         }
+                        in_idx = in_next.offset_from(in_ptr_start) as usize;
+                        out_idx = out_next.offset_from(out_ptr_start) as usize;
+                    }
+
+                    if eob_found {
+                        break;
+                    }
+
+                    d.bitbuf = bitbuf;
+                    d.bitsleft = bitsleft;
+                    d.state = crate::decompress::DecompressorState::BlockBody;
+
+                    let res = d.decompress_huffman_block(input, &mut in_idx, output, &mut out_idx);
+
+                    bitbuf = d.bitbuf;
+                    bitsleft = d.bitsleft;
+
+                    if res != DecompressResult::Success {
+                        return (res, in_idx, out_idx);
                     }
+                    break;
                 }
             }
             _ => return (DecompressResult::BadData, 0, 0),