Skip to main content

firecloud_storage/
compression.rs

1//! Adaptive compression using Zstd and LZ4
2
3use crate::{StorageError, StorageResult};
4
5/// Compression level presets
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum CompressionLevel {
8    /// No compression (for already compressed data like images, videos)
9    None,
10    /// Fast compression using LZ4 (for real-time transfers)
11    Fast,
12    /// Balanced Zstd level 3
13    Balanced,
14    /// Best ratio Zstd level 9 (for storage)
15    Best,
16}
17
18impl Default for CompressionLevel {
19    fn default() -> Self {
20        Self::Balanced
21    }
22}
23
24/// Compress data with the specified level
25pub fn compress(data: &[u8], level: CompressionLevel) -> StorageResult<Vec<u8>> {
26    match level {
27        CompressionLevel::None => Ok(data.to_vec()),
28
29        CompressionLevel::Fast => {
30            // LZ4 compression
31            Ok(lz4_flex::compress_prepend_size(data))
32        }
33
34        CompressionLevel::Balanced => {
35            // Zstd level 3
36            zstd::encode_all(data, 3).map_err(|e| StorageError::Compression(e.to_string()))
37        }
38
39        CompressionLevel::Best => {
40            // Zstd level 9
41            zstd::encode_all(data, 9).map_err(|e| StorageError::Compression(e.to_string()))
42        }
43    }
44}
45
46/// Decompress data (auto-detects format based on magic bytes)
47pub fn decompress(data: &[u8], was_lz4: bool) -> StorageResult<Vec<u8>> {
48    if was_lz4 {
49        // LZ4 decompression
50        lz4_flex::decompress_size_prepended(data)
51            .map_err(|e| StorageError::Decompression(e.to_string()))
52    } else {
53        // Zstd decompression
54        zstd::decode_all(data).map_err(|e| StorageError::Decompression(e.to_string()))
55    }
56}
57
58/// Detect if a file type is already compressed (skip compression)
59/// 
60/// TODO: Use this heuristic in chunker.rs to avoid redundant compression
61/// of already-compressed formats (images, videos, archives)
62#[allow(dead_code)]
63pub fn should_compress(mime_type: Option<&str>, data: &[u8]) -> bool {
64    // Check MIME type first
65    if let Some(mime) = mime_type {
66        let skip_types = [
67            "image/jpeg",
68            "image/png",
69            "image/gif",
70            "image/webp",
71            "video/",
72            "audio/",
73            "application/zip",
74            "application/gzip",
75            "application/x-7z-compressed",
76            "application/x-rar-compressed",
77        ];
78
79        for skip in &skip_types {
80            if mime.starts_with(skip) {
81                return false;
82            }
83        }
84    }
85
86    // Check magic bytes for common compressed formats
87    if data.len() >= 4 {
88        // JPEG
89        if data.starts_with(&[0xFF, 0xD8, 0xFF]) {
90            return false;
91        }
92        // PNG
93        if data.starts_with(&[0x89, 0x50, 0x4E, 0x47]) {
94            return false;
95        }
96        // GIF
97        if data.starts_with(b"GIF8") {
98            return false;
99        }
100        // ZIP/DOCX/XLSX
101        if data.starts_with(&[0x50, 0x4B, 0x03, 0x04]) {
102            return false;
103        }
104        // GZIP
105        if data.starts_with(&[0x1F, 0x8B]) {
106            return false;
107        }
108        // MP4/MOV
109        if data.len() >= 8 && &data[4..8] == b"ftyp" {
110            return false;
111        }
112    }
113
114    true
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    #[test]
122    fn test_zstd_roundtrip() {
123        let original = b"Hello, FireCloud! This is test data for compression.";
124
125        let compressed = compress(original, CompressionLevel::Balanced).unwrap();
126        let decompressed = decompress(&compressed, false).unwrap();
127
128        assert_eq!(decompressed, original);
129    }
130
131    #[test]
132    fn test_lz4_roundtrip() {
133        let original = b"Hello, FireCloud! This is test data for LZ4 compression.";
134
135        let compressed = compress(original, CompressionLevel::Fast).unwrap();
136        let decompressed = decompress(&compressed, true).unwrap();
137
138        assert_eq!(decompressed, original);
139    }
140
141    #[test]
142    fn test_no_compression() {
143        let original = b"Raw data";
144
145        let result = compress(original, CompressionLevel::None).unwrap();
146
147        assert_eq!(result, original);
148    }
149
150    #[test]
151    fn test_should_compress() {
152        // JPEG magic bytes - should NOT compress
153        let jpeg = [0xFF, 0xD8, 0xFF, 0xE0, 0x00, 0x10];
154        assert!(!should_compress(None, &jpeg));
155
156        // Text data - should compress
157        let text = b"This is plain text that should be compressed";
158        assert!(should_compress(None, text));
159
160        // MIME type check
161        assert!(!should_compress(Some("image/jpeg"), &[]));
162        assert!(should_compress(Some("text/plain"), &[]));
163    }
164}