Skip to main content

pjson_rs/compression/
secure.rs

1//! Secure compression with bomb protection and real byte-level codecs.
2//!
3//! This module provides [`SecureCompressor`], which applies byte-level compression (Layer B)
4//! to arbitrary `&[u8]` payloads. It is distinct from `SchemaCompressor` in `compression/mod.rs`,
5//! which operates on `serde_json::Value` (Layer A / structural compression).
6//!
7//! # Security
8//!
9//! Every decompression is routed through `CompressionBombProtector`, which streams the decoder
10//! output and aborts if decompressed size or ratio exceeds configured limits.
11//!
12//! # In-process only
13//!
14//! [`SecureCompressedData`] carries the codec tag and is intended for in-process use only.
15//! It is not a wire format. If cross-process transport is needed in a future PR, a versioned
16//! framing header must be designed separately.
17
18use crate::{
19    Error, Result,
20    security::{CompressionBombDetector, CompressionStats},
21};
22#[cfg(feature = "compression")]
23use std::io::Write;
24use std::io::{Cursor, Read};
25use tracing::{debug, info, warn};
26#[cfg(all(feature = "compression", not(target_arch = "wasm32")))]
27use zstd;
28
29/// Byte-level compression algorithms used by [`SecureCompressor`].
30///
31/// This is distinct from [`CompressionStrategy`](super::CompressionStrategy), which operates on
32/// `serde_json::Value` (Layer A). `ByteCodec` operates on raw bytes after JSON serialization
33/// (Layer B).
34///
35/// Codecs other than `None` require the `compression` feature.
36///
37/// # Breaking change (pre-1.0)
38///
39/// `Copy` was removed from this enum when `ZstdDict` was added (it carries an
40/// `Arc<ZstdDictionary>`).  Code that relied on implicit copy can use `.clone()`
41/// (one atomic refcount bump for `ZstdDict`; a no-op for the other variants).
42#[derive(Debug, Clone, PartialEq, Eq, Default)]
43pub enum ByteCodec {
44    /// No compression — bytes stored verbatim. Always available.
45    #[default]
46    None,
47    /// Raw deflate (RFC 1951). Low framing overhead.
48    ///
49    /// Note: raw deflate has no magic header, so codec mismatch during decompression will
50    /// produce a decoder error rather than a guaranteed clean failure. The codec tag embedded
51    /// in [`SecureCompressedData`] prevents this for in-process round-trips.
52    ///
53    /// Requires `feature = "compression"`.
54    Deflate,
55    /// Gzip (RFC 1952). Self-identifying via `1f 8b` magic header.
56    ///
57    /// Requires `feature = "compression"`.
58    Gzip,
59    /// Brotli. Best ratio for repetitive JSON.
60    ///
61    /// Requires `feature = "compression"`.
62    Brotli,
63    /// Trained zstd dictionary compression.
64    ///
65    /// A single `Arc<ZstdDictionary>` is the canonical sharing primitive. The inner
66    /// `Vec<u8>` inside [`crate::compression::zstd::ZstdDictionary`] is **not**
67    /// `Arc`-wrapped — sharing happens exactly once at this enum level (avoids
68    /// double indirection). Cloning this variant performs one atomic refcount
69    /// increment and no allocation.
70    ///
71    /// Equality compares the underlying bytes via `Arc<T>: PartialEq where T: PartialEq`.
72    /// When both sides share the same `Arc` allocation, `Arc::ptr_eq` provides a fast path.
73    ///
74    /// Requires `feature = "compression"` on a non-`wasm32` target.
75    #[cfg(all(feature = "compression", not(target_arch = "wasm32")))]
76    ZstdDict(std::sync::Arc<crate::compression::zstd::ZstdDictionary>),
77}
78
79/// Quality knob for byte-level codecs.
80///
81/// Maps to codec-specific levels: deflate 1/6/9 and brotli quality 1/5/11.
82#[derive(Debug, Clone, Copy, Default)]
83pub enum CompressionQuality {
84    /// Speed-optimised: deflate level 1, brotli quality 1.
85    Fast,
86    /// Balanced speed/ratio (default): deflate level 6, brotli quality 5.
87    #[default]
88    Balanced,
89    /// Maximum ratio: deflate level 9, brotli quality 11.
90    Best,
91}
92
93impl CompressionQuality {
94    #[cfg(feature = "compression")]
95    fn flate2_level(self) -> flate2::Compression {
96        match self {
97            Self::Fast => flate2::Compression::fast(),
98            Self::Balanced => flate2::Compression::default(),
99            Self::Best => flate2::Compression::best(),
100        }
101    }
102
103    #[cfg(feature = "compression")]
104    fn brotli_quality(self) -> i32 {
105        match self {
106            Self::Fast => 1,
107            Self::Balanced => 5,
108            Self::Best => 11,
109        }
110    }
111}
112
113/// Compressed bytes with security metadata and codec identification.
114///
115/// # In-process only
116///
117/// This struct is intended for in-process use only and is not a wire format.
118/// The `codec` field is carried alongside `data` so that [`SecureCompressor::decompress_protected`]
119/// always uses the correct decoder. If this type must cross process boundaries in the future,
120/// design a versioned framing header as a separate concern.
121#[derive(Debug, Clone)]
122pub struct SecureCompressedData {
123    /// The compressed (or verbatim) payload.
124    pub data: Vec<u8>,
125    /// Original uncompressed size in bytes.
126    pub original_size: usize,
127    /// Compression ratio: `original_size / compressed_size`.
128    ///
129    /// A value of `2.0` means the compressed payload is half the original size (50% size reduction).
130    /// For `ByteCodec::None` this is always `1.0`; for incompressible data it can be `< 1.0`
131    /// because most codecs add a small framing header.
132    pub compression_ratio: f64,
133    /// Codec used to produce `data`. Must be passed back to [`SecureCompressor::decompress_protected`].
134    pub codec: ByteCodec,
135}
136
137/// Secure byte-level compressor with integrated bomb protection.
138///
139/// Wraps a [`CompressionBombDetector`] to ensure decompressed output never exceeds configured
140/// size and ratio limits, regardless of which codec is active.
141///
142/// # Examples
143///
144/// ```rust
145/// use pjson_rs::compression::secure::{SecureCompressor, ByteCodec};
146///
147/// let compressor = SecureCompressor::with_default_security(ByteCodec::None);
148/// let compressed = compressor.compress(b"hello world").unwrap();
149/// let decompressed = compressor.decompress_protected(&compressed).unwrap();
150/// assert_eq!(decompressed, b"hello world");
151/// ```
152pub struct SecureCompressor {
153    detector: CompressionBombDetector,
154    codec: ByteCodec,
155    #[cfg_attr(not(feature = "compression"), allow(dead_code))]
156    quality: CompressionQuality,
157}
158
159impl SecureCompressor {
160    /// Create a new secure compressor with the given detector and codec.
161    pub fn new(detector: CompressionBombDetector, codec: ByteCodec) -> Self {
162        Self {
163            detector,
164            codec,
165            quality: CompressionQuality::default(),
166        }
167    }
168
169    /// Create with default security settings and the given codec.
170    pub fn with_default_security(codec: ByteCodec) -> Self {
171        Self::new(CompressionBombDetector::default(), codec)
172    }
173
174    /// Create with explicit quality setting.
175    pub fn with_quality(
176        detector: CompressionBombDetector,
177        codec: ByteCodec,
178        quality: CompressionQuality,
179    ) -> Self {
180        Self {
181            detector,
182            codec,
183            quality,
184        }
185    }
186
187    /// Compress `data` using the configured codec.
188    ///
189    /// Validates the input size against `max_compressed_size` before encoding.
190    pub fn compress(&self, data: &[u8]) -> Result<SecureCompressedData> {
191        self.detector.validate_pre_decompression(data.len())?;
192
193        let compressed_bytes = self.encode(data)?;
194
195        let compression_ratio = data.len() as f64 / compressed_bytes.len().max(1) as f64;
196        info!("Compression completed: {:.2}x ratio", compression_ratio);
197
198        Ok(SecureCompressedData {
199            original_size: data.len(),
200            compression_ratio,
201            codec: self.codec.clone(),
202            data: compressed_bytes,
203        })
204    }
205
206    /// Decompress `compressed` using the codec recorded in `compressed.codec`.
207    ///
208    /// Decoder output is streamed through `CompressionBombProtector` — decompression aborts
209    /// early if size or ratio limits are exceeded.
210    pub fn decompress_protected(&self, compressed: &SecureCompressedData) -> Result<Vec<u8>> {
211        self.detector
212            .validate_pre_decompression(compressed.data.len())?;
213        self.decode_with_protection(&compressed.data, compressed.codec.clone(), None)
214    }
215
216    /// Decompress nested/chained compression with depth tracking.
217    ///
218    /// Equivalent to [`decompress_protected`](Self::decompress_protected) but additionally enforces
219    /// `max_compression_depth` via [`CompressionBombDetector::protect_nested_reader`].
220    pub fn decompress_nested(
221        &self,
222        compressed: &SecureCompressedData,
223        depth: usize,
224    ) -> Result<Vec<u8>> {
225        self.detector
226            .validate_pre_decompression(compressed.data.len())?;
227        self.decode_with_protection(&compressed.data, compressed.codec.clone(), Some(depth))
228    }
229
230    /// Encode `data` with the configured codec. Returns compressed bytes only.
231    fn encode(&self, data: &[u8]) -> Result<Vec<u8>> {
232        match &self.codec {
233            ByteCodec::None => {
234                debug!("No compression applied");
235                Ok(data.to_vec())
236            }
237
238            #[cfg(feature = "compression")]
239            ByteCodec::Deflate => {
240                use flate2::write::DeflateEncoder;
241                let mut enc = DeflateEncoder::new(Vec::new(), self.quality.flate2_level());
242                enc.write_all(data)
243                    .map_err(|e| Error::CompressionError(format!("deflate encode: {e}")))?;
244                enc.finish()
245                    .map_err(|e| Error::CompressionError(format!("deflate finish: {e}")))
246            }
247
248            #[cfg(feature = "compression")]
249            ByteCodec::Gzip => {
250                use flate2::write::GzEncoder;
251                let mut enc = GzEncoder::new(Vec::new(), self.quality.flate2_level());
252                enc.write_all(data)
253                    .map_err(|e| Error::CompressionError(format!("gzip encode: {e}")))?;
254                enc.finish()
255                    .map_err(|e| Error::CompressionError(format!("gzip finish: {e}")))
256            }
257
258            #[cfg(feature = "compression")]
259            ByteCodec::Brotli => {
260                let params = brotli::enc::BrotliEncoderParams {
261                    quality: self.quality.brotli_quality(),
262                    ..Default::default()
263                };
264                let mut out = Vec::new();
265                brotli::BrotliCompress(&mut Cursor::new(data), &mut out, &params)
266                    .map_err(|e| Error::CompressionError(format!("brotli encode: {e}")))?;
267                Ok(out)
268            }
269
270            #[cfg(all(feature = "compression", not(target_arch = "wasm32")))]
271            ByteCodec::ZstdDict(dict) => {
272                crate::compression::zstd::ZstdDictCompressor::compress(data, dict.as_ref())
273            }
274
275            #[cfg(not(feature = "compression"))]
276            ByteCodec::Deflate | ByteCodec::Gzip | ByteCodec::Brotli => Err(
277                Error::CompressionError("feature `compression` is not enabled".into()),
278            ),
279        }
280    }
281
282    /// Decode `data` through a bomb-protected reader.
283    ///
284    /// `depth` is `Some(n)` for nested decompression (depth-limited) or `None` for a flat call.
285    fn decode_with_protection(
286        &self,
287        data: &[u8],
288        codec: ByteCodec,
289        depth: Option<usize>,
290    ) -> Result<Vec<u8>> {
291        // Macro-free helper: executes the read loop with any `impl Read` decoder.
292        // Avoids boxing across a lifetime boundary by keeping decoder + protector in one scope.
293        macro_rules! run {
294            ($decoder:expr) => {{
295                let compressed_size = data.len();
296                let mut out = Vec::new();
297                let result = if let Some(d) = depth {
298                    let mut protector =
299                        self.detector
300                            .protect_nested_reader($decoder, compressed_size, d)?;
301                    let r = protector.read_to_end(&mut out);
302                    let stats = protector.stats();
303                    self.log_decompression_stats(&stats);
304                    if stats.compression_depth > 0 {
305                        warn!(
306                            "Nested decompression detected at depth {}",
307                            stats.compression_depth
308                        );
309                    }
310                    r
311                } else {
312                    let mut protector = self.detector.protect_reader($decoder, compressed_size);
313                    let r = protector.read_to_end(&mut out);
314                    let stats = protector.stats();
315                    self.log_decompression_stats(&stats);
316                    r
317                };
318                match result {
319                    Ok(_) => {
320                        self.detector.validate_result(compressed_size, out.len())?;
321                        Ok(out)
322                    }
323                    Err(e) => {
324                        warn!("Decompression failed: {}", e);
325                        Err(Error::SecurityError(format!(
326                            "Protected decompression failed: {}",
327                            e
328                        )))
329                    }
330                }
331            }};
332        }
333
334        match codec {
335            ByteCodec::None => run!(Cursor::new(data)),
336
337            #[cfg(feature = "compression")]
338            ByteCodec::Deflate => run!(flate2::read::DeflateDecoder::new(Cursor::new(data))),
339
340            #[cfg(feature = "compression")]
341            ByteCodec::Gzip => run!(flate2::read::GzDecoder::new(Cursor::new(data))),
342
343            #[cfg(feature = "compression")]
344            ByteCodec::Brotli => run!(brotli::Decompressor::new(Cursor::new(data), 4096)),
345
346            // ZstdDict uses the streaming decoder so every decompressed byte
347            // passes through the CompressionBombProtector's read loop (run!).
348            // Bulk `zstd::bulk::Decompressor::decompress` is intentionally
349            // avoided here — it would bypass the byte-level output cap.
350            #[cfg(all(feature = "compression", not(target_arch = "wasm32")))]
351            ByteCodec::ZstdDict(dict) => {
352                let decoder = zstd::stream::read::Decoder::with_dictionary(
353                    Cursor::new(data),
354                    dict.as_bytes(),
355                )
356                .map_err(|e| Error::CompressionError(format!("zstd decoder init: {e}")))?;
357                run!(decoder)
358            }
359
360            #[cfg(not(feature = "compression"))]
361            ByteCodec::Deflate | ByteCodec::Gzip | ByteCodec::Brotli => Err(
362                Error::CompressionError("feature `compression` is not enabled".into()),
363            ),
364        }
365    }
366
367    fn log_decompression_stats(&self, stats: &CompressionStats) {
368        info!(
369            "Decompression stats: {}B -> {}B (ratio: {:.2}x, depth: {})",
370            stats.compressed_size, stats.decompressed_size, stats.ratio, stats.compression_depth
371        );
372    }
373}
374
375/// Secure decompression context for streaming operations.
376pub struct SecureDecompressionContext {
377    detector: CompressionBombDetector,
378    current_depth: usize,
379    max_concurrent_streams: usize,
380    active_streams: usize,
381}
382
383impl SecureDecompressionContext {
384    /// Create new secure decompression context.
385    pub fn new(detector: CompressionBombDetector, max_concurrent_streams: usize) -> Self {
386        Self {
387            detector,
388            current_depth: 0,
389            max_concurrent_streams,
390            active_streams: 0,
391        }
392    }
393
394    /// Start a new protected decompression stream.
395    ///
396    /// Returns an error if the concurrent stream limit would be exceeded.
397    ///
398    /// # Note
399    ///
400    /// The returned `CompressionBombProtector` wraps an empty in-memory cursor. Callers are
401    /// responsible for writing compressed bytes into the underlying buffer before reading. This API
402    /// is a concurrency-limit scaffold; true streaming wire integration is left for a future PR.
403    pub fn start_stream(
404        &mut self,
405        compressed_size: usize,
406    ) -> Result<crate::security::CompressionBombProtector<Cursor<Vec<u8>>>> {
407        if self.active_streams >= self.max_concurrent_streams {
408            return Err(Error::SecurityError(format!(
409                "Too many concurrent decompression streams: {}/{}",
410                self.active_streams, self.max_concurrent_streams
411            )));
412        }
413
414        let cursor = Cursor::new(Vec::new());
415        let protector =
416            self.detector
417                .protect_nested_reader(cursor, compressed_size, self.current_depth)?;
418
419        self.active_streams += 1;
420        info!(
421            "Started secure decompression stream (active: {})",
422            self.active_streams
423        );
424
425        Ok(protector)
426    }
427
428    /// Finish a decompression stream and decrement the active count.
429    pub fn finish_stream(&mut self) {
430        if self.active_streams > 0 {
431            self.active_streams -= 1;
432            info!(
433                "Finished secure decompression stream (active: {})",
434                self.active_streams
435            );
436        }
437    }
438
439    /// Get current context statistics.
440    pub fn stats(&self) -> DecompressionContextStats {
441        DecompressionContextStats {
442            current_depth: self.current_depth,
443            active_streams: self.active_streams,
444            max_concurrent_streams: self.max_concurrent_streams,
445        }
446    }
447}
448
449/// Statistics for a [`SecureDecompressionContext`].
450#[derive(Debug, Clone)]
451pub struct DecompressionContextStats {
452    /// Current nested decompression depth.
453    pub current_depth: usize,
454    /// Number of decompression streams currently in flight.
455    pub active_streams: usize,
456    /// Configured maximum number of concurrent streams.
457    pub max_concurrent_streams: usize,
458}
459
460#[cfg(test)]
461mod tests {
462    use super::*;
463    use crate::security::CompressionBombConfig;
464
465    #[test]
466    fn test_secure_compressor_creation() {
467        let detector = CompressionBombDetector::default();
468        let compressor = SecureCompressor::new(detector, ByteCodec::None);
469        // Verify the compressor is created (not null pointer).
470        assert!(!std::ptr::addr_of!(compressor).cast::<u8>().is_null());
471    }
472
473    #[test]
474    fn test_secure_compression_none() {
475        let compressor = SecureCompressor::with_default_security(ByteCodec::None);
476        let data = b"Hello, world! This is test data for compression.";
477
478        let result = compressor.compress(data);
479        assert!(result.is_ok());
480
481        let compressed = result.unwrap();
482        assert_eq!(compressed.original_size, data.len());
483        assert_eq!(compressed.codec, ByteCodec::None);
484    }
485
486    #[test]
487    fn test_none_roundtrip() {
488        let compressor = SecureCompressor::with_default_security(ByteCodec::None);
489        let data = b"round-trip test";
490
491        let compressed = compressor.compress(data).unwrap();
492        let decompressed = compressor.decompress_protected(&compressed).unwrap();
493        assert_eq!(decompressed, data);
494    }
495
496    #[test]
497    fn test_compression_size_limit() {
498        let config = CompressionBombConfig {
499            max_compressed_size: 100, // Very small limit
500            ..Default::default()
501        };
502        let detector = CompressionBombDetector::new(config);
503        let compressor = SecureCompressor::new(detector, ByteCodec::None);
504
505        let large_data = vec![0u8; 1000]; // 1 KiB data
506        let result = compressor.compress(&large_data);
507
508        // Should fail pre-compression validation (compressed_size > max_compressed_size).
509        assert!(result.is_err());
510    }
511
512    #[test]
513    fn test_different_codecs_none() {
514        let compressor = SecureCompressor::with_default_security(ByteCodec::None);
515        let data = b"test data";
516
517        let result = compressor.compress(data);
518        assert!(result.is_ok());
519
520        let compressed = result.unwrap();
521        assert_eq!(compressed.compression_ratio, 1.0);
522        assert_eq!(compressed.codec, ByteCodec::None);
523    }
524
525    #[cfg(feature = "compression")]
526    mod compression_tests {
527        use super::*;
528
529        // ~4 KiB of repetitive JSON-like payload — should compress well.
530        fn repetitive_json() -> Vec<u8> {
531            let item = br#"{"id":1,"name":"test","value":42,"active":true}"#;
532            item.repeat(100)
533        }
534
535        #[test]
536        fn test_deflate_roundtrip() {
537            let compressor = SecureCompressor::with_default_security(ByteCodec::Deflate);
538            let data = repetitive_json();
539
540            let compressed = compressor.compress(&data).unwrap();
541            assert_eq!(compressed.codec, ByteCodec::Deflate);
542            assert!(
543                compressed.data.len() < data.len(),
544                "deflate must reduce size"
545            );
546
547            let decompressed = compressor.decompress_protected(&compressed).unwrap();
548            assert_eq!(decompressed, data);
549        }
550
551        #[test]
552        fn test_gzip_roundtrip() {
553            let compressor = SecureCompressor::with_default_security(ByteCodec::Gzip);
554            let data = repetitive_json();
555
556            let compressed = compressor.compress(&data).unwrap();
557            assert_eq!(compressed.codec, ByteCodec::Gzip);
558            assert!(compressed.data.len() < data.len(), "gzip must reduce size");
559
560            let decompressed = compressor.decompress_protected(&compressed).unwrap();
561            assert_eq!(decompressed, data);
562        }
563
564        #[test]
565        fn test_brotli_roundtrip() {
566            let compressor = SecureCompressor::with_default_security(ByteCodec::Brotli);
567            let data = repetitive_json();
568
569            let compressed = compressor.compress(&data).unwrap();
570            assert_eq!(compressed.codec, ByteCodec::Brotli);
571            assert!(
572                compressed.data.len() < data.len(),
573                "brotli must reduce size"
574            );
575
576            let decompressed = compressor.decompress_protected(&compressed).unwrap();
577            assert_eq!(decompressed, data);
578        }
579
580        #[test]
581        fn test_all_qualities_deflate() {
582            let data = repetitive_json();
583            for quality in [
584                CompressionQuality::Fast,
585                CompressionQuality::Balanced,
586                CompressionQuality::Best,
587            ] {
588                let c = SecureCompressor::with_quality(
589                    CompressionBombDetector::default(),
590                    ByteCodec::Deflate,
591                    quality,
592                );
593                let compressed = c.compress(&data).unwrap();
594                let decompressed = c.decompress_protected(&compressed).unwrap();
595                assert_eq!(decompressed, data);
596            }
597        }
598
599        #[test]
600        fn test_all_qualities_brotli() {
601            // Use Fast only to keep test time reasonable (quality 11 is slow).
602            let data = repetitive_json();
603            let c = SecureCompressor::with_quality(
604                CompressionBombDetector::default(),
605                ByteCodec::Brotli,
606                CompressionQuality::Fast,
607            );
608            let compressed = c.compress(&data).unwrap();
609            let decompressed = c.decompress_protected(&compressed).unwrap();
610            assert_eq!(decompressed, data);
611        }
612
613        #[test]
614        fn test_codec_mismatch_returns_error() {
615            // Compress with Brotli, but tell decompressor it is Gzip.
616            let c = SecureCompressor::with_default_security(ByteCodec::Brotli);
617            let data = b"codec mismatch test data";
618            let mut compressed = c.compress(data).unwrap();
619            compressed.codec = ByteCodec::Gzip; // wrong codec tag
620
621            let result = c.decompress_protected(&compressed);
622            assert!(
623                result.is_err(),
624                "wrong codec must produce an error, not garbage"
625            );
626        }
627
628        #[test]
629        fn test_bomb_detection_on_real_codec() {
630            // A very tight max_decompressed_size so any real inflation trips the guard.
631            let config = CompressionBombConfig {
632                max_decompressed_size: 200,  // Only 200 bytes allowed out
633                max_compressed_size: 10_000, // Allow the compressed input
634                max_ratio: 300.0,
635                check_interval_bytes: 64,
636                ..Default::default()
637            };
638            let detector = CompressionBombDetector::new(config);
639            let compressor =
640                SecureCompressor::new(CompressionBombDetector::default(), ByteCodec::Gzip);
641
642            // Produce a real gzip payload of ~4 KiB.
643            let data = repetitive_json();
644            let compressed = compressor.compress(&data).unwrap();
645
646            // Now decompress with a detector that caps at 200 bytes.
647            let strict_compressor = SecureCompressor::new(detector, ByteCodec::Gzip);
648            let result = strict_compressor.decompress_protected(&compressed);
649            assert!(
650                result.is_err(),
651                "bomb detector must stop oversized decompression"
652            );
653        }
654    }
655
656    #[test]
657    fn test_secure_decompression_context() {
658        let detector = CompressionBombDetector::default();
659        let mut context = SecureDecompressionContext::new(detector, 2);
660
661        assert!(context.start_stream(1024).is_ok());
662        assert!(context.start_stream(1024).is_ok());
663
664        // Third stream exceeds limit.
665        assert!(context.start_stream(1024).is_err());
666
667        context.finish_stream();
668        assert!(context.start_stream(1024).is_ok());
669    }
670
671    #[test]
672    fn test_context_stats() {
673        let detector = CompressionBombDetector::default();
674        let context = SecureDecompressionContext::new(detector, 5);
675
676        let stats = context.stats();
677        assert_eq!(stats.current_depth, 0);
678        assert_eq!(stats.active_streams, 0);
679        assert_eq!(stats.max_concurrent_streams, 5);
680    }
681
682    #[test]
683    fn test_context_finish_stream_underflow_safe() {
684        let detector = CompressionBombDetector::default();
685        let mut context = SecureDecompressionContext::new(detector, 5);
686
687        // finish_stream when active_streams == 0 must not underflow.
688        context.finish_stream();
689        let stats = context.stats();
690        assert_eq!(stats.active_streams, 0);
691    }
692
693    #[test]
694    fn test_byte_codec_default_is_none() {
695        assert_eq!(ByteCodec::default(), ByteCodec::None);
696    }
697
698    #[test]
699    fn test_byte_codec_clone() {
700        let codec = ByteCodec::None;
701        let cloned = codec.clone();
702        assert_eq!(codec, cloned);
703    }
704
705    #[test]
706    fn test_compression_quality_default_is_balanced() {
707        // Default quality must produce a valid compressor without error.
708        let c = SecureCompressor::with_default_security(ByteCodec::None);
709        let data = b"quality default test";
710        let compressed = c.compress(data).unwrap();
711        let decompressed = c.decompress_protected(&compressed).unwrap();
712        assert_eq!(decompressed.as_slice(), data);
713    }
714
715    #[test]
716    fn test_secure_compressed_data_clone() {
717        let c = SecureCompressor::with_default_security(ByteCodec::None);
718        let compressed = c.compress(b"clone test").unwrap();
719        let cloned = compressed.clone();
720        assert_eq!(compressed.data, cloned.data);
721        assert_eq!(compressed.original_size, cloned.original_size);
722        assert_eq!(compressed.codec, cloned.codec);
723    }
724
725    #[test]
726    fn test_none_roundtrip_empty_payload() {
727        let c = SecureCompressor::with_default_security(ByteCodec::None);
728        let compressed = c.compress(b"").unwrap();
729        let decompressed = c.decompress_protected(&compressed).unwrap();
730        assert_eq!(decompressed, b"");
731    }
732
733    #[test]
734    fn test_decompress_nested_none() {
735        let c = SecureCompressor::with_default_security(ByteCodec::None);
736        let data = b"nested roundtrip";
737        let compressed = c.compress(data).unwrap();
738        let decompressed = c.decompress_nested(&compressed, 0).unwrap();
739        assert_eq!(decompressed.as_slice(), data);
740    }
741
742    #[cfg(all(feature = "compression", not(target_arch = "wasm32")))]
743    mod zstd_dict_tests {
744        use super::*;
745        use crate::compression::zstd::{MAX_DICT_SIZE, N_TRAIN, ZstdDictCompressor};
746        use crate::security::CompressionBombConfig;
747        use std::sync::Arc;
748
749        fn repetitive_json() -> Vec<u8> {
750            let item = br#"{"id":1,"name":"test","value":42,"active":true}"#;
751            item.repeat(100)
752        }
753
754        fn trained_dict() -> crate::compression::zstd::ZstdDictionary {
755            let samples: Vec<Vec<u8>> = (0..N_TRAIN)
756                .map(|i| {
757                    format!(
758                        r#"{{"id":{i},"name":"item-{i}","value":{},"active":true}}"#,
759                        i * 10
760                    )
761                    .into_bytes()
762                })
763                .collect();
764            ZstdDictCompressor::train(&samples, MAX_DICT_SIZE).unwrap()
765        }
766
767        #[test]
768        fn test_zstd_dict_roundtrip_via_secure_compressor() {
769            let dict = Arc::new(trained_dict());
770            let compressor =
771                SecureCompressor::with_default_security(ByteCodec::ZstdDict(dict.clone()));
772            let data = repetitive_json();
773
774            let compressed = compressor.compress(&data).unwrap();
775            assert!(matches!(compressed.codec, ByteCodec::ZstdDict(_)));
776            assert!(
777                compressed.data.len() < data.len(),
778                "zstd dict must reduce size on repetitive data"
779            );
780
781            let decompressed = compressor.decompress_protected(&compressed).unwrap();
782            assert_eq!(decompressed, data);
783        }
784
785        #[test]
786        fn test_zstd_dict_bomb_detection() {
787            let dict = Arc::new(trained_dict());
788            let producer =
789                SecureCompressor::with_default_security(ByteCodec::ZstdDict(dict.clone()));
790            let data = repetitive_json();
791            let compressed = producer.compress(&data).unwrap();
792
793            let config = CompressionBombConfig {
794                max_decompressed_size: 200,
795                max_compressed_size: 10_000,
796                max_ratio: 300.0,
797                check_interval_bytes: 64,
798                ..Default::default()
799            };
800            let strict = SecureCompressor::new(
801                crate::security::CompressionBombDetector::new(config),
802                ByteCodec::ZstdDict(dict),
803            );
804            let result = strict.decompress_protected(&compressed);
805            assert!(
806                result.is_err(),
807                "bomb detector must block oversized zstd dict output"
808            );
809        }
810
811        #[test]
812        fn test_zstd_dict_codec_mismatch_errors() {
813            let dict = Arc::new(trained_dict());
814            let c = SecureCompressor::with_default_security(ByteCodec::ZstdDict(dict));
815            let data = b"codec mismatch test data";
816            let mut compressed = c.compress(data).unwrap();
817            // Lie about the codec — decoding as Gzip must fail.
818            compressed.codec = ByteCodec::Gzip;
819            assert!(
820                c.decompress_protected(&compressed).is_err(),
821                "wrong codec must produce an error"
822            );
823        }
824
825        #[test]
826        fn test_zstd_dict_empty_payload_roundtrip() {
827            let dict = Arc::new(trained_dict());
828            let c = SecureCompressor::with_default_security(ByteCodec::ZstdDict(dict));
829            let compressed = c.compress(b"").unwrap();
830            let decompressed = c.decompress_protected(&compressed).unwrap();
831            assert_eq!(decompressed, b"");
832        }
833
834        #[test]
835        fn test_zstd_dict_wrong_dictionary_errors() {
836            // Build two independent dictionaries from distinct corpora.
837            let samples_a: Vec<Vec<u8>> = (0..N_TRAIN)
838                .map(|i| format!(r#"{{"corpus":"alpha","id":{i},"score":{}}}"#, i * 7).into_bytes())
839                .collect();
840            let samples_b: Vec<Vec<u8>> = (0..N_TRAIN)
841                .map(|i| format!(r#"{{"corpus":"beta","seq":{i},"label":"x-{i}"}}"#).into_bytes())
842                .collect();
843
844            let dict_a = ZstdDictCompressor::train(&samples_a, MAX_DICT_SIZE).unwrap();
845            let dict_b = ZstdDictCompressor::train(&samples_b, MAX_DICT_SIZE).unwrap();
846
847            let data = b"some representative payload data";
848            let compressed =
849                ZstdDictCompressor::compress(data, &dict_a).expect("compress with dict_a");
850
851            // Decompressing dict_a-compressed bytes with dict_b must fail at libzstd level.
852            let result = ZstdDictCompressor::decompress(&compressed, &dict_b, data.len() * 4);
853            assert!(
854                result.is_err(),
855                "wrong dictionary must produce a libzstd error"
856            );
857        }
858    }
859
860    #[cfg(feature = "compression")]
861    mod extended_compression_tests {
862        use super::*;
863
864        // Non-repetitive payload: pseudo-random bytes unlikely to compress well.
865        fn incompressible_payload() -> Vec<u8> {
866            // Simple LCG to generate pseudo-random bytes without extra deps.
867            let mut state: u64 = 0x_dead_beef_cafe_babe;
868            (0..512)
869                .map(|_| {
870                    state = state
871                        .wrapping_mul(6_364_136_223_846_793_005)
872                        .wrapping_add(1);
873                    (state >> 33) as u8
874                })
875                .collect()
876        }
877
878        #[test]
879        fn test_deflate_roundtrip_incompressible() {
880            let c = SecureCompressor::with_default_security(ByteCodec::Deflate);
881            let data = incompressible_payload();
882            let compressed = c.compress(&data).unwrap();
883            assert_eq!(compressed.codec, ByteCodec::Deflate);
884            let decompressed = c.decompress_protected(&compressed).unwrap();
885            assert_eq!(decompressed, data);
886        }
887
888        #[test]
889        fn test_gzip_roundtrip_incompressible() {
890            let c = SecureCompressor::with_default_security(ByteCodec::Gzip);
891            let data = incompressible_payload();
892            let compressed = c.compress(&data).unwrap();
893            assert_eq!(compressed.codec, ByteCodec::Gzip);
894            let decompressed = c.decompress_protected(&compressed).unwrap();
895            assert_eq!(decompressed, data);
896        }
897
898        #[test]
899        fn test_brotli_roundtrip_incompressible() {
900            let c = SecureCompressor::with_default_security(ByteCodec::Brotli);
901            let data = incompressible_payload();
902            let compressed = c.compress(&data).unwrap();
903            assert_eq!(compressed.codec, ByteCodec::Brotli);
904            let decompressed = c.decompress_protected(&compressed).unwrap();
905            assert_eq!(decompressed, data);
906        }
907
908        #[test]
909        fn test_gzip_all_qualities() {
910            let item = br#"{"id":1,"name":"test","value":42}"#;
911            let data: Vec<u8> = item.repeat(50);
912            for quality in [
913                CompressionQuality::Fast,
914                CompressionQuality::Balanced,
915                CompressionQuality::Best,
916            ] {
917                let c = SecureCompressor::with_quality(
918                    CompressionBombDetector::default(),
919                    ByteCodec::Gzip,
920                    quality,
921                );
922                let compressed = c.compress(&data).unwrap();
923                let decompressed = c.decompress_protected(&compressed).unwrap();
924                assert_eq!(
925                    decompressed, data,
926                    "gzip quality {quality:?} roundtrip failed"
927                );
928            }
929        }
930
931        #[test]
932        fn test_brotli_balanced_quality() {
933            let item = br#"{"key":"value","n":99}"#;
934            let data: Vec<u8> = item.repeat(80);
935            let c = SecureCompressor::with_quality(
936                CompressionBombDetector::default(),
937                ByteCodec::Brotli,
938                CompressionQuality::Balanced,
939            );
940            let compressed = c.compress(&data).unwrap();
941            let decompressed = c.decompress_protected(&compressed).unwrap();
942            assert_eq!(decompressed, data);
943        }
944
945        #[test]
946        fn test_decompress_nested_with_depth() {
947            let c = SecureCompressor::with_default_security(ByteCodec::Deflate);
948            let item = br#"{"x":1}"#;
949            let data: Vec<u8> = item.repeat(100);
950            let compressed = c.compress(&data).unwrap();
951            let decompressed = c.decompress_nested(&compressed, 1).unwrap();
952            assert_eq!(decompressed, data);
953        }
954
955        #[test]
956        fn test_decompress_nested_depth_exceeded_returns_error() {
957            use crate::security::CompressionBombConfig;
958            let config = CompressionBombConfig {
959                max_compression_depth: 2,
960                ..Default::default()
961            };
962            let c = SecureCompressor::new(CompressionBombDetector::new(config), ByteCodec::Deflate);
963            let item = br#"{"x":1}"#;
964            let data: Vec<u8> = item.repeat(100);
965            let compressed = c.compress(&data).unwrap();
966            // depth 3 exceeds max_compression_depth 2 — must error.
967            let result = c.decompress_nested(&compressed, 3);
968            assert!(result.is_err(), "depth beyond limit must return an error");
969        }
970
971        #[test]
972        fn test_bomb_detection_deflate() {
973            use crate::security::CompressionBombConfig;
974            let config = CompressionBombConfig {
975                max_decompressed_size: 200,
976                max_compressed_size: 10_000,
977                max_ratio: 300.0,
978                check_interval_bytes: 64,
979                ..Default::default()
980            };
981            let item = br#"{"id":1,"name":"test","value":42,"active":true}"#;
982            let data: Vec<u8> = item.repeat(100);
983            let producer = SecureCompressor::with_default_security(ByteCodec::Deflate);
984            let compressed = producer.compress(&data).unwrap();
985
986            let strict =
987                SecureCompressor::new(CompressionBombDetector::new(config), ByteCodec::Deflate);
988            let result = strict.decompress_protected(&compressed);
989            assert!(
990                result.is_err(),
991                "bomb detector must block oversized deflate output"
992            );
993        }
994
995        #[test]
996        fn test_bomb_detection_brotli() {
997            use crate::security::CompressionBombConfig;
998            let config = CompressionBombConfig {
999                max_decompressed_size: 200,
1000                max_compressed_size: 10_000,
1001                max_ratio: 300.0,
1002                check_interval_bytes: 64,
1003                ..Default::default()
1004            };
1005            let item = br#"{"id":1,"name":"test","value":42,"active":true}"#;
1006            let data: Vec<u8> = item.repeat(100);
1007            let producer = SecureCompressor::with_default_security(ByteCodec::Brotli);
1008            let compressed = producer.compress(&data).unwrap();
1009
1010            let strict =
1011                SecureCompressor::new(CompressionBombDetector::new(config), ByteCodec::Brotli);
1012            let result = strict.decompress_protected(&compressed);
1013            assert!(
1014                result.is_err(),
1015                "bomb detector must block oversized brotli output"
1016            );
1017        }
1018
1019        #[test]
1020        fn test_codec_mismatch_deflate_as_gzip() {
1021            let c = SecureCompressor::with_default_security(ByteCodec::Deflate);
1022            let data = b"deflate mismatch test payload";
1023            let mut compressed = c.compress(data).unwrap();
1024            compressed.codec = ByteCodec::Gzip;
1025            let result = c.decompress_protected(&compressed);
1026            assert!(result.is_err(), "Deflate data decoded as Gzip must fail");
1027        }
1028
1029        #[test]
1030        fn test_empty_payload_all_codecs() {
1031            for codec in [ByteCodec::Deflate, ByteCodec::Gzip, ByteCodec::Brotli] {
1032                let label = format!("{codec:?}");
1033                let c = SecureCompressor::with_default_security(codec);
1034                let compressed = c.compress(b"").unwrap();
1035                let decompressed = c.decompress_protected(&compressed).unwrap();
1036                assert_eq!(decompressed, b"", "empty roundtrip failed for {label}");
1037            }
1038        }
1039    }
1040}