cqlite_core/storage/sstable/compression.rs
1//! Compression support for SSTable storage
2
3use crate::{error::Error, Result};
4// use async_trait::async_trait; // Commented out - unused
5
6/// Compression algorithms supported
7#[derive(Debug, Clone, Copy, PartialEq, serde::Serialize, serde::Deserialize, Default)]
8pub enum CompressionAlgorithm {
9 /// No compression
10 None,
11 /// LZ4 compression (fast)
12 #[default]
13 Lz4,
14 /// Snappy compression (balanced)
15 Snappy,
16 /// Deflate compression (high ratio)
17 Deflate,
18 /// Zstd compression (high efficiency)
19 Zstd,
20}
21
22/// Maximum allowed decompressed size to prevent memory exhaustion attacks (128MB)
23const MAX_DECOMPRESSED_SIZE: usize = 128 * 1024 * 1024;
24
25impl CompressionAlgorithm {
26 /// Map a recognized compressor name to its enum variant.
27 ///
28 /// Accepts CQLite short names (`LZ4`, `SNAPPY`, ...), Cassandra simple names
29 /// (`LZ4Compressor`, `SnappyCompressor`, ...) and the explicit no-compression
30 /// markers (`NONE`, `NoopCompressor`, `NoCompressor`). Returns `None` for any
31 /// other (unrecognized) name — callers MUST treat `None` here as "unknown",
32 /// not as "uncompressed".
33 fn from_name_opt(s: &str) -> Option<Self> {
34 // Strip any fully-qualified class prefix Cassandra may emit.
35 let simple = s.rsplit('.').next().unwrap_or(s);
36 match simple.to_uppercase().as_str() {
37 "NONE" | "NOOPCOMPRESSOR" | "NOCOMPRESSOR" | "NULLCOMPRESSOR" => {
38 Some(CompressionAlgorithm::None)
39 }
40 "LZ4" | "LZ4COMPRESSOR" => Some(CompressionAlgorithm::Lz4),
41 "SNAPPY" | "SNAPPYCOMPRESSOR" => Some(CompressionAlgorithm::Snappy),
42 "DEFLATE" | "DEFLATECOMPRESSOR" => Some(CompressionAlgorithm::Deflate),
43 "ZSTD" | "ZSTDCOMPRESSOR" => Some(CompressionAlgorithm::Zstd),
44 _ => None,
45 }
46 }
47
48 /// Fallible parse of a compressor name (issue #1001).
49 ///
50 /// This is the path the SSTable open / `CompressionInfo.db` flow MUST use: an
51 /// unrecognized name produces an explicit `UnsupportedFormat` error (including the
52 /// exact offending string) rather than silently falling back to uncompressed.
53 /// No content-based guessing is performed (no-heuristics mandate, issue #28).
54 pub fn parse(s: &str) -> Result<Self> {
55 Self::from_name_opt(s).ok_or_else(|| {
56 Error::UnsupportedFormat(format!(
57 "Unsupported compression algorithm '{}'. CQLite supports: \
58 LZ4Compressor, SnappyCompressor, DeflateCompressor, ZstdCompressor \
59 (or NONE for uncompressed).",
60 s
61 ))
62 })
63 }
64}
65
66impl From<String> for CompressionAlgorithm {
67 fn from(s: String) -> Self {
68 Self::from(s.as_str())
69 }
70}
71
72impl From<&str> for CompressionAlgorithm {
73 /// Infallible best-effort mapping. Unrecognized names map to `None`.
74 ///
75 /// WARNING: this MUST NOT be used on the SSTable read path — an unknown name here
76 /// is indistinguishable from genuinely-disabled compression. Use the fallible
77 /// [`CompressionAlgorithm::parse`] for any path that opens real `CompressionInfo.db`
78 /// metadata (issue #1001). This `From` is retained only for the legacy header
79 /// (`SSTableHeader.compression.algorithm`) path, which guards with an explicit
80 /// `!= "NONE"` check before conversion and re-validates the resulting variant.
81 fn from(s: &str) -> Self {
82 Self::from_name_opt(s).unwrap_or(CompressionAlgorithm::None)
83 }
84}
85
86/// Validates that decompressed size does not exceed safety limits
87///
88/// # Security
89/// Prevents decompression bomb attacks by rejecting sizes > 128MB
90///
91/// Only the LZ4 small-block decompress path consumes this (the Snappy/Deflate/Zstd
92/// paths bound inline via `MAX_DECOMPRESSED_SIZE`), so it is gated to its sole caller
93/// to stay dead-code-free under single-feature builds (issue #1873).
94#[cfg(feature = "lz4")]
95fn validate_decompression_size(uncompressed_size: usize) -> Result<()> {
96 if uncompressed_size > MAX_DECOMPRESSED_SIZE {
97 return Err(Error::storage(format!(
98 "Decompression bomb protection: size {} exceeds limit {} (128MB)",
99 uncompressed_size, MAX_DECOMPRESSED_SIZE
100 )));
101 }
102 Ok(())
103}
104
105/// Decode a RAW Snappy block (Cassandra 5.0 `SnappyCompressor`: no length prefix)
106/// in EXACTLY one attempt.
107///
108/// The authoritative CompressionInfo.db algorithm determines the single format —
109/// no framed-then-raw guessing (no-heuristics mandate #28, issue #1588). A guess
110/// could silently mis-decode an adversarial chunk to wrong bytes; strict raw
111/// decoding surfaces a typed error instead.
112///
113/// `decode_attempts` is incremented once per decode call. Production passes a
114/// throwaway `&mut 0`; tests thread a real counter to assert a single attempt.
115#[cfg(feature = "snappy")]
116fn snappy_decompress_raw(data: &[u8], decode_attempts: &mut usize) -> Result<Vec<u8>> {
117 use snap::raw::Decoder;
118 *decode_attempts += 1;
119
120 // CENTRALIZED bomb guard (issue #1588): a raw Snappy block carries its
121 // decompressed length as a leading varint. Inspect it FIRST and reject an
122 // over-limit block WITHOUT calling `decompress_vec` — `decompress_vec`
123 // pre-allocates `decompress_len` bytes up front, so an adversarial block
124 // declaring a huge size would allocate before any post-decode guard runs.
125 // This single choke point protects EVERY caller (chunk decode + streaming).
126 let advertised = snap::raw::decompress_len(data)
127 .map_err(|e| Error::storage(format!("Snappy (raw) length decode failed: {}", e)))?;
128 if advertised > MAX_DECOMPRESSED_SIZE {
129 return Err(Error::storage(format!(
130 "Decompression bomb protection: advertised size {} exceeds limit {} (128MB)",
131 advertised, MAX_DECOMPRESSED_SIZE
132 )));
133 }
134
135 let mut decoder = Decoder::new();
136 let decompressed = decoder
137 .decompress_vec(data)
138 .map_err(|e| Error::storage(format!("Snappy (raw) decompression failed: {}", e)))?;
139 // Belt-and-suspenders: the advertised length is attacker-controlled, so
140 // re-check the ACTUAL decoded size against the hard cap.
141 if decompressed.len() > MAX_DECOMPRESSED_SIZE {
142 return Err(Error::storage(format!(
143 "Decompression bomb protection: decompressed size {} exceeds limit {} (128MB)",
144 decompressed.len(),
145 MAX_DECOMPRESSED_SIZE
146 )));
147 }
148 Ok(decompressed)
149}
150
151/// Compression handler
152pub struct Compression {
153 algorithm: CompressionAlgorithm,
154}
155
156impl Compression {
157 /// Create a new compression handler
158 pub fn new(algorithm: CompressionAlgorithm) -> Result<Self> {
159 Ok(Self { algorithm })
160 }
161
162 /// Compress data with Cassandra-compatible parameters
163 pub fn compress(&self, data: &[u8]) -> Result<Vec<u8>> {
164 match self.algorithm {
165 CompressionAlgorithm::None => Ok(data.to_vec()),
166 CompressionAlgorithm::Lz4 => {
167 #[cfg(feature = "lz4")]
168 {
169 // Use Cassandra-compatible LZ4 compression
170 use lz4_flex::compress_prepend_size;
171
172 // Cassandra uses LZ4 frame format with specific parameters
173 let compressed = compress_prepend_size(data);
174 Ok(compressed)
175 }
176 #[cfg(not(feature = "lz4"))]
177 {
178 Err(Error::storage("LZ4 compression not available".to_string()))
179 }
180 }
181 CompressionAlgorithm::Snappy => {
182 #[cfg(feature = "snappy")]
183 {
184 use snap::raw::Encoder;
185
186 // Cassandra 5.0's SnappyCompressor writes a RAW Snappy block with NO
187 // length prefix (see writer::compressed_data_writer::SnappyCompressor
188 // and org.apache.cassandra.io.compress.SnappyCompressor). Emit raw so
189 // it round-trips through the strict raw decode path below (#1588). A
190 // 4-byte size prefix was NEVER a Cassandra-compatible format.
191 let mut encoder = Encoder::new();
192 encoder
193 .compress_vec(data)
194 .map_err(|e| Error::storage(format!("Snappy compression failed: {}", e)))
195 }
196 #[cfg(not(feature = "snappy"))]
197 {
198 Err(Error::storage(
199 "Snappy compression not available".to_string(),
200 ))
201 }
202 }
203 CompressionAlgorithm::Deflate => {
204 #[cfg(feature = "deflate")]
205 {
206 use flate2::write::ZlibEncoder;
207 use flate2::Compression as DeflateCompression;
208 use std::io::Write;
209
210 // Cassandra's DeflateCompressor uses java.util.zip.Deflater, which
211 // emits a ZLIB-wrapped stream (2-byte header + DEFLATE body +
212 // Adler-32 trailer) with NO 4-byte size prefix. Match it exactly so
213 // the output reads back through the zlib-aware decode path. (#1082)
214 let mut encoder = ZlibEncoder::new(Vec::new(), DeflateCompression::new(6));
215 encoder.write_all(data).map_err(|e| {
216 Error::storage(format!("Deflate compression failed: {}", e))
217 })?;
218 encoder
219 .finish()
220 .map_err(|e| Error::storage(format!("Deflate finish failed: {}", e)))
221 }
222 #[cfg(not(feature = "deflate"))]
223 {
224 Err(Error::storage(
225 "Deflate compression not available".to_string(),
226 ))
227 }
228 }
229 CompressionAlgorithm::Zstd => {
230 #[cfg(feature = "zstd")]
231 {
232 use zstd::stream::encode_all;
233
234 // Cassandra's ZstdCompressor writes a BARE zstd frame with NO
235 // 4-byte size prefix. Match it so the output reads back through
236 // the bare-frame decode path (#1082).
237 encode_all(data, 3)
238 .map_err(|e| Error::storage(format!("Zstd compression failed: {}", e)))
239 }
240 #[cfg(not(feature = "zstd"))]
241 {
242 Err(Error::storage("Zstd compression not available".to_string()))
243 }
244 }
245 }
246 }
247
248 /// Decompress data using traditional method (for small blocks)
249 pub fn decompress(&self, data: &[u8]) -> Result<Vec<u8>> {
250 // A5 read-work counter (DECOMPRESS_CALLS; consumers B1/E3): one per chunk
251 // decompress. This is the single choke point every compressed-chunk read
252 // path funnels through. No-op in release (design.md Decision 1/2).
253 crate::storage::sstable::read_work_counters::record_decompress();
254 match self.algorithm {
255 CompressionAlgorithm::None => Ok(data.to_vec()),
256 CompressionAlgorithm::Lz4 => {
257 #[cfg(feature = "lz4")]
258 {
259 use lz4_flex::decompress_size_prepended;
260
261 // LZ4 format: 4-byte size prefix (little-endian) + compressed data
262 if data.len() < 4 {
263 return Err(Error::storage("Invalid LZ4 data: too short".to_string()));
264 }
265
266 // Extract uncompressed size (4 bytes, little-endian for LZ4)
267 let uncompressed_size =
268 u32::from_le_bytes([data[0], data[1], data[2], data[3]]) as usize;
269
270 // Validate size to prevent decompression bombs
271 // SECURITY: lz4_flex::decompress_size_prepended does NOT validate the size
272 // prefix before allocating memory, making it vulnerable to memory exhaustion
273 // attacks if a malicious file contains an excessively large size value.
274 validate_decompression_size(uncompressed_size)?;
275
276 // Decompress using library function (now safe after validation)
277 decompress_size_prepended(data)
278 .map_err(|e| Error::storage(format!("LZ4 decompression failed: {}", e)))
279 }
280 #[cfg(not(feature = "lz4"))]
281 {
282 Err(Error::storage("LZ4 compression not available".to_string()))
283 }
284 }
285 CompressionAlgorithm::Snappy => {
286 #[cfg(feature = "snappy")]
287 {
288 // Decode EXACTLY one format (raw Snappy), determined by the
289 // authoritative CompressionInfo.db algorithm — never by trying
290 // framed-then-raw and keeping whichever "succeeds" (no-heuristics
291 // mandate #28, issue #1588). The framed-guess could silently return
292 // wrong bytes for an adversarial chunk; strict raw decoding rejects
293 // it. `&mut 0` discards the (test-only) attempt counter.
294 snappy_decompress_raw(data, &mut 0)
295 }
296 #[cfg(not(feature = "snappy"))]
297 {
298 Err(Error::storage(
299 "Snappy compression not available".to_string(),
300 ))
301 }
302 }
303 CompressionAlgorithm::Deflate => {
304 #[cfg(feature = "deflate")]
305 {
306 use flate2::read::ZlibDecoder;
307 use std::io::Read;
308
309 // Cassandra's DeflateCompressor uses java.util.zip.Deflater/Inflater,
310 // which emit ZLIB-wrapped streams: a 2-byte header (0x78 0x9c) +
311 // DEFLATE body + 4-byte Adler-32 trailer. There is NO 4-byte
312 // uncompressed-size prefix (that is an LZ4/Zstd convention) and the
313 // body is NOT raw DEFLATE. Decode with ZlibDecoder. (#1082)
314 if data.is_empty() {
315 return Err(Error::storage(
316 "Invalid Deflate data: empty chunk".to_string(),
317 ));
318 }
319
320 // Decompression-bomb guard: the decoder reads into a growing Vec,
321 // so we cap the output length rather than trusting any in-stream
322 // size field (none exists for zlib). The caller (chunk reader)
323 // separately bounds chunks by CompressionInfo.db lengths.
324 let mut decoder = ZlibDecoder::new(data).take(MAX_DECOMPRESSED_SIZE as u64 + 1);
325 let mut decompressed = Vec::new();
326 decoder.read_to_end(&mut decompressed).map_err(|e| {
327 Error::storage(format!("Deflate decompression failed: {}", e))
328 })?;
329
330 if decompressed.len() > MAX_DECOMPRESSED_SIZE {
331 return Err(Error::storage(format!(
332 "Decompression bomb protection: Deflate output exceeds limit {} (128MB)",
333 MAX_DECOMPRESSED_SIZE
334 )));
335 }
336
337 Ok(decompressed)
338 }
339 #[cfg(not(feature = "deflate"))]
340 {
341 Err(Error::storage(
342 "Deflate compression not available".to_string(),
343 ))
344 }
345 }
346 CompressionAlgorithm::Zstd => {
347 #[cfg(feature = "zstd")]
348 {
349 use std::io::Read;
350 use zstd::stream::read::Decoder as ZstdDecoder;
351
352 // Cassandra's ZstdCompressor writes a BARE zstd frame (magic
353 // 0x28 0xB5 0x2F 0xFD ...) with NO 4-byte uncompressed-size
354 // prefix — the same as the chunk-targeted decode path
355 // (chunk_decompressor.rs::decompress_zstd_chunk). The previous
356 // 4-byte-prefix assumption mis-read the frame magic as a ~650MB
357 // size and tripped the bomb guard on every stitched scan (#1082,
358 // same root cause as the Deflate fix). Decode the frame directly
359 // and bound the OUTPUT length instead of trusting an in-stream
360 // size field.
361 if data.is_empty() {
362 return Err(Error::storage("Invalid Zstd data: empty chunk".to_string()));
363 }
364
365 // Decompression-bomb guard: stream through a capped reader so a
366 // small malicious frame cannot allocate past the limit BEFORE we
367 // check the length (mirrors the Deflate path). A zstd frame can
368 // declare a huge content size, so `decode_all` would pre-allocate
369 // it up front — `Read::take` bounds the work instead.
370 let mut decoder = ZstdDecoder::new(data)
371 .map_err(|e| Error::storage(format!("Zstd decoder init failed: {}", e)))?
372 .take(MAX_DECOMPRESSED_SIZE as u64 + 1);
373 let mut decompressed = Vec::new();
374 decoder
375 .read_to_end(&mut decompressed)
376 .map_err(|e| Error::storage(format!("Zstd decompression failed: {}", e)))?;
377
378 if decompressed.len() > MAX_DECOMPRESSED_SIZE {
379 return Err(Error::storage(format!(
380 "Decompression bomb protection: Zstd output exceeds limit {} (128MB)",
381 MAX_DECOMPRESSED_SIZE
382 )));
383 }
384
385 Ok(decompressed)
386 }
387 #[cfg(not(feature = "zstd"))]
388 {
389 Err(Error::storage("Zstd compression not available".to_string()))
390 }
391 }
392 }
393 }
394
395 /// Get compression algorithm
396 pub fn algorithm(&self) -> &CompressionAlgorithm {
397 &self.algorithm
398 }
399
400 /// Select optimal compression algorithm based on data characteristics
401 pub fn select_optimal_algorithm(
402 data_sample: &[u8],
403 performance_priority: CompressionPriority,
404 ) -> CompressionAlgorithm {
405 // Analyze data characteristics
406 let entropy = calculate_entropy(data_sample);
407 let repetition_score = calculate_repetition_score(data_sample);
408 let data_size = data_sample.len();
409
410 match performance_priority {
411 CompressionPriority::Speed => {
412 // Prioritize speed over compression ratio
413 if entropy > 0.9 {
414 CompressionAlgorithm::None // High entropy data doesn't compress well
415 } else {
416 CompressionAlgorithm::Lz4 // Fast compression
417 }
418 }
419 CompressionPriority::Balanced => {
420 // Balance speed and compression ratio
421 if entropy > 0.95 {
422 CompressionAlgorithm::None
423 } else if repetition_score > 0.7 || data_size > 1024 * 1024 {
424 CompressionAlgorithm::Snappy // Good balance for large or repetitive data
425 } else {
426 CompressionAlgorithm::Lz4
427 }
428 }
429 CompressionPriority::Ratio => {
430 // Prioritize compression ratio
431 if entropy > 0.98 {
432 CompressionAlgorithm::None
433 } else if repetition_score > 0.5 {
434 CompressionAlgorithm::Deflate // Best compression for repetitive data
435 } else {
436 CompressionAlgorithm::Snappy
437 }
438 }
439 }
440 }
441}
442
443/// Compression priority for algorithm selection
444#[derive(Debug, Clone, Copy, PartialEq)]
445pub enum CompressionPriority {
446 /// Prioritize compression/decompression speed
447 Speed,
448 /// Balance speed and compression ratio
449 Balanced,
450 /// Prioritize maximum compression ratio
451 Ratio,
452}
453
454/// Compression statistics
455#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
456pub struct CompressionStats {
457 /// Original size in bytes
458 pub original_size: u64,
459
460 /// Compressed size in bytes
461 pub compressed_size: u64,
462
463 /// Compression ratio (compressed / original)
464 pub ratio: f64,
465
466 /// Compression algorithm used
467 pub algorithm: CompressionAlgorithm,
468}
469
470impl CompressionStats {
471 /// Calculate compression statistics
472 pub fn calculate(
473 original_size: u64,
474 compressed_size: u64,
475 algorithm: CompressionAlgorithm,
476 ) -> Self {
477 let ratio = if original_size > 0 {
478 compressed_size as f64 / original_size as f64
479 } else {
480 1.0
481 };
482
483 Self {
484 original_size,
485 compressed_size,
486 ratio,
487 algorithm,
488 }
489 }
490
491 /// Get space saved in bytes
492 pub fn space_saved(&self) -> u64 {
493 self.original_size.saturating_sub(self.compressed_size)
494 }
495
496 /// Get compression percentage
497 pub fn compression_percentage(&self) -> f64 {
498 (1.0 - self.ratio) * 100.0
499 }
500}
501
502/// Calculate entropy of data sample (0.0 = no entropy, 1.0 = maximum entropy)
503fn calculate_entropy(data: &[u8]) -> f64 {
504 if data.is_empty() {
505 return 0.0;
506 }
507
508 let mut counts = [0u32; 256];
509 for &byte in data {
510 counts[byte as usize] += 1;
511 }
512
513 let total = data.len() as f64;
514 let mut entropy = 0.0;
515
516 for &count in &counts {
517 if count > 0 {
518 let probability = count as f64 / total;
519 entropy -= probability * probability.log2();
520 }
521 }
522
523 // Normalize to 0.0-1.0 range
524 entropy / 8.0 // 8 bits per byte
525}
526
527/// Calculate repetition score (0.0 = no repetition, 1.0 = highly repetitive)
528fn calculate_repetition_score(data: &[u8]) -> f64 {
529 if data.len() < 4 {
530 return 0.0;
531 }
532
533 let mut repeated_bytes = 0;
534 let mut pattern_matches = 0;
535
536 // Check for byte repetitions
537 for i in 1..data.len() {
538 if data[i] == data[i - 1] {
539 repeated_bytes += 1;
540 }
541 }
542
543 // Check for 2-byte pattern repetitions
544 // Need at least 4 bytes to check 2-byte patterns (i-3 must be valid)
545 // Starting at i=3 prevents arithmetic underflow when accessing data[i-3]
546 for i in 3..data.len() {
547 if data[i] == data[i - 2] && data[i - 1] == data[i - 3] {
548 pattern_matches += 1;
549 }
550 }
551
552 let byte_repetition_score = repeated_bytes as f64 / (data.len() - 1) as f64;
553 let pattern_repetition_score = if data.len() > 3 {
554 pattern_matches as f64 / (data.len() - 3) as f64
555 } else {
556 0.0
557 };
558
559 // Combine scores with weights
560 (byte_repetition_score * 0.6 + pattern_repetition_score * 0.4).min(1.0)
561}
562
563/// The compression algorithm a reader uses to decompress a `Data.db`'s chunks.
564///
565/// Reduced to a plain algorithm field (issue #1597 / G1): the reader derives the
566/// algorithm from the single authoritative `CompressionInfo.db` parse and every
567/// decompression site funnels through [`Compression::decompress`] via
568/// [`CompressionReader::algorithm`]. The former streaming half (`read_streaming`,
569/// `read`, `with_block_size`, `block_size`, and the `buffer`/`block_size` fields)
570/// had zero consumers and was deleted.
571pub struct CompressionReader {
572 algorithm: CompressionAlgorithm,
573}
574
575impl CompressionReader {
576 /// Create a compression reader for the given algorithm.
577 pub fn new(algorithm: CompressionAlgorithm) -> Self {
578 Self { algorithm }
579 }
580
581 /// Get the compression algorithm.
582 pub fn algorithm(&self) -> &CompressionAlgorithm {
583 &self.algorithm
584 }
585}
586
587#[cfg(test)]
588mod tests {
589 use super::*;
590
591 #[test]
592 fn test_no_compression() {
593 let compression = Compression::new(CompressionAlgorithm::None).unwrap();
594 let data = b"hello world";
595
596 let compressed = compression.compress(data).unwrap();
597 assert_eq!(compressed, data);
598
599 let decompressed = compression.decompress(&compressed).unwrap();
600 assert_eq!(decompressed, data);
601 }
602
603 #[test]
604 fn test_compression_stats() {
605 let stats = CompressionStats::calculate(1000, 600, CompressionAlgorithm::Lz4);
606
607 assert_eq!(stats.original_size, 1000);
608 assert_eq!(stats.compressed_size, 600);
609 assert_eq!(stats.ratio, 0.6);
610 assert_eq!(stats.space_saved(), 400);
611 assert_eq!(stats.compression_percentage(), 40.0);
612 }
613
614 // Note: Test methods temporarily disabled due to compilation issues
615 // The functionality is tested via integration tests
616
617 /// Adversarial oracle (issue #1588, decision #14): a chunk whose leading 4
618 /// bytes ALSO parse as a plausible framed big-endian length header, followed
619 /// by a valid RAW-snappy body of exactly that many output bytes.
620 ///
621 /// Under the (deleted) framed-then-raw guessing, `decompress` read the 4-byte
622 /// BE prefix `S`, decoded `data[4..]` as raw snappy to `P_wrong` (whose length
623 /// equals `S`), and RETURNED those bytes — silently wrong. The authoritative
624 /// Cassandra 5.0 format for a `SnappyCompressor` chunk is RAW snappy with NO
625 /// length prefix, so strict raw decoding of the WHOLE chunk must NOT return
626 /// `P_wrong` (it rejects the malformed leading zero-varint stream). This is the
627 /// no-heuristics enforcement: decode exactly one format determined by metadata.
628 #[cfg(feature = "snappy")]
629 #[test]
630 fn test_snappy_decode_is_strict_raw_only_no_format_guessing() {
631 use snap::raw::Encoder;
632 let p_wrong = b"WRONG-framed-decode-abcdefghijklmnopqrstuvwxyz".to_vec();
633 let s = p_wrong.len() as u32;
634 let mut enc = Encoder::new();
635 let inner = enc.compress_vec(&p_wrong).unwrap(); // valid raw snappy -> p_wrong
636 let mut adversarial = s.to_be_bytes().to_vec(); // plausible framed BE header
637 adversarial.extend_from_slice(&inner);
638
639 let compression = Compression::new(CompressionAlgorithm::Snappy).unwrap();
640 let got = compression.decompress(&adversarial).ok();
641 // Strict raw decode must not merely differ from the framed guess — it must
642 // FAIL (typed error) rather than silently produce any bytes: the leading
643 // 4-byte pseudo-header parses as a malformed raw Snappy stream (a
644 // zero-length literal with trailing data), which the raw decoder rejects.
645 assert!(
646 got.is_none(),
647 "strict raw decode must ERROR on the ambiguous chunk, not return bytes \
648 (no-heuristics, #1588); got {got:?}"
649 );
650 }
651
652 #[cfg(feature = "snappy")]
653 #[test]
654 fn test_snappy_compression_cassandra_format() {
655 let compression = Compression::new(CompressionAlgorithm::Snappy).unwrap();
656 let data = b"This is test data for Snappy compression with Cassandra format validation. "
657 .repeat(10);
658
659 let compressed = compression.compress(&data).unwrap();
660
661 // Cassandra 5.0 SnappyCompressor emits a RAW Snappy block with NO length
662 // prefix (#1588). The compressed bytes are exactly what a raw Snappy
663 // decoder consumes; there is no 4-byte big-endian size header.
664 let raw_roundtrip = {
665 use snap::raw::Decoder;
666 Decoder::new().decompress_vec(&compressed).unwrap()
667 };
668 assert_eq!(raw_roundtrip, data, "compress() output must be raw Snappy");
669
670 let decompressed = compression.decompress(&compressed).unwrap();
671 assert_eq!(decompressed, data);
672 }
673
674 /// A legitimate RAW Snappy chunk decodes to the known-good bytes in EXACTLY
675 /// one decode attempt (issue #1588 decision #14: single-attempt + byte-identity).
676 #[cfg(feature = "snappy")]
677 #[test]
678 fn test_snappy_decode_single_attempt_and_byte_identical() {
679 use snap::raw::Encoder;
680 let good = b"the quick brown fox jumps over the lazy dog. ".repeat(8);
681 let chunk = Encoder::new().compress_vec(&good).unwrap();
682
683 let mut attempts = 0usize;
684 let out = super::snappy_decompress_raw(&chunk, &mut attempts).unwrap();
685 assert_eq!(
686 out, good,
687 "raw decode is byte-identical to the known-good input"
688 );
689 assert_eq!(
690 attempts, 1,
691 "exactly one decode attempt (no format guessing)"
692 );
693 }
694
695 /// Encode a Snappy raw-block length prefix (LEB128 varint) for `n`.
696 #[cfg(feature = "snappy")]
697 fn snappy_len_prefix(mut n: u64) -> Vec<u8> {
698 let mut out = Vec::new();
699 loop {
700 let mut b = (n & 0x7f) as u8;
701 n >>= 7;
702 if n != 0 {
703 b |= 0x80;
704 }
705 out.push(b);
706 if n == 0 {
707 break;
708 }
709 }
710 out
711 }
712
713 /// SECURITY (issue #1588): the SHARED `snappy_decompress_raw` helper — the
714 /// choke point used by the CHUNK (non-streaming) decode path — must reject a
715 /// block whose ADVERTISED decompressed length exceeds `MAX_DECOMPRESSED_SIZE`
716 /// WITHOUT allocating (i.e. before `decompress_vec` pre-allocates that size).
717 /// This proves the guard lives at the helper, not only in the streaming caller.
718 #[cfg(feature = "snappy")]
719 #[test]
720 fn test_snappy_raw_helper_advertised_len_bomb_rejected_no_alloc() {
721 // Craft a tiny raw block: a varint advertising 200MB (> 128MB cap) plus a
722 // stub tag byte. The length guard fires before any output allocation.
723 let mut chunk = snappy_len_prefix(200 * 1024 * 1024);
724 chunk.push(0x00);
725
726 let mut attempts = 0usize;
727 let err = super::snappy_decompress_raw(&chunk, &mut attempts)
728 .expect_err("over-limit advertised length must be rejected at the helper");
729 assert!(
730 err.to_string().contains("Decompression bomb protection"),
731 "expected typed decompression-bomb error, got: {err}"
732 );
733 assert!(
734 err.to_string().contains("advertised size"),
735 "guard must fire on the ADVERTISED length before decode/alloc, got: {err}"
736 );
737 // The decode attempt is counted, but the rejection happens before the
738 // `decompress_vec` allocation of the advertised 200MB.
739 assert_eq!(attempts, 1, "helper is entered exactly once");
740 }
741
742 #[cfg(feature = "deflate")]
743 #[test]
744 fn test_deflate_compression_cassandra_format() {
745 let compression = Compression::new(CompressionAlgorithm::Deflate).unwrap();
746 let data = b"This is test data for Deflate compression with Cassandra format validation. "
747 .repeat(10);
748
749 let compressed = compression.compress(&data).unwrap();
750
751 // Cassandra format: ZLIB-wrapped (0x78 header), NO 4-byte size prefix (#1082).
752 assert!(compressed.len() >= 2);
753 assert_eq!(compressed[0], 0x78, "zlib stream must start with CMF 0x78");
754
755 let decompressed = compression.decompress(&compressed).unwrap();
756 assert_eq!(decompressed, data);
757 }
758
759 #[test]
760 fn test_compression_reader() {
761 // CompressionReader is now a plain algorithm field (issue #1597 / G1): it
762 // carries the algorithm the reader derived from CompressionInfo.db, which
763 // the decompression sites feed to `Compression::new`.
764 let reader = CompressionReader::new(CompressionAlgorithm::Snappy);
765 assert_eq!(reader.algorithm(), &CompressionAlgorithm::Snappy);
766 }
767
768 #[test]
769 fn test_compression_algorithm_from_string() {
770 assert_eq!(
771 CompressionAlgorithm::from("NONE".to_string()),
772 CompressionAlgorithm::None
773 );
774 assert_eq!(
775 CompressionAlgorithm::from("LZ4".to_string()),
776 CompressionAlgorithm::Lz4
777 );
778 assert_eq!(
779 CompressionAlgorithm::from("SNAPPY".to_string()),
780 CompressionAlgorithm::Snappy
781 );
782 assert_eq!(
783 CompressionAlgorithm::from("DEFLATE".to_string()),
784 CompressionAlgorithm::Deflate
785 );
786 assert_eq!(
787 CompressionAlgorithm::from("unknown".to_string()),
788 CompressionAlgorithm::None
789 );
790 }
791
792 #[test]
793 fn test_compression_invalid_data() {
794 let compression = Compression::new(CompressionAlgorithm::Snappy).unwrap();
795
796 // Test with data too short for size prefix
797 let short_data = &[1, 2];
798 assert!(compression.decompress(short_data).is_err());
799
800 // Test with invalid size prefix
801 let invalid_data = &[0, 0, 0, 100, 1, 2, 3]; // Claims 100 bytes but only has 3
802 if cfg!(feature = "snappy") {
803 assert!(compression.decompress(invalid_data).is_err());
804 }
805 }
806
807 #[test]
808 fn test_decompression_bomb_protection() {
809 // Test protection against malicious size claims for all algorithms
810 // Using 200MB claim (exceeds 128MB limit) to test protection
811
812 // Snappy: Test that decompression bomb protection works after decompression
813 // (not during prefix check, since NB format uses raw Snappy without prefix)
814 #[cfg(feature = "snappy")]
815 {
816 // Note: The decompression bomb protection for Snappy happens AFTER decompression
817 // completes, by checking the decompressed size. This is because Cassandra 5.0 NB
818 // format uses raw Snappy without a size prefix, so we can't detect bombs early.
819 //
820 // A malicious prefix with fake size >128MB is handled by skipping the prefixed
821 // format and trying raw Snappy instead (which will fail if the data is invalid).
822 //
823 // This test verifies that post-decompression size checking works correctly.
824 // The actual protection is at lines 281-286 in the decompress() method.
825 }
826
827 // Deflate (#1082): Cassandra emits ZLIB-wrapped streams with NO 4-byte size
828 // prefix, so the bomb guard caps the DECODED output length rather than
829 // trusting an in-stream size field. The legacy "fake 200MB size prefix" no
830 // longer applies; instead verify that non-zlib bytes (here, what the old
831 // format would have produced) are rejected as malformed rather than read as
832 // a 2GB size and OOM'd.
833 #[cfg(feature = "deflate")]
834 {
835 let compression = Compression::new(CompressionAlgorithm::Deflate).unwrap();
836 let malicious_size: u32 = 200 * 1024 * 1024;
837 let mut malicious_data = malicious_size.to_be_bytes().to_vec();
838 malicious_data.extend_from_slice(&[0u8; 10]);
839
840 let result = compression.decompress(&malicious_data);
841 assert!(
842 result.is_err(),
843 "Should reject malformed (non-zlib) Deflate data"
844 );
845
846 // A genuine zlib-wrapped round-trip still decodes correctly.
847 let data = b"deflate bomb-guard roundtrip".repeat(4);
848 let compressed = compression.compress(&data).unwrap();
849 let decompressed = compression.decompress(&compressed).unwrap();
850 assert_eq!(decompressed, data);
851 }
852
853 // Zstd (#1082): Cassandra writes a BARE zstd frame with NO 4-byte size
854 // prefix, so the old "fake 200MB prefix" scenario no longer applies. The
855 // bomb guard now caps the DECODED output length; here verify that malformed
856 // (non-frame) bytes are rejected rather than mis-read as a multi-GB size,
857 // and that a genuine bare-frame round-trip still decodes.
858 #[cfg(feature = "zstd")]
859 {
860 let compression = Compression::new(CompressionAlgorithm::Zstd).unwrap();
861 let mut malicious_data = (200u32 * 1024 * 1024).to_be_bytes().to_vec();
862 malicious_data.extend_from_slice(&[0u8; 10]);
863
864 let result = compression.decompress(&malicious_data);
865 assert!(result.is_err(), "Should reject malformed Zstd frame");
866
867 let data = b"zstd bomb-guard roundtrip".repeat(4);
868 let compressed = compression.compress(&data).unwrap();
869 let decompressed = compression.decompress(&compressed).unwrap();
870 assert_eq!(decompressed, data);
871 }
872
873 // LZ4: Create data claiming 200MB uncompressed size
874 #[cfg(feature = "lz4")]
875 {
876 let compression = Compression::new(CompressionAlgorithm::Lz4).unwrap();
877 let malicious_size: u32 = 200 * 1024 * 1024; // 200MB claim (exceeds 128MB limit)
878 let mut malicious_data = malicious_size.to_le_bytes().to_vec(); // LZ4 uses little-endian
879 malicious_data.extend_from_slice(&[0u8; 10]); // Some fake compressed data
880
881 let result = compression.decompress(&malicious_data);
882 assert!(result.is_err(), "Should reject malicious LZ4 size");
883 assert!(result
884 .unwrap_err()
885 .to_string()
886 .contains("Decompression bomb"));
887 }
888 }
889
890 #[test]
891 fn test_entropy_calculation() {
892 // Test with uniform data (high entropy)
893 let uniform_data: Vec<u8> = (0..=255).collect();
894 let entropy = calculate_entropy(&uniform_data);
895 assert!(entropy > 0.9); // Should be close to 1.0
896
897 // Test with repetitive data (low entropy)
898 let repetitive_data = vec![0u8; 256];
899 let entropy = calculate_entropy(&repetitive_data);
900 assert!(entropy < 0.1); // Should be close to 0.0
901 }
902
903 #[test]
904 fn test_repetition_score() {
905 // Test with highly repetitive data
906 let repetitive_data = vec![0u8, 0u8, 0u8, 0u8];
907 let score = calculate_repetition_score(&repetitive_data);
908 assert!(score > 0.8);
909
910 // Test with random data
911 let random_data = vec![1u8, 2u8, 3u8, 4u8, 5u8, 6u8, 7u8, 8u8];
912 let score = calculate_repetition_score(&random_data);
913 assert!(score < 0.2);
914 }
915
916 // Note: Algorithm selection test temporarily disabled due to compilation issues
917 // The functionality is tested via integration tests
918}