From 7ac151a5dc2c35df96842501e2fc96a9d75b8f2c Mon Sep 17 00:00:00 2001
From: 404Setup <153366651+404Setup@users.noreply.github.com>
Date: Wed, 29 Apr 2026 06:37:32 +0000
Subject: [PATCH] fix: handle missing match finder in Compressor safely

Replaced unsafe .unwrap() calls on self.mf.take() with safe checks that return a new InternalError variant.
This prevents a potential Denial of Service (DoS) panic if the match finder is unexpectedly None, for example, due to a previous panic poisoning the Compressor state.
Updated the return signature of Compressor::compress_to_size to support error propagation.
---
 src/api.rs          |  8 ++++----
 src/compress/mod.rs | 37 ++++++++++++++++++++++++++++++++-----
 2 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/src/api.rs b/src/api.rs
index 5ce58fc..62d506c 100644
--- a/src/api.rs
+++ b/src/api.rs
@@ -91,6 +91,7 @@ impl Compressor {
                 Ok(output)
             }
             CompressResult::InsufficientSpace => Err(io::Error::other("Insufficient space")),
+            CompressResult::InternalError => Err(io::Error::other("Compression failed")),
         }
     }
 
@@ -116,10 +117,9 @@ impl Compressor {
         }
         let out_uninit = crate::common::slice_as_uninit_mut(output);
         let (res, size) = f(&mut self.inner, data, out_uninit);
-        if res == CompressResult::Success {
-            Ok(size)
-        } else {
-            Err(io::Error::other(error_msg))
+        match res {
+            CompressResult::Success => Ok(size),
+            _ => Err(io::Error::other(error_msg)),
         }
     }
 }
diff --git a/src/compress/mod.rs b/src/compress/mod.rs
index 11567e8..b7342af 100644
--- a/src/compress/mod.rs
+++ b/src/compress/mod.rs
@@ -210,6 +210,7 @@ fn compute_static_tables() -> StaticTables {
 pub enum CompressResult {
     Success,
     InsufficientSpace,
+    InternalError,
 }
 
 #[derive(Clone, Copy)]
@@ -740,7 +741,10 @@ impl Compressor {
 
         let mut bs = Bitstream::new(output);
 
-        let mut mf_enum = self.mf.take().unwrap();
+        let mut mf_enum = match self.mf.take() {
+            Some(mf) => mf,
+            None => return (CompressResult::InternalError, 0, 0),
+        };
 
         let res = match &mut mf_enum {
             MatchFinderEnum::Chain(mf) => self.compress_loop(mf, input, &mut bs, flush_mode),
@@ -1030,7 +1034,7 @@ impl Compressor {
         (processed, bits)
     }
 
-    pub fn compress_to_size(&mut self, input: &[u8], final_block: bool) -> usize {
+    pub fn compress_to_size(&mut self, input: &[u8], final_block: bool) -> (CompressResult, usize) {
         if self.compression_level == 0 {
             let num_blocks = input.len() / 65535
                 + if !input.len().is_multiple_of(65535) || (input.is_empty() && final_block) {
@@ -1038,10 +1042,13 @@ impl Compressor {
                 } else {
                     0
                 };
-            return input.len() + num_blocks * 5;
+            return (CompressResult::Success, input.len() + num_blocks * 5);
         }
 
-        let mut mf_enum = self.mf.take().unwrap();
+        let mut mf_enum = match self.mf.take() {
+            Some(mf) => mf,
+            None => return (CompressResult::InternalError, 0),
+        };
 
         let res = match &mut mf_enum {
             MatchFinderEnum::Chain(mf) => self.compress_to_size_loop(mf, input, final_block),
@@ -1050,7 +1057,7 @@ impl Compressor {
         };
 
         self.mf = Some(mf_enum);
-        res
+        (CompressResult::Success, res)
     }
 
     fn accumulate_greedy_frequencies<T: MatchFinderTrait>(
@@ -2407,3 +2414,23 @@ impl Compressor {
         (CompressResult::Success, out_idx)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_compressor_missing_mf() {
+        let mut compressor = Compressor::new(6);
+        compressor.mf = None;
+
+        let input = b"hello world";
+        let mut output = [MaybeUninit::uninit(); 100];
+
+        let (res, _, _) = compressor.compress(input, &mut output, FlushMode::Finish);
+        assert_eq!(res, CompressResult::InternalError);
+
+        let (res, _) = compressor.compress_to_size(input, true);
+        assert_eq!(res, CompressResult::InternalError);
+    }
+}