Skip to content

Commit 7ac151a

Browse files
committed
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.
1 parent 90c8d41 commit 7ac151a

2 files changed

Lines changed: 36 additions & 9 deletions

File tree

src/api.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@ impl Compressor {
9191
Ok(output)
9292
}
9393
CompressResult::InsufficientSpace => Err(io::Error::other("Insufficient space")),
94+
CompressResult::InternalError => Err(io::Error::other("Compression failed")),
9495
}
9596
}
9697

@@ -116,10 +117,9 @@ impl Compressor {
116117
}
117118
let out_uninit = crate::common::slice_as_uninit_mut(output);
118119
let (res, size) = f(&mut self.inner, data, out_uninit);
119-
if res == CompressResult::Success {
120-
Ok(size)
121-
} else {
122-
Err(io::Error::other(error_msg))
120+
match res {
121+
CompressResult::Success => Ok(size),
122+
_ => Err(io::Error::other(error_msg)),
123123
}
124124
}
125125
}

src/compress/mod.rs

Lines changed: 32 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -210,6 +210,7 @@ fn compute_static_tables() -> StaticTables {
210210
pub enum CompressResult {
211211
Success,
212212
InsufficientSpace,
213+
InternalError,
213214
}
214215

215216
#[derive(Clone, Copy)]
@@ -740,7 +741,10 @@ impl Compressor {
740741

741742
let mut bs = Bitstream::new(output);
742743

743-
let mut mf_enum = self.mf.take().unwrap();
744+
let mut mf_enum = match self.mf.take() {
745+
Some(mf) => mf,
746+
None => return (CompressResult::InternalError, 0, 0),
747+
};
744748

745749
let res = match &mut mf_enum {
746750
MatchFinderEnum::Chain(mf) => self.compress_loop(mf, input, &mut bs, flush_mode),
@@ -1030,18 +1034,21 @@ impl Compressor {
10301034
(processed, bits)
10311035
}
10321036

1033-
pub fn compress_to_size(&mut self, input: &[u8], final_block: bool) -> usize {
1037+
pub fn compress_to_size(&mut self, input: &[u8], final_block: bool) -> (CompressResult, usize) {
10341038
if self.compression_level == 0 {
10351039
let num_blocks = input.len() / 65535
10361040
+ if !input.len().is_multiple_of(65535) || (input.is_empty() && final_block) {
10371041
1
10381042
} else {
10391043
0
10401044
};
1041-
return input.len() + num_blocks * 5;
1045+
return (CompressResult::Success, input.len() + num_blocks * 5);
10421046
}
10431047

1044-
let mut mf_enum = self.mf.take().unwrap();
1048+
let mut mf_enum = match self.mf.take() {
1049+
Some(mf) => mf,
1050+
None => return (CompressResult::InternalError, 0),
1051+
};
10451052

10461053
let res = match &mut mf_enum {
10471054
MatchFinderEnum::Chain(mf) => self.compress_to_size_loop(mf, input, final_block),
@@ -1050,7 +1057,7 @@ impl Compressor {
10501057
};
10511058

10521059
self.mf = Some(mf_enum);
1053-
res
1060+
(CompressResult::Success, res)
10541061
}
10551062

10561063
fn accumulate_greedy_frequencies<T: MatchFinderTrait>(
@@ -2407,3 +2414,23 @@ impl Compressor {
24072414
(CompressResult::Success, out_idx)
24082415
}
24092416
}
2417+
2418+
#[cfg(test)]
2419+
mod tests {
2420+
use super::*;
2421+
2422+
#[test]
2423+
fn test_compressor_missing_mf() {
2424+
let mut compressor = Compressor::new(6);
2425+
compressor.mf = None;
2426+
2427+
let input = b"hello world";
2428+
let mut output = [MaybeUninit::uninit(); 100];
2429+
2430+
let (res, _, _) = compressor.compress(input, &mut output, FlushMode::Finish);
2431+
assert_eq!(res, CompressResult::InternalError);
2432+
2433+
let (res, _) = compressor.compress_to_size(input, true);
2434+
assert_eq!(res, CompressResult::InternalError);
2435+
}
2436+
}

0 commit comments

Comments
 (0)