heddle_format/compression/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2#![deny(clippy::cast_possible_truncation)]
3
4//! Compression utilities for Heddle storage.
5//!
6//! Provides configurable compression with support for:
7//! - zstd: High compression ratio, good speed
8//! - Delta encoding: For similar versions of the same file
9
10#[cfg(feature = "zstd")]
11use std::io::Read;
12
13const COMPRESSED_HEADER_LEN: usize = 9;
14const MAX_DECOMPRESSED_SIZE: u64 = 256 * 1024 * 1024;
15const ZSTD_MAGIC: [u8; 4] = [0x28, 0xB5, 0x2F, 0xFD];
16
17/// Compression algorithm selection.
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19#[repr(u8)]
20enum CompressionType {
21    /// Zstandard compression.
22    Zstd = 1,
23}
24
25impl CompressionType {
26    /// Convert from byte value.
27    fn from_u8(value: u8) -> Option<Self> {
28        match value {
29            1 => Some(CompressionType::Zstd),
30            _ => None,
31        }
32    }
33}
34
35/// Compression configuration.
36#[derive(Debug, Clone, Copy)]
37pub struct CompressionConfig {
38    /// Whether compression is enabled.
39    pub enabled: bool,
40    /// Compression level (algorithm-specific).
41    /// For zstd: 1-22 (1=fast, 22=best, 3=default)
42    pub level: i32,
43    /// Minimum size to compress (smaller objects aren't worth it).
44    pub min_size: usize,
45    /// Maximum size for delta compression base.
46    pub max_delta_size: usize,
47}
48
49impl Default for CompressionConfig {
50    fn default() -> Self {
51        Self {
52            enabled: cfg!(feature = "zstd"),
53            level: 3,                   // zstd default
54            min_size: 256,              // Don't compress tiny objects
55            max_delta_size: 10_000_000, // 10MB max for delta base
56        }
57    }
58}
59
60impl CompressionConfig {
61    /// Create configuration from environment variables.
62    pub fn from_env() -> Self {
63        let mut config = Self::default();
64
65        if let Ok(val) = std::env::var("HEDDLE_COMPRESSION") {
66            let requested = val != "0" && val.to_lowercase() != "false";
67            config.enabled = requested && cfg!(feature = "zstd");
68        }
69
70        if let Ok(val) = std::env::var("HEDDLE_COMPRESSION_LEVEL")
71            && let Ok(level) = val.parse::<i32>()
72        {
73            config.level = level.clamp(1, 22);
74        }
75
76        if let Ok(val) = std::env::var("HEDDLE_COMPRESSION_MIN_SIZE")
77            && let Ok(size) = val.parse::<usize>()
78        {
79            config.min_size = size;
80        }
81
82        config
83    }
84
85    /// Disable compression.
86    pub fn disabled() -> Self {
87        Self {
88            enabled: false,
89            level: 0,
90            min_size: usize::MAX,
91            max_delta_size: 0,
92        }
93    }
94}
95
96/// Compression error type.
97#[derive(Debug, thiserror::Error)]
98pub enum CompressionError {
99    #[error("decompression failed: {0}")]
100    DecompressionFailed(String),
101    #[error("compression failed: {0}")]
102    CompressionFailed(String),
103    #[error("invalid compression type: {0}")]
104    InvalidType(u8),
105    #[error("corrupted data: {0}")]
106    CorruptedData(String),
107    #[error("invalid operation: {0}")]
108    InvalidOperation(String),
109    #[error("object size {size} exceeds maximum {max}")]
110    SizeLimitExceeded { size: u64, max: u64 },
111}
112
113#[cfg(feature = "zstd")]
114/// Compress data using zstd.
115fn compress_zstd_impl(data: &[u8], level: i32) -> Result<Vec<u8>, CompressionError> {
116    zstd::encode_all(data, level).map_err(|e| CompressionError::CompressionFailed(e.to_string()))
117}
118
119#[cfg(not(feature = "zstd"))]
120fn compress_zstd_impl(_data: &[u8], _level: i32) -> Result<Vec<u8>, CompressionError> {
121    Err(CompressionError::InvalidOperation(
122        "zstd compression support not compiled into this build".to_string(),
123    ))
124}
125
126#[cfg(feature = "bench")]
127/// Compress data using zstd.
128pub fn compress_zstd(data: &[u8], level: i32) -> Result<Vec<u8>, CompressionError> {
129    compress_zstd_impl(data, level)
130}
131
132#[cfg(feature = "zstd")]
133/// Decompress zstd data while enforcing the recorded output size.
134fn decompress_zstd_impl(data: &[u8], expected_size: u64) -> Result<Vec<u8>, CompressionError> {
135    validate_size(expected_size)?;
136    let expected_capacity = checked_size_to_usize("zstd expected size", expected_size)?;
137
138    let mut decoder = zstd::stream::read::Decoder::new(data)
139        .map_err(|e| CompressionError::DecompressionFailed(e.to_string()))?;
140    let mut decompressed = Vec::with_capacity(expected_capacity);
141    let mut buffer = [0u8; 8192];
142
143    loop {
144        let bytes_read = decoder
145            .read(&mut buffer)
146            .map_err(|e| CompressionError::DecompressionFailed(e.to_string()))?;
147        if bytes_read == 0 {
148            break;
149        }
150
151        let next_size = decompressed.len().checked_add(bytes_read).ok_or_else(|| {
152            CompressionError::CorruptedData("decompressed size overflows".to_string())
153        })?;
154        let next_size = u64::try_from(next_size).map_err(|_| {
155            CompressionError::CorruptedData("decompressed size exceeds platform limits".to_string())
156        })?;
157        if next_size > expected_size {
158            return Err(CompressionError::CorruptedData(format!(
159                "decompressed size exceeds recorded header size: expected {expected_size}, got at least {next_size}",
160            )));
161        }
162
163        decompressed.extend_from_slice(&buffer[..bytes_read]);
164    }
165
166    Ok(decompressed)
167}
168
169#[cfg(not(feature = "zstd"))]
170fn decompress_zstd_impl(_data: &[u8], expected_size: u64) -> Result<Vec<u8>, CompressionError> {
171    validate_size(expected_size)?;
172    Err(CompressionError::InvalidOperation(
173        "zstd-compressed data is unsupported in this build".to_string(),
174    ))
175}
176
177#[cfg(feature = "bench")]
178/// Decompress zstd data while enforcing the recorded output size.
179pub fn decompress_zstd(data: &[u8], expected_size: u64) -> Result<Vec<u8>, CompressionError> {
180    decompress_zstd_impl(data, expected_size)
181}
182
183/// Compress data with automatic algorithm selection.
184///
185/// Returns the compressed data with header, or None if compression
186/// doesn't help (compressed would be larger).
187pub fn compress(
188    data: &[u8],
189    config: &CompressionConfig,
190) -> Result<Option<Vec<u8>>, CompressionError> {
191    if !config.enabled || data.len() < config.min_size {
192        return Ok(None);
193    }
194
195    validate_size(data.len() as u64)?;
196
197    // Try zstd compression
198    let compressed = compress_zstd_impl(data, config.level)?;
199
200    // Only use compression if it actually helps
201    if compressed.len() >= data.len() {
202        return Ok(None);
203    }
204
205    // Build header: [type][size][data]
206    let mut result = Vec::with_capacity(COMPRESSED_HEADER_LEN + compressed.len());
207    result.push(CompressionType::Zstd as u8);
208    result.extend_from_slice(&(data.len() as u64).to_be_bytes());
209    result.extend_from_slice(&compressed);
210
211    Ok(Some(result))
212}
213
214/// Decompress data based on header.
215///
216/// Returns the decompressed data, or original data if uncompressed.
217pub fn decompress(data: &[u8]) -> Result<Vec<u8>, CompressionError> {
218    if data.len() < COMPRESSED_HEADER_LEN {
219        // Too short for header, assume uncompressed
220        return Ok(data.to_vec());
221    }
222
223    let compression_type =
224        CompressionType::from_u8(data[0]).ok_or_else(|| CompressionError::InvalidType(data[0]))?;
225
226    match compression_type {
227        CompressionType::Zstd if zstd_header_len(data).is_some() => {
228            decompress_zstd_with_header(data)
229        }
230        CompressionType::Zstd => Ok(data.to_vec()),
231    }
232}
233
234/// Check if data is compressed (has compression header).
235pub fn is_compressed(data: &[u8]) -> bool {
236    if data.len() < COMPRESSED_HEADER_LEN {
237        return false;
238    }
239
240    matches!(
241        CompressionType::from_u8(data[0]),
242        Some(CompressionType::Zstd)
243    ) && zstd_header_len(data).is_some()
244}
245
246/// Peek at the recorded *uncompressed* size in a header-prefixed blob,
247/// without decompressing the payload. Returns `None` for short or
248/// unprefixed inputs (the caller can then fall back to the file length).
249///
250/// Used by header-only size queries (e.g. [`ObjectStore::blob_size`])
251/// where reading the full blob would dominate. Only the first 9 bytes
252/// of the input are consulted.
253pub fn header_uncompressed_size(data: &[u8]) -> Option<u64> {
254    if data.len() < COMPRESSED_HEADER_LEN {
255        return None;
256    }
257    let CompressionType::Zstd = CompressionType::from_u8(data[0])?;
258    zstd_header_len(data)?;
259    Some(u64::from_be_bytes(
260        data[1..COMPRESSED_HEADER_LEN].try_into().ok()?,
261    ))
262}
263
264#[cfg(test)]
265/// Get compression info from header.
266fn compression_info(data: &[u8]) -> Option<(CompressionType, u64)> {
267    if data.len() < COMPRESSED_HEADER_LEN {
268        return None;
269    }
270
271    let compression_type = CompressionType::from_u8(data[0])?;
272    let uncompressed_size = u64::from_be_bytes(data[1..COMPRESSED_HEADER_LEN].try_into().ok()?);
273
274    Some((compression_type, uncompressed_size))
275}
276
277fn decompress_zstd_with_header(data: &[u8]) -> Result<Vec<u8>, CompressionError> {
278    try_decompress_zstd(data, COMPRESSED_HEADER_LEN, read_u64_size)
279}
280
281fn zstd_header_len(data: &[u8]) -> Option<usize> {
282    if has_magic_at(data, COMPRESSED_HEADER_LEN, ZSTD_MAGIC) {
283        Some(COMPRESSED_HEADER_LEN)
284    } else {
285        None
286    }
287}
288
289fn try_decompress_zstd<F>(
290    data: &[u8],
291    header_len: usize,
292    read_size: F,
293) -> Result<Vec<u8>, CompressionError>
294where
295    F: Fn(&[u8]) -> Result<u64, CompressionError>,
296{
297    let uncompressed_size = read_size(data)?;
298    let decompressed = decompress_zstd_impl(&data[header_len..], uncompressed_size)?;
299    validate_decompressed_len(uncompressed_size, decompressed.len())?;
300    Ok(decompressed)
301}
302
303fn read_u64_size(data: &[u8]) -> Result<u64, CompressionError> {
304    if data.len() < COMPRESSED_HEADER_LEN {
305        return Err(CompressionError::CorruptedData(
306            "compression header truncated".to_string(),
307        ));
308    }
309
310    let recorded_size =
311        u64::from_be_bytes(data[1..COMPRESSED_HEADER_LEN].try_into().map_err(|_| {
312            CompressionError::CorruptedData("compression header truncated".to_string())
313        })?);
314    validate_size(recorded_size)?;
315    Ok(recorded_size)
316}
317
318fn validate_size(size: u64) -> Result<(), CompressionError> {
319    if size > MAX_DECOMPRESSED_SIZE {
320        return Err(CompressionError::SizeLimitExceeded {
321            size,
322            max: MAX_DECOMPRESSED_SIZE,
323        });
324    }
325
326    Ok(())
327}
328
329#[cfg(feature = "zstd")]
330fn checked_size_to_usize(field: &str, size: u64) -> Result<usize, CompressionError> {
331    usize::try_from(size)
332        .map_err(|_| CompressionError::CorruptedData(format!("{field} exceeds platform limits")))
333}
334
335fn validate_decompressed_len(expected: u64, actual: usize) -> Result<(), CompressionError> {
336    if actual as u64 != expected {
337        return Err(CompressionError::CorruptedData(format!(
338            "decompressed size mismatch: expected {expected}, got {actual}",
339        )));
340    }
341
342    Ok(())
343}
344
345fn has_magic_at(data: &[u8], offset: usize, magic: [u8; 4]) -> bool {
346    data.get(offset..offset + magic.len()) == Some(magic.as_slice())
347}
348
349#[cfg(test)]
350mod compression_tests;
heddle_format/compression/mod.rs

heddle_format/compression/
mod.rs
