threatflux_binary_analysis/utils/
compression.rs

1use crate::{Result, error::BinaryError};
2use flate2::read::{MultiGzDecoder, ZlibDecoder};
3use std::io::Read;
4
5/// Maximum decompressed size accepted by [`decompress`]: 64 MiB.
6pub const DEFAULT_MAX_DECOMPRESSED_SIZE: usize = 64 * 1024 * 1024;
7
8/// Decompress gzip or zlib data with a 64 MiB output limit.
9///
10/// Use [`decompress_with_limit`] when a different application budget is needed.
11#[cfg(feature = "compression")]
12pub fn decompress(data: &[u8]) -> Result<Vec<u8>> {
13    decompress_with_limit(data, DEFAULT_MAX_DECOMPRESSED_SIZE)
14}
15
16/// Decompress gzip or zlib data without allowing the output to exceed `max_size`.
17///
18/// The limit is enforced while streaming, before bytes are appended to the output
19/// allocation, which protects callers from high-expansion compressed inputs.
20#[cfg(feature = "compression")]
21pub fn decompress_with_limit(data: &[u8], max_size: usize) -> Result<Vec<u8>> {
22    if data.starts_with(b"\x1f\x8b") {
23        read_bounded(MultiGzDecoder::new(data), max_size)
24    } else if is_zlib_header(data) {
25        read_bounded(ZlibDecoder::new(data), max_size)
26    } else {
27        Err(BinaryError::invalid_data("unsupported compression format"))
28    }
29}
30
31fn is_zlib_header(data: &[u8]) -> bool {
32    let Some((&cmf, rest)) = data.split_first() else {
33        return false;
34    };
35    let Some(&flg) = rest.first() else {
36        return false;
37    };
38
39    // RFC 1950: DEFLATE (CM=8), a window no larger than 32 KiB (CINFO<=7),
40    // and a header divisible by 31. This accepts every valid compression-level
41    // encoding instead of recognizing only three common byte pairs.
42    cmf & 0x0f == 8 && cmf >> 4 <= 7 && (u16::from(cmf) << 8 | u16::from(flg)) % 31 == 0
43}
44
45fn read_bounded(mut reader: impl Read, max_size: usize) -> Result<Vec<u8>> {
46    const BUFFER_SIZE: usize = 8 * 1024;
47
48    let mut output = Vec::with_capacity(max_size.min(BUFFER_SIZE));
49    let mut buffer = [0_u8; BUFFER_SIZE];
50
51    loop {
52        let count = reader.read(&mut buffer)?;
53        if count == 0 {
54            return Ok(output);
55        }
56
57        let new_len = output
58            .len()
59            .checked_add(count)
60            .ok_or_else(|| BinaryError::invalid_data("decompressed size overflows usize"))?;
61        if new_len > max_size {
62            return Err(BinaryError::invalid_data(format!(
63                "decompressed data exceeds the {max_size}-byte limit"
64            )));
65        }
66
67        output
68            .try_reserve(count)
69            .map_err(|_| BinaryError::invalid_data("unable to allocate decompression output"))?;
70        output.extend_from_slice(&buffer[..count]);
71    }
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77    use flate2::Compression;
78    use flate2::write::{GzEncoder, ZlibEncoder};
79
80    #[cfg(feature = "compression")]
81    #[test]
82    fn test_decompress_gzip() {
83        let data = b"hello world";
84        let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
85        use std::io::Write;
86        encoder.write_all(data).unwrap();
87        let compressed = encoder.finish().unwrap();
88
89        let decompressed = decompress(&compressed).unwrap();
90        assert_eq!(decompressed, data);
91    }
92
93    #[cfg(feature = "compression")]
94    #[test]
95    fn test_decompress_zlib() {
96        let data = b"another test";
97        let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
98        use std::io::Write;
99        encoder.write_all(data).unwrap();
100        let compressed = encoder.finish().unwrap();
101
102        let decompressed = decompress(&compressed).unwrap();
103        assert_eq!(decompressed, data);
104    }
105
106    #[cfg(feature = "compression")]
107    #[test]
108    fn rejects_output_larger_than_limit() {
109        let data = vec![b'A'; 64 * 1024];
110        let mut encoder = GzEncoder::new(Vec::new(), Compression::best());
111        use std::io::Write;
112        encoder.write_all(&data).unwrap();
113        let compressed = encoder.finish().unwrap();
114
115        let error = decompress_with_limit(&compressed, 1024).unwrap_err();
116        assert!(error.to_string().contains("1024-byte limit"));
117    }
118
119    #[cfg(feature = "compression")]
120    #[test]
121    fn accepts_exact_output_limit() {
122        let data = b"exact limit";
123        let mut encoder = ZlibEncoder::new(Vec::new(), Compression::none());
124        use std::io::Write;
125        encoder.write_all(data).unwrap();
126        let compressed = encoder.finish().unwrap();
127
128        assert_eq!(
129            decompress_with_limit(&compressed, data.len()).unwrap(),
130            data
131        );
132    }
133
134    #[test]
135    fn recognizes_all_valid_rfc_1950_headers() {
136        assert!(is_zlib_header(&[0x78, 0x01]));
137        assert!(is_zlib_header(&[0x78, 0x5e]));
138        assert!(is_zlib_header(&[0x78, 0x9c]));
139        assert!(is_zlib_header(&[0x78, 0xda]));
140        assert!(!is_zlib_header(&[0x78, 0x00]));
141        assert!(!is_zlib_header(&[0x88, 0x1c]));
142    }
143}
threatflux_binary_analysis/utils/compression.rs

threatflux_binary_analysis/utils/
compression.rs
