nodedb_codec/
lz4.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Bounded LZ4 block compression for byte columns.
4//!
5//! Frames contain a declared total, block size, block count, a length table,
6//! and independently compressed blocks. Decoding validates all framing before
7//! allocation and each LZ4 block before extending the result.
8
9use crate::bounds::{
10    checked_add, checked_capacity, checked_mul, checked_range, decoded_len, encode_input_len,
11    encode_u32_len, u32_to_usize,
12};
13use crate::error::CodecError;
14
15/// Default block size for LZ4 compression (4 KiB).
16const DEFAULT_BLOCK_SIZE: usize = 4096;
17const HEADER_SIZE: usize = 12;
18const MIN_BLOCK_SIZE: usize = 64;
19
20/// Compress raw bytes using LZ4 block compression.
21pub fn encode(data: &[u8]) -> Result<Vec<u8>, CodecError> {
22    encode_with_block_size(data, DEFAULT_BLOCK_SIZE)
23}
24
25/// Compress with a custom block size (useful for testing or tuning).
26///
27/// Inputs and framing fields are rejected with typed errors when they exceed
28/// the application or wire-format limits.
29pub fn encode_with_block_size(data: &[u8], block_size: usize) -> Result<Vec<u8>, CodecError> {
30    let total_len = encode_input_len(data.len(), "LZ4")?;
31    let block_size = block_size.max(MIN_BLOCK_SIZE);
32    decoded_len(block_size, "LZ4 block")?;
33    let block_size_u32 = encode_u32_len(block_size, "LZ4 block size")?;
34    let block_count = if data.is_empty() {
35        0
36    } else {
37        data.len().div_ceil(block_size)
38    };
39    let block_count_u32 = encode_u32_len(block_count, "LZ4 block count")?;
40    let table_len = checked_mul(block_count, 4, "LZ4 block-length table")?;
41    let initial_capacity = checked_add(
42        checked_add(HEADER_SIZE, table_len, "LZ4 output header")?,
43        data.len(),
44        "LZ4 output capacity",
45    )?;
46    let mut out = Vec::with_capacity(initial_capacity);
47    out.extend_from_slice(&total_len.to_le_bytes());
48    out.extend_from_slice(&block_size_u32.to_le_bytes());
49    out.extend_from_slice(&block_count_u32.to_le_bytes());
50    let lengths_offset = out.len();
51    out.resize(lengths_offset + table_len, 0);
52
53    for (index, chunk) in data.chunks(block_size).enumerate() {
54        let compressed = lz4_flex::compress_prepend_size(chunk);
55        let compressed_len = encode_u32_len(compressed.len(), "LZ4 compressed block")?;
56        let table_pos = lengths_offset + index * 4;
57        out[table_pos..table_pos + 4].copy_from_slice(&compressed_len.to_le_bytes());
58        out.extend_from_slice(&compressed);
59    }
60    Ok(out)
61}
62
63/// Decompress LZ4 block-compressed bytes back to raw data.
64pub fn decode(data: &[u8]) -> Result<Vec<u8>, CodecError> {
65    let header = read_header(data)?;
66    let result_capacity = checked_capacity(header.uncompressed_size, 1, "LZ4 decoded bytes")?;
67    let mut result = Vec::with_capacity(result_capacity);
68    let mut block_offset = header.data_offset;
69
70    for (index, &compressed_len) in header.block_lengths.iter().enumerate() {
71        let block = checked_range(data, block_offset, compressed_len, "LZ4 block range")?;
72        let expected_block_len = expected_block_len(&header, index)?;
73        validate_prepended_size(block, expected_block_len, index)?;
74        let decompressed = lz4_flex::decompress_size_prepended(block).map_err(|error| {
75            CodecError::DecompressFailed {
76                detail: format!("LZ4 block {index}: {error}"),
77            }
78        })?;
79        if decompressed.len() != expected_block_len {
80            return Err(CodecError::Corrupt {
81                detail: format!(
82                    "LZ4 block {index} decoded {} bytes, expected {expected_block_len}",
83                    decompressed.len()
84                ),
85            });
86        }
87        result.extend_from_slice(&decompressed);
88        block_offset = checked_add(block_offset, compressed_len, "LZ4 block offset")?;
89    }
90    if result.len() != header.uncompressed_size {
91        return Err(CodecError::Corrupt {
92            detail: format!(
93                "LZ4 uncompressed size mismatch: header says {}, got {}",
94                header.uncompressed_size,
95                result.len()
96            ),
97        });
98    }
99    Ok(result)
100}
101
102/// Decompress a single block by index (for random access).
103pub fn decode_block(data: &[u8], block_idx: usize) -> Result<Vec<u8>, CodecError> {
104    let header = read_header(data)?;
105    let compressed_len =
106        *header
107            .block_lengths
108            .get(block_idx)
109            .ok_or_else(|| CodecError::Corrupt {
110                detail: format!("LZ4 block index {block_idx} out of range"),
111            })?;
112    let mut block_offset = header.data_offset;
113    for &len in &header.block_lengths[..block_idx] {
114        block_offset = checked_add(block_offset, len, "LZ4 preceding block offset")?;
115    }
116    let block = checked_range(data, block_offset, compressed_len, "LZ4 block range")?;
117    let expected_len = expected_block_len(&header, block_idx)?;
118    validate_prepended_size(block, expected_len, block_idx)?;
119    let decoded = lz4_flex::decompress_size_prepended(block).map_err(|error| {
120        CodecError::DecompressFailed {
121            detail: format!("LZ4 block {block_idx}: {error}"),
122        }
123    })?;
124    if decoded.len() != expected_len {
125        return Err(CodecError::Corrupt {
126            detail: format!(
127                "LZ4 block {block_idx} decoded {} bytes, expected {expected_len}",
128                decoded.len()
129            ),
130        });
131    }
132    Ok(decoded)
133}
134
135struct Lz4Header {
136    uncompressed_size: usize,
137    block_size: usize,
138    block_lengths: Vec<usize>,
139    data_offset: usize,
140}
141
142fn read_header(data: &[u8]) -> Result<Lz4Header, CodecError> {
143    checked_range(data, 0, HEADER_SIZE, "LZ4 header")?;
144    let uncompressed_size = decoded_len(
145        u32_to_usize(read_u32(data, 0, "LZ4 size")?, "LZ4 size")?,
146        "LZ4",
147    )?;
148    let block_size = u32_to_usize(read_u32(data, 4, "LZ4 block size")?, "LZ4 block size")?;
149    let block_count = u32_to_usize(read_u32(data, 8, "LZ4 block count")?, "LZ4 block count")?;
150    if !(MIN_BLOCK_SIZE..=crate::bounds::MAX_DECODED_BYTES).contains(&block_size) {
151        return Err(CodecError::Corrupt {
152            detail: "invalid LZ4 block size".into(),
153        });
154    }
155    let expected_count = if uncompressed_size == 0 {
156        0
157    } else {
158        uncompressed_size.div_ceil(block_size)
159    };
160    if block_count != expected_count {
161        return Err(CodecError::Corrupt {
162            detail: "LZ4 block count does not match declared output".into(),
163        });
164    }
165    let table_len = checked_mul(block_count, 4, "LZ4 block-length table")?;
166    let data_offset = checked_add(HEADER_SIZE, table_len, "LZ4 data offset")?;
167    let table = checked_range(data, HEADER_SIZE, table_len, "LZ4 block-length table")?;
168    let block_length_capacity = checked_capacity(
169        block_count,
170        std::mem::size_of::<usize>(),
171        "LZ4 block lengths",
172    )?;
173    let mut block_lengths = Vec::with_capacity(block_length_capacity);
174    let mut total_compressed = 0usize;
175    for bytes in table.as_chunks::<4>().0 {
176        let len = u32_to_usize(u32::from_le_bytes(*bytes), "LZ4 block length")?;
177        total_compressed = checked_add(total_compressed, len, "LZ4 compressed payload")?;
178        block_lengths.push(len);
179    }
180    checked_range(
181        data,
182        data_offset,
183        total_compressed,
184        "LZ4 compressed payload",
185    )?;
186    if checked_add(data_offset, total_compressed, "LZ4 frame end")? != data.len() {
187        return Err(CodecError::Corrupt {
188            detail: "trailing bytes after LZ4 frame".into(),
189        });
190    }
191    Ok(Lz4Header {
192        uncompressed_size,
193        block_size,
194        block_lengths,
195        data_offset,
196    })
197}
198
199fn read_u32(data: &[u8], offset: usize, context: &str) -> Result<u32, CodecError> {
200    let bytes = checked_range(data, offset, 4, context)?;
201    Ok(u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
202}
203
204fn expected_block_len(header: &Lz4Header, index: usize) -> Result<usize, CodecError> {
205    let position = checked_mul(index, header.block_size, "LZ4 block position")?;
206    let remaining = header
207        .uncompressed_size
208        .checked_sub(position)
209        .ok_or_else(|| CodecError::Corrupt {
210            detail: "LZ4 block position exceeds declared output".into(),
211        })?;
212    Ok(remaining.min(header.block_size))
213}
214
215fn validate_prepended_size(block: &[u8], expected: usize, index: usize) -> Result<(), CodecError> {
216    let size = u32_to_usize(
217        read_u32(block, 0, "LZ4 block size prefix")?,
218        "LZ4 prepended output size",
219    )?;
220    if size != expected {
221        return Err(CodecError::Corrupt {
222            detail: format!("LZ4 block {index} declares {size} bytes, expected {expected}"),
223        });
224    }
225    Ok(())
226}
227
228/// Streaming LZ4 encoder. Accumulates data and compresses on `finish()`.
229pub struct Lz4Encoder {
230    buf: Vec<u8>,
231    block_size: usize,
232}
233impl Lz4Encoder {
234    pub fn new() -> Self {
235        Self {
236            buf: Vec::with_capacity(DEFAULT_BLOCK_SIZE),
237            block_size: DEFAULT_BLOCK_SIZE,
238        }
239    }
240    pub fn with_block_size(block_size: usize) -> Result<Self, CodecError> {
241        let block_size = block_size.max(MIN_BLOCK_SIZE);
242        decoded_len(block_size, "LZ4 block")?;
243        Ok(Self {
244            buf: Vec::new(),
245            block_size,
246        })
247    }
248    pub fn push(&mut self, data: &[u8]) -> Result<(), CodecError> {
249        let next_len = checked_add(self.buf.len(), data.len(), "LZ4 streaming input")?;
250        decoded_len(next_len, "LZ4 streaming input")?;
251        self.buf.extend_from_slice(data);
252        Ok(())
253    }
254    pub fn len(&self) -> usize {
255        self.buf.len()
256    }
257    pub fn is_empty(&self) -> bool {
258        self.buf.is_empty()
259    }
260    pub fn finish(self) -> Result<Vec<u8>, CodecError> {
261        encode_with_block_size(&self.buf, self.block_size)
262    }
263}
264impl Default for Lz4Encoder {
265    fn default() -> Self {
266        Self::new()
267    }
268}
269pub struct Lz4Decoder;
270impl Lz4Decoder {
271    pub fn decode_all(data: &[u8]) -> Result<Vec<u8>, CodecError> {
272        decode(data)
273    }
274    pub fn decode_block(data: &[u8], block_idx: usize) -> Result<Vec<u8>, CodecError> {
275        decode_block(data, block_idx)
276    }
277    pub fn block_count(data: &[u8]) -> Result<usize, CodecError> {
278        Ok(read_header(data)?.block_lengths.len())
279    }
280}
281
282#[cfg(test)]
283mod tests {
284    use super::*;
285    use crate::bounds::MAX_DECODED_BYTES;
286
287    #[test]
288    fn empty_roundtrip() {
289        let encoded = encode(&[]).expect("encode");
290        assert!(decode(&encoded).expect("decode").is_empty());
291    }
292
293    #[test]
294    fn small_data_roundtrip() {
295        let data = b"hello world, this is a log message";
296        let encoded = encode(data).expect("encode");
297        assert_eq!(decode(&encoded).expect("decode"), data);
298    }
299
300    #[test]
301    fn large_data_multiple_blocks() {
302        let mut data = Vec::new();
303        for i in 0..1000 {
304            let line = format!(
305                "2024-01-15T10:30:{:02}.000Z INFO request_id={} status=200 duration_ms={}\n",
306                i % 60,
307                10000 + i,
308                i * 3 + 1
309            );
310            data.extend_from_slice(line.as_bytes());
311        }
312        let encoded = encode(&data).expect("encode");
313        assert_eq!(decode(&encoded).expect("decode"), data);
314        assert!(data.len() as f64 / encoded.len() as f64 > 2.0);
315    }
316
317    #[test]
318    fn random_access_block() {
319        let data: Vec<u8> = (0..20_000).map(|i| (i % 256) as u8).collect();
320        let block_size = 4096;
321        let encoded = encode_with_block_size(&data, block_size).expect("encode");
322        let block_count = Lz4Decoder::block_count(&encoded).expect("count");
323        assert_eq!(block_count, data.len().div_ceil(block_size));
324        let mut reassembled = Vec::new();
325        for index in 0..block_count {
326            reassembled.extend_from_slice(&decode_block(&encoded, index).expect("block"));
327        }
328        assert_eq!(reassembled, data);
329    }
330
331    #[test]
332    fn out_of_range_block_index() {
333        let encoded = encode(b"some data here").expect("encode");
334        assert!(decode_block(&encoded, 999).is_err());
335    }
336
337    #[test]
338    fn compressible_log_data_exceeds_three_to_one() {
339        let line = "2024-01-15 ERROR database connection timeout host=db-prod-01 retry=3\n";
340        let data = line.as_bytes().repeat(500);
341        let encoded = encode(&data).expect("encode");
342        assert_eq!(decode(&encoded).expect("decode"), data);
343        assert!(data.len() as f64 / encoded.len() as f64 > 3.0);
344    }
345
346    #[test]
347    fn incompressible_data_roundtrip() {
348        let mut data = vec![0u8; 10_000];
349        let mut rng: u64 = 9999;
350        for byte in &mut data {
351            rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1);
352            *byte = (rng >> 33) as u8;
353        }
354        let encoded = encode(&data).expect("encode");
355        assert_eq!(decode(&encoded).expect("decode"), data);
356    }
357
358    #[test]
359    fn streaming_encoder() {
360        let mut encoder = Lz4Encoder::new();
361        encoder.push(b"hello ").expect("push");
362        encoder.push(b"world").expect("push");
363        let encoded = encoder.finish().expect("finish");
364        assert_eq!(decode(&encoded).expect("decode"), b"hello world");
365    }
366
367    #[test]
368    fn custom_block_size() {
369        let data = vec![42u8; 10_000];
370        let encoded = encode_with_block_size(&data, 1024).expect("encode");
371        assert_eq!(decode(&encoded).expect("decode"), data);
372        assert_eq!(Lz4Decoder::block_count(&encoded).expect("count"), 10);
373    }
374
375    #[test]
376    fn rejects_huge_count_and_output_before_allocation() {
377        let mut frame = vec![0; HEADER_SIZE];
378        frame[0..4].copy_from_slice(&u32::MAX.to_le_bytes());
379        frame[4..8].copy_from_slice(&64u32.to_le_bytes());
380        frame[8..12].copy_from_slice(&u32::MAX.to_le_bytes());
381        assert!(matches!(
382            decode(&frame),
383            Err(CodecError::ResourceLimit { .. })
384        ));
385        assert!(decoded_len(MAX_DECODED_BYTES + 1, "test").is_err());
386    }
387
388    #[test]
389    fn rejects_truncated_table_payload_and_oversized_block_output() {
390        assert!(decode(&[]).is_err());
391        assert!(decode(&[0; 8]).is_err());
392        let mut frame = vec![0; HEADER_SIZE];
393        frame[0..4].copy_from_slice(&64u32.to_le_bytes());
394        frame[4..8].copy_from_slice(&64u32.to_le_bytes());
395        frame[8..12].copy_from_slice(&1u32.to_le_bytes());
396        assert!(matches!(decode(&frame), Err(CodecError::Truncated { .. })));
397        frame.extend_from_slice(&4u32.to_le_bytes());
398        frame.extend_from_slice(&65u32.to_le_bytes());
399        assert!(matches!(decode(&frame), Err(CodecError::Corrupt { .. })));
400    }
401
402    #[test]
403    fn random_access_rejects_short_nonfinal_block() {
404        let encoded = encode_with_block_size(&[7; 128], 64).expect("encode");
405        let first_len = read_u32(&encoded, HEADER_SIZE, "first length").expect("length") as usize;
406        let second_start = HEADER_SIZE + 8 + first_len;
407        let empty = lz4_flex::compress_prepend_size(&[]);
408        let mut malformed = encoded[..HEADER_SIZE + 8].to_vec();
409        malformed[HEADER_SIZE..HEADER_SIZE + 4]
410            .copy_from_slice(&(empty.len() as u32).to_le_bytes());
411        malformed.extend_from_slice(&empty);
412        malformed.extend_from_slice(&encoded[second_start..]);
413        assert!(matches!(
414            decode_block(&malformed, 0),
415            Err(CodecError::Corrupt { .. })
416        ));
417    }
418
419    #[test]
420    fn encoder_constructor_and_stream_reject_oversize_before_allocation() {
421        assert!(matches!(
422            Lz4Encoder::with_block_size(usize::MAX),
423            Err(CodecError::ResourceLimit { .. })
424        ));
425        let mut encoder = Lz4Encoder {
426            buf: vec![0; MAX_DECODED_BYTES],
427            block_size: DEFAULT_BLOCK_SIZE,
428        };
429        assert!(matches!(
430            encoder.push(&[1]),
431            Err(CodecError::ResourceLimit { .. })
432        ));
433        assert_eq!(encoder.len(), MAX_DECODED_BYTES);
434    }
435}
nodedb_codec/lz4.rs

nodedb_codec/
lz4.rs
