apfsds_obfuscation/
compression.rs

1//! Compression utilities
2
3use thiserror::Error;
4
5/// Compression threshold in bytes
6pub const COMPRESSION_THRESHOLD: usize = 1024;
7
8/// Default compression level
9pub const DEFAULT_COMPRESSION_LEVEL: i32 = 3;
10
11#[derive(Error, Debug)]
12pub enum CompressionError {
13    #[error("Compression failed: {0}")]
14    CompressionFailed(String),
15
16    #[error("Decompression failed: {0}")]
17    DecompressionFailed(String),
18
19    #[error("Data is not compressed")]
20    NotCompressed,
21}
22
23/// Compress data using zstd if above threshold
24pub fn compress_if_needed(data: &[u8]) -> Result<(Vec<u8>, bool), CompressionError> {
25    if data.len() < COMPRESSION_THRESHOLD {
26        return Ok((data.to_vec(), false));
27    }
28
29    compress(data).map(|compressed| (compressed, true))
30}
31
32/// Compress data using zstd
33pub fn compress(data: &[u8]) -> Result<Vec<u8>, CompressionError> {
34    compress_with_level(data, DEFAULT_COMPRESSION_LEVEL)
35}
36
37/// Compress data with specific level (1-22)
38pub fn compress_with_level(data: &[u8], level: i32) -> Result<Vec<u8>, CompressionError> {
39    zstd::encode_all(data, level).map_err(|e| CompressionError::CompressionFailed(e.to_string()))
40}
41
42/// Decompress zstd data
43pub fn decompress(data: &[u8]) -> Result<Vec<u8>, CompressionError> {
44    zstd::decode_all(data).map_err(|e| CompressionError::DecompressionFailed(e.to_string()))
45}
46
47/// Decompress with maximum size limit (for safety)
48pub fn decompress_with_limit(data: &[u8], max_size: usize) -> Result<Vec<u8>, CompressionError> {
49    let mut decoder = zstd::Decoder::new(data)
50        .map_err(|e| CompressionError::DecompressionFailed(e.to_string()))?;
51
52    let mut result = Vec::new();
53    let mut buf = [0u8; 8192];
54
55    loop {
56        use std::io::Read;
57        let n = decoder
58            .read(&mut buf)
59            .map_err(|e| CompressionError::DecompressionFailed(e.to_string()))?;
60
61        if n == 0 {
62            break;
63        }
64
65        if result.len() + n > max_size {
66            return Err(CompressionError::DecompressionFailed(format!(
67                "Decompressed size exceeds limit of {} bytes",
68                max_size
69            )));
70        }
71
72        result.extend_from_slice(&buf[..n]);
73    }
74
75    Ok(result)
76}
77
78/// Check if data might be zstd compressed (magic number: 0x28 0xB5 0x2F 0xFD)
79pub fn is_compressed(data: &[u8]) -> bool {
80    data.len() >= 4 && data[0] == 0x28 && data[1] == 0xB5 && data[2] == 0x2F && data[3] == 0xFD
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    #[test]
88    fn test_compress_decompress() {
89        let data = b"Hello, APFSDS! This is a test message that should be compressed.";
90        let compressed = compress(data).unwrap();
91        let decompressed = decompress(&compressed).unwrap();
92
93        assert_eq!(data.as_slice(), decompressed.as_slice());
94    }
95
96    #[test]
97    fn test_compress_if_needed_small() {
98        let small_data = vec![1, 2, 3, 4, 5];
99        let (result, compressed) = compress_if_needed(&small_data).unwrap();
100
101        assert!(!compressed);
102        assert_eq!(small_data, result);
103    }
104
105    #[test]
106    fn test_compress_if_needed_large() {
107        let large_data: Vec<u8> = (0..2000).map(|i| (i % 256) as u8).collect();
108        let (result, compressed) = compress_if_needed(&large_data).unwrap();
109
110        assert!(compressed);
111        assert_ne!(large_data, result);
112
113        // Verify we can decompress
114        let decompressed = decompress(&result).unwrap();
115        assert_eq!(large_data, decompressed);
116    }
117
118    #[test]
119    fn test_is_compressed() {
120        let data = b"uncompressed data";
121        assert!(!is_compressed(data));
122
123        let compressed = compress(data).unwrap();
124        assert!(is_compressed(&compressed));
125    }
126
127    #[test]
128    fn test_decompress_with_limit() {
129        let data: Vec<u8> = (0..10000).map(|i| (i % 256) as u8).collect();
130        let compressed = compress(&data).unwrap();
131
132        // Should fail with small limit
133        let result = decompress_with_limit(&compressed, 1000);
134        assert!(result.is_err());
135
136        // Should succeed with large limit
137        let result = decompress_with_limit(&compressed, 100000);
138        assert!(result.is_ok());
139        assert_eq!(result.unwrap(), data);
140    }
141}
apfsds_obfuscation/compression.rs

apfsds_obfuscation/
compression.rs
