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    pub current_depth: usize,
453    pub active_streams: usize,
454    pub max_concurrent_streams: usize,
455}
456
457#[cfg(test)]
458mod tests {
459    use super::*;
460    use crate::security::CompressionBombConfig;
461
462    #[test]
463    fn test_secure_compressor_creation() {
464        let detector = CompressionBombDetector::default();
465        let compressor = SecureCompressor::new(detector, ByteCodec::None);
466        // Verify the compressor is created (not null pointer).
467        assert!(!std::ptr::addr_of!(compressor).cast::<u8>().is_null());
468    }
469
470    #[test]
471    fn test_secure_compression_none() {
472        let compressor = SecureCompressor::with_default_security(ByteCodec::None);
473        let data = b"Hello, world! This is test data for compression.";
474
475        let result = compressor.compress(data);
476        assert!(result.is_ok());
477
478        let compressed = result.unwrap();
479        assert_eq!(compressed.original_size, data.len());
480        assert_eq!(compressed.codec, ByteCodec::None);
481    }
482
483    #[test]
484    fn test_none_roundtrip() {
485        let compressor = SecureCompressor::with_default_security(ByteCodec::None);
486        let data = b"round-trip test";
487
488        let compressed = compressor.compress(data).unwrap();
489        let decompressed = compressor.decompress_protected(&compressed).unwrap();
490        assert_eq!(decompressed, data);
491    }
492
493    #[test]
494    fn test_compression_size_limit() {
495        let config = CompressionBombConfig {
496            max_compressed_size: 100, // Very small limit
497            ..Default::default()
498        };
499        let detector = CompressionBombDetector::new(config);
500        let compressor = SecureCompressor::new(detector, ByteCodec::None);
501
502        let large_data = vec![0u8; 1000]; // 1 KiB data
503        let result = compressor.compress(&large_data);
504
505        // Should fail pre-compression validation (compressed_size > max_compressed_size).
506        assert!(result.is_err());
507    }
508
509    #[test]
510    fn test_different_codecs_none() {
511        let compressor = SecureCompressor::with_default_security(ByteCodec::None);
512        let data = b"test data";
513
514        let result = compressor.compress(data);
515        assert!(result.is_ok());
516
517        let compressed = result.unwrap();
518        assert_eq!(compressed.compression_ratio, 1.0);
519        assert_eq!(compressed.codec, ByteCodec::None);
520    }
521
522    #[cfg(feature = "compression")]
523    mod compression_tests {
524        use super::*;
525
526        // ~4 KiB of repetitive JSON-like payload — should compress well.
527        fn repetitive_json() -> Vec<u8> {
528            let item = br#"{"id":1,"name":"test","value":42,"active":true}"#;
529            item.repeat(100)
530        }
531
532        #[test]
533        fn test_deflate_roundtrip() {
534            let compressor = SecureCompressor::with_default_security(ByteCodec::Deflate);
535            let data = repetitive_json();
536
537            let compressed = compressor.compress(&data).unwrap();
538            assert_eq!(compressed.codec, ByteCodec::Deflate);
539            assert!(
540                compressed.data.len() < data.len(),
541                "deflate must reduce size"
542            );
543
544            let decompressed = compressor.decompress_protected(&compressed).unwrap();
545            assert_eq!(decompressed, data);
546        }
547
548        #[test]
549        fn test_gzip_roundtrip() {
550            let compressor = SecureCompressor::with_default_security(ByteCodec::Gzip);
551            let data = repetitive_json();
552
553            let compressed = compressor.compress(&data).unwrap();
554            assert_eq!(compressed.codec, ByteCodec::Gzip);
555            assert!(compressed.data.len() < data.len(), "gzip must reduce size");
556
557            let decompressed = compressor.decompress_protected(&compressed).unwrap();
558            assert_eq!(decompressed, data);
559        }
560
561        #[test]
562        fn test_brotli_roundtrip() {
563            let compressor = SecureCompressor::with_default_security(ByteCodec::Brotli);
564            let data = repetitive_json();
565
566            let compressed = compressor.compress(&data).unwrap();
567            assert_eq!(compressed.codec, ByteCodec::Brotli);
568            assert!(
569                compressed.data.len() < data.len(),
570                "brotli must reduce size"
571            );
572
573            let decompressed = compressor.decompress_protected(&compressed).unwrap();
574            assert_eq!(decompressed, data);
575        }
576
577        #[test]
578        fn test_all_qualities_deflate() {
579            let data = repetitive_json();
580            for quality in [
581                CompressionQuality::Fast,
582                CompressionQuality::Balanced,
583                CompressionQuality::Best,
584            ] {
585                let c = SecureCompressor::with_quality(
586                    CompressionBombDetector::default(),
587                    ByteCodec::Deflate,
588                    quality,
589                );
590                let compressed = c.compress(&data).unwrap();
591                let decompressed = c.decompress_protected(&compressed).unwrap();
592                assert_eq!(decompressed, data);
593            }
594        }
595
596        #[test]
597        fn test_all_qualities_brotli() {
598            // Use Fast only to keep test time reasonable (quality 11 is slow).
599            let data = repetitive_json();
600            let c = SecureCompressor::with_quality(
601                CompressionBombDetector::default(),
602                ByteCodec::Brotli,
603                CompressionQuality::Fast,
604            );
605            let compressed = c.compress(&data).unwrap();
606            let decompressed = c.decompress_protected(&compressed).unwrap();
607            assert_eq!(decompressed, data);
608        }
609
610        #[test]
611        fn test_codec_mismatch_returns_error() {
612            // Compress with Brotli, but tell decompressor it is Gzip.
613            let c = SecureCompressor::with_default_security(ByteCodec::Brotli);
614            let data = b"codec mismatch test data";
615            let mut compressed = c.compress(data).unwrap();
616            compressed.codec = ByteCodec::Gzip; // wrong codec tag
617
618            let result = c.decompress_protected(&compressed);
619            assert!(
620                result.is_err(),
621                "wrong codec must produce an error, not garbage"
622            );
623        }
624
625        #[test]
626        fn test_bomb_detection_on_real_codec() {
627            // A very tight max_decompressed_size so any real inflation trips the guard.
628            let config = CompressionBombConfig {
629                max_decompressed_size: 200,  // Only 200 bytes allowed out
630                max_compressed_size: 10_000, // Allow the compressed input
631                max_ratio: 300.0,
632                check_interval_bytes: 64,
633                ..Default::default()
634            };
635            let detector = CompressionBombDetector::new(config);
636            let compressor =
637                SecureCompressor::new(CompressionBombDetector::default(), ByteCodec::Gzip);
638
639            // Produce a real gzip payload of ~4 KiB.
640            let data = repetitive_json();
641            let compressed = compressor.compress(&data).unwrap();
642
643            // Now decompress with a detector that caps at 200 bytes.
644            let strict_compressor = SecureCompressor::new(detector, ByteCodec::Gzip);
645            let result = strict_compressor.decompress_protected(&compressed);
646            assert!(
647                result.is_err(),
648                "bomb detector must stop oversized decompression"
649            );
650        }
651    }
652
653    #[test]
654    fn test_secure_decompression_context() {
655        let detector = CompressionBombDetector::default();
656        let mut context = SecureDecompressionContext::new(detector, 2);
657
658        assert!(context.start_stream(1024).is_ok());
659        assert!(context.start_stream(1024).is_ok());
660
661        // Third stream exceeds limit.
662        assert!(context.start_stream(1024).is_err());
663
664        context.finish_stream();
665        assert!(context.start_stream(1024).is_ok());
666    }
667
668    #[test]
669    fn test_context_stats() {
670        let detector = CompressionBombDetector::default();
671        let context = SecureDecompressionContext::new(detector, 5);
672
673        let stats = context.stats();
674        assert_eq!(stats.current_depth, 0);
675        assert_eq!(stats.active_streams, 0);
676        assert_eq!(stats.max_concurrent_streams, 5);
677    }
678
679    #[test]
680    fn test_context_finish_stream_underflow_safe() {
681        let detector = CompressionBombDetector::default();
682        let mut context = SecureDecompressionContext::new(detector, 5);
683
684        // finish_stream when active_streams == 0 must not underflow.
685        context.finish_stream();
686        let stats = context.stats();
687        assert_eq!(stats.active_streams, 0);
688    }
689
690    #[test]
691    fn test_byte_codec_default_is_none() {
692        assert_eq!(ByteCodec::default(), ByteCodec::None);
693    }
694
695    #[test]
696    fn test_byte_codec_clone() {
697        let codec = ByteCodec::None;
698        let cloned = codec.clone();
699        assert_eq!(codec, cloned);
700    }
701
702    #[test]
703    fn test_compression_quality_default_is_balanced() {
704        // Default quality must produce a valid compressor without error.
705        let c = SecureCompressor::with_default_security(ByteCodec::None);
706        let data = b"quality default test";
707        let compressed = c.compress(data).unwrap();
708        let decompressed = c.decompress_protected(&compressed).unwrap();
709        assert_eq!(decompressed.as_slice(), data);
710    }
711
712    #[test]
713    fn test_secure_compressed_data_clone() {
714        let c = SecureCompressor::with_default_security(ByteCodec::None);
715        let compressed = c.compress(b"clone test").unwrap();
716        let cloned = compressed.clone();
717        assert_eq!(compressed.data, cloned.data);
718        assert_eq!(compressed.original_size, cloned.original_size);
719        assert_eq!(compressed.codec, cloned.codec);
720    }
721
722    #[test]
723    fn test_none_roundtrip_empty_payload() {
724        let c = SecureCompressor::with_default_security(ByteCodec::None);
725        let compressed = c.compress(b"").unwrap();
726        let decompressed = c.decompress_protected(&compressed).unwrap();
727        assert_eq!(decompressed, b"");
728    }
729
730    #[test]
731    fn test_decompress_nested_none() {
732        let c = SecureCompressor::with_default_security(ByteCodec::None);
733        let data = b"nested roundtrip";
734        let compressed = c.compress(data).unwrap();
735        let decompressed = c.decompress_nested(&compressed, 0).unwrap();
736        assert_eq!(decompressed.as_slice(), data);
737    }
738
739    #[cfg(all(feature = "compression", not(target_arch = "wasm32")))]
740    mod zstd_dict_tests {
741        use super::*;
742        use crate::compression::zstd::{MAX_DICT_SIZE, N_TRAIN, ZstdDictCompressor};
743        use crate::security::CompressionBombConfig;
744        use std::sync::Arc;
745
746        fn repetitive_json() -> Vec<u8> {
747            let item = br#"{"id":1,"name":"test","value":42,"active":true}"#;
748            item.repeat(100)
749        }
750
751        fn trained_dict() -> crate::compression::zstd::ZstdDictionary {
752            let samples: Vec<Vec<u8>> = (0..N_TRAIN)
753                .map(|i| {
754                    format!(
755                        r#"{{"id":{i},"name":"item-{i}","value":{},"active":true}}"#,
756                        i * 10
757                    )
758                    .into_bytes()
759                })
760                .collect();
761            ZstdDictCompressor::train(&samples, MAX_DICT_SIZE).unwrap()
762        }
763
764        #[test]
765        fn test_zstd_dict_roundtrip_via_secure_compressor() {
766            let dict = Arc::new(trained_dict());
767            let compressor =
768                SecureCompressor::with_default_security(ByteCodec::ZstdDict(dict.clone()));
769            let data = repetitive_json();
770
771            let compressed = compressor.compress(&data).unwrap();
772            assert!(matches!(compressed.codec, ByteCodec::ZstdDict(_)));
773            assert!(
774                compressed.data.len() < data.len(),
775                "zstd dict must reduce size on repetitive data"
776            );
777
778            let decompressed = compressor.decompress_protected(&compressed).unwrap();
779            assert_eq!(decompressed, data);
780        }
781
782        #[test]
783        fn test_zstd_dict_bomb_detection() {
784            let dict = Arc::new(trained_dict());
785            let producer =
786                SecureCompressor::with_default_security(ByteCodec::ZstdDict(dict.clone()));
787            let data = repetitive_json();
788            let compressed = producer.compress(&data).unwrap();
789
790            let config = CompressionBombConfig {
791                max_decompressed_size: 200,
792                max_compressed_size: 10_000,
793                max_ratio: 300.0,
794                check_interval_bytes: 64,
795                ..Default::default()
796            };
797            let strict = SecureCompressor::new(
798                crate::security::CompressionBombDetector::new(config),
799                ByteCodec::ZstdDict(dict),
800            );
801            let result = strict.decompress_protected(&compressed);
802            assert!(
803                result.is_err(),
804                "bomb detector must block oversized zstd dict output"
805            );
806        }
807
808        #[test]
809        fn test_zstd_dict_codec_mismatch_errors() {
810            let dict = Arc::new(trained_dict());
811            let c = SecureCompressor::with_default_security(ByteCodec::ZstdDict(dict));
812            let data = b"codec mismatch test data";
813            let mut compressed = c.compress(data).unwrap();
814            // Lie about the codec — decoding as Gzip must fail.
815            compressed.codec = ByteCodec::Gzip;
816            assert!(
817                c.decompress_protected(&compressed).is_err(),
818                "wrong codec must produce an error"
819            );
820        }
821
822        #[test]
823        fn test_zstd_dict_empty_payload_roundtrip() {
824            let dict = Arc::new(trained_dict());
825            let c = SecureCompressor::with_default_security(ByteCodec::ZstdDict(dict));
826            let compressed = c.compress(b"").unwrap();
827            let decompressed = c.decompress_protected(&compressed).unwrap();
828            assert_eq!(decompressed, b"");
829        }
830
831        #[test]
832        fn test_zstd_dict_wrong_dictionary_errors() {
833            // Build two independent dictionaries from distinct corpora.
834            let samples_a: Vec<Vec<u8>> = (0..N_TRAIN)
835                .map(|i| format!(r#"{{"corpus":"alpha","id":{i},"score":{}}}"#, i * 7).into_bytes())
836                .collect();
837            let samples_b: Vec<Vec<u8>> = (0..N_TRAIN)
838                .map(|i| format!(r#"{{"corpus":"beta","seq":{i},"label":"x-{i}"}}"#).into_bytes())
839                .collect();
840
841            let dict_a = ZstdDictCompressor::train(&samples_a, MAX_DICT_SIZE).unwrap();
842            let dict_b = ZstdDictCompressor::train(&samples_b, MAX_DICT_SIZE).unwrap();
843
844            let data = b"some representative payload data";
845            let compressed =
846                ZstdDictCompressor::compress(data, &dict_a).expect("compress with dict_a");
847
848            // Decompressing dict_a-compressed bytes with dict_b must fail at libzstd level.
849            let result = ZstdDictCompressor::decompress(&compressed, &dict_b, data.len() * 4);
850            assert!(
851                result.is_err(),
852                "wrong dictionary must produce a libzstd error"
853            );
854        }
855    }
856
857    #[cfg(feature = "compression")]
858    mod extended_compression_tests {
859        use super::*;
860
861        // Non-repetitive payload: pseudo-random bytes unlikely to compress well.
862        fn incompressible_payload() -> Vec<u8> {
863            // Simple LCG to generate pseudo-random bytes without extra deps.
864            let mut state: u64 = 0x_dead_beef_cafe_babe;
865            (0..512)
866                .map(|_| {
867                    state = state
868                        .wrapping_mul(6_364_136_223_846_793_005)
869                        .wrapping_add(1);
870                    (state >> 33) as u8
871                })
872                .collect()
873        }
874
875        #[test]
876        fn test_deflate_roundtrip_incompressible() {
877            let c = SecureCompressor::with_default_security(ByteCodec::Deflate);
878            let data = incompressible_payload();
879            let compressed = c.compress(&data).unwrap();
880            assert_eq!(compressed.codec, ByteCodec::Deflate);
881            let decompressed = c.decompress_protected(&compressed).unwrap();
882            assert_eq!(decompressed, data);
883        }
884
885        #[test]
886        fn test_gzip_roundtrip_incompressible() {
887            let c = SecureCompressor::with_default_security(ByteCodec::Gzip);
888            let data = incompressible_payload();
889            let compressed = c.compress(&data).unwrap();
890            assert_eq!(compressed.codec, ByteCodec::Gzip);
891            let decompressed = c.decompress_protected(&compressed).unwrap();
892            assert_eq!(decompressed, data);
893        }
894
895        #[test]
896        fn test_brotli_roundtrip_incompressible() {
897            let c = SecureCompressor::with_default_security(ByteCodec::Brotli);
898            let data = incompressible_payload();
899            let compressed = c.compress(&data).unwrap();
900            assert_eq!(compressed.codec, ByteCodec::Brotli);
901            let decompressed = c.decompress_protected(&compressed).unwrap();
902            assert_eq!(decompressed, data);
903        }
904
905        #[test]
906        fn test_gzip_all_qualities() {
907            let item = br#"{"id":1,"name":"test","value":42}"#;
908            let data: Vec<u8> = item.repeat(50);
909            for quality in [
910                CompressionQuality::Fast,
911                CompressionQuality::Balanced,
912                CompressionQuality::Best,
913            ] {
914                let c = SecureCompressor::with_quality(
915                    CompressionBombDetector::default(),
916                    ByteCodec::Gzip,
917                    quality,
918                );
919                let compressed = c.compress(&data).unwrap();
920                let decompressed = c.decompress_protected(&compressed).unwrap();
921                assert_eq!(
922                    decompressed, data,
923                    "gzip quality {quality:?} roundtrip failed"
924                );
925            }
926        }
927
928        #[test]
929        fn test_brotli_balanced_quality() {
930            let item = br#"{"key":"value","n":99}"#;
931            let data: Vec<u8> = item.repeat(80);
932            let c = SecureCompressor::with_quality(
933                CompressionBombDetector::default(),
934                ByteCodec::Brotli,
935                CompressionQuality::Balanced,
936            );
937            let compressed = c.compress(&data).unwrap();
938            let decompressed = c.decompress_protected(&compressed).unwrap();
939            assert_eq!(decompressed, data);
940        }
941
942        #[test]
943        fn test_decompress_nested_with_depth() {
944            let c = SecureCompressor::with_default_security(ByteCodec::Deflate);
945            let item = br#"{"x":1}"#;
946            let data: Vec<u8> = item.repeat(100);
947            let compressed = c.compress(&data).unwrap();
948            let decompressed = c.decompress_nested(&compressed, 1).unwrap();
949            assert_eq!(decompressed, data);
950        }
951
952        #[test]
953        fn test_decompress_nested_depth_exceeded_returns_error() {
954            use crate::security::CompressionBombConfig;
955            let config = CompressionBombConfig {
956                max_compression_depth: 2,
957                ..Default::default()
958            };
959            let c = SecureCompressor::new(CompressionBombDetector::new(config), ByteCodec::Deflate);
960            let item = br#"{"x":1}"#;
961            let data: Vec<u8> = item.repeat(100);
962            let compressed = c.compress(&data).unwrap();
963            // depth 3 exceeds max_compression_depth 2 — must error.
964            let result = c.decompress_nested(&compressed, 3);
965            assert!(result.is_err(), "depth beyond limit must return an error");
966        }
967
968        #[test]
969        fn test_bomb_detection_deflate() {
970            use crate::security::CompressionBombConfig;
971            let config = CompressionBombConfig {
972                max_decompressed_size: 200,
973                max_compressed_size: 10_000,
974                max_ratio: 300.0,
975                check_interval_bytes: 64,
976                ..Default::default()
977            };
978            let item = br#"{"id":1,"name":"test","value":42,"active":true}"#;
979            let data: Vec<u8> = item.repeat(100);
980            let producer = SecureCompressor::with_default_security(ByteCodec::Deflate);
981            let compressed = producer.compress(&data).unwrap();
982
983            let strict =
984                SecureCompressor::new(CompressionBombDetector::new(config), ByteCodec::Deflate);
985            let result = strict.decompress_protected(&compressed);
986            assert!(
987                result.is_err(),
988                "bomb detector must block oversized deflate output"
989            );
990        }
991
992        #[test]
993        fn test_bomb_detection_brotli() {
994            use crate::security::CompressionBombConfig;
995            let config = CompressionBombConfig {
996                max_decompressed_size: 200,
997                max_compressed_size: 10_000,
998                max_ratio: 300.0,
999                check_interval_bytes: 64,
1000                ..Default::default()
1001            };
1002            let item = br#"{"id":1,"name":"test","value":42,"active":true}"#;
1003            let data: Vec<u8> = item.repeat(100);
1004            let producer = SecureCompressor::with_default_security(ByteCodec::Brotli);
1005            let compressed = producer.compress(&data).unwrap();
1006
1007            let strict =
1008                SecureCompressor::new(CompressionBombDetector::new(config), ByteCodec::Brotli);
1009            let result = strict.decompress_protected(&compressed);
1010            assert!(
1011                result.is_err(),
1012                "bomb detector must block oversized brotli output"
1013            );
1014        }
1015
1016        #[test]
1017        fn test_codec_mismatch_deflate_as_gzip() {
1018            let c = SecureCompressor::with_default_security(ByteCodec::Deflate);
1019            let data = b"deflate mismatch test payload";
1020            let mut compressed = c.compress(data).unwrap();
1021            compressed.codec = ByteCodec::Gzip;
1022            let result = c.decompress_protected(&compressed);
1023            assert!(result.is_err(), "Deflate data decoded as Gzip must fail");
1024        }
1025
1026        #[test]
1027        fn test_empty_payload_all_codecs() {
1028            for codec in [ByteCodec::Deflate, ByteCodec::Gzip, ByteCodec::Brotli] {
1029                let label = format!("{codec:?}");
1030                let c = SecureCompressor::with_default_security(codec);
1031                let compressed = c.compress(b"").unwrap();
1032                let decompressed = c.decompress_protected(&compressed).unwrap();
1033                assert_eq!(decompressed, b"", "empty roundtrip failed for {label}");
1034            }
1035        }
1036    }
1037}