dig_message/
compression.rs

1//! The compression layer (SPEC §1.1) — the crypto-free codec that runs BEFORE the seal (WU2). A
2//! payload is compressed in its OWN fresh context (never a shared running context — the §5.5 CRIME/
3//! BREACH boundary), then WU2 seals the compressed bytes.
4//!
5//! Two codecs ship in WU1: raw/identity (id 0, mandatory) and zstd level-3 (id 1, recommended). The
6//! reserved id bands (2..=63 standard, 64..=255 experimental) are rejected cleanly, never mis-decoded.
7
8use crate::constants::{MAX_DECOMPRESSED_BYTES, MIN_COMPRESS_BYTES, ZSTD_LEVEL};
9use crate::error::{MessageError, Result};
10
11/// Raw / identity codec: the payload bytes verbatim. Mandatory (SPEC §1.1).
12pub const COMPRESSION_NONE: u8 = 0;
13
14/// zstd codec, pinned to level 3, single-frame, no dictionary (SPEC §1.1/§1.2).
15pub const COMPRESSION_ZSTD: u8 = 1;
16
17/// A compressed payload ready to be sealed: the algorithm id, the compressed bytes, and the declared
18/// original length (the bomb-guard bound carried in the sealed `InnerMessage`, SPEC §5.2).
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct CompressedPayload {
21    /// The algorithm id ([`COMPRESSION_NONE`] or [`COMPRESSION_ZSTD`]).
22    pub compression: u8,
23    /// The on-wire compressed bytes (identical to the input for the raw codec).
24    pub bytes: Vec<u8>,
25    /// The original uncompressed length — the receiver's decompression-bomb bound.
26    pub uncompressed_len: u32,
27}
28
29/// Compress a payload for sealing, choosing the codec per the SPEC §1.1 raw threshold.
30///
31/// Uses zstd (id 1) only when the payload is at least [`MIN_COMPRESS_BYTES`] AND zstd actually shrinks
32/// it; otherwise falls back to raw (id 0) so small or incompressible payloads never pay a header/
33/// expansion penalty.
34///
35/// # Errors
36/// [`MessageError::PayloadTooLarge`] if the payload does not fit the `u32` length field;
37/// [`MessageError::Codec`] on a zstd encoder failure.
38pub fn compress_payload(payload: &[u8]) -> Result<CompressedPayload> {
39    let uncompressed_len =
40        u32::try_from(payload.len()).map_err(|_| MessageError::PayloadTooLarge(payload.len()))?;
41
42    if payload.len() < MIN_COMPRESS_BYTES {
43        return Ok(raw(payload, uncompressed_len));
44    }
45
46    let compressed = zstd::bulk::compress(payload, ZSTD_LEVEL)
47        .map_err(|e| MessageError::Codec(e.to_string()))?;
48
49    // Fall back to raw when zstd does not shrink (already-compressed / incompressible data), so we
50    // never emit a larger frame than the plaintext (SPEC §1.1).
51    if compressed.len() >= payload.len() {
52        return Ok(raw(payload, uncompressed_len));
53    }
54
55    Ok(CompressedPayload {
56        compression: COMPRESSION_ZSTD,
57        bytes: compressed,
58        uncompressed_len,
59    })
60}
61
62/// Decompress a sealed payload under the SPEC §1.1 bomb guard.
63///
64/// Rejects a declared length over [`MAX_DECOMPRESSED_BYTES`] BEFORE decoding, bounds the decoder
65/// output to that cap DURING, and rejects on any length mismatch — a hostile peer cannot OOM the host.
66///
67/// # Errors
68/// [`MessageError::DecompressionBomb`] if `uncompressed_len` exceeds the cap;
69/// [`MessageError::UnsupportedCompression`] for an unknown id;
70/// [`MessageError::DecompressedLengthMismatch`] on an overrun or a decoded-length disagreement;
71/// [`MessageError::Codec`] on a zstd decoder failure.
72pub fn decompress_payload(compression: u8, data: &[u8], uncompressed_len: u32) -> Result<Vec<u8>> {
73    let declared = uncompressed_len as usize;
74    if declared > MAX_DECOMPRESSED_BYTES {
75        return Err(MessageError::DecompressionBomb {
76            declared,
77            max: MAX_DECOMPRESSED_BYTES,
78        });
79    }
80
81    match compression {
82        COMPRESSION_NONE => {
83            if data.len() != declared {
84                return Err(MessageError::DecompressedLengthMismatch {
85                    expected: declared,
86                    actual: data.len(),
87                });
88            }
89            Ok(data.to_vec())
90        }
91        COMPRESSION_ZSTD => {
92            // `capacity = declared` bounds the allocation to the (already-capped) declared size; zstd
93            // errors if the frame decodes to more, defeating a bomb that under-declares its output.
94            let decoded = zstd::bulk::decompress(data, declared)
95                .map_err(|e| MessageError::Codec(e.to_string()))?;
96            if decoded.len() != declared {
97                return Err(MessageError::DecompressedLengthMismatch {
98                    expected: declared,
99                    actual: decoded.len(),
100                });
101            }
102            Ok(decoded)
103        }
104        other => Err(MessageError::UnsupportedCompression(other)),
105    }
106}
107
108/// The raw/identity result: the payload verbatim under [`COMPRESSION_NONE`].
109fn raw(payload: &[u8], uncompressed_len: u32) -> CompressedPayload {
110    CompressedPayload {
111        compression: COMPRESSION_NONE,
112        bytes: payload.to_vec(),
113        uncompressed_len,
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    /// Deterministic, compressible test bytes derived from a seed — never a hard-coded literal (CodeQL).
122    fn compressible(len: usize) -> Vec<u8> {
123        // A repeating low-entropy pattern zstd shrinks well.
124        (0..len).map(|i| (i % 7) as u8).collect()
125    }
126
127    #[test]
128    fn small_payload_stays_raw() {
129        let payload = compressible(MIN_COMPRESS_BYTES - 1);
130        let out = compress_payload(&payload).unwrap();
131        assert_eq!(out.compression, COMPRESSION_NONE);
132        assert_eq!(out.bytes, payload);
133        assert_eq!(out.uncompressed_len as usize, payload.len());
134    }
135
136    #[test]
137    fn large_compressible_payload_uses_zstd_and_round_trips() {
138        let payload = compressible(4096);
139        let out = compress_payload(&payload).unwrap();
140        assert_eq!(out.compression, COMPRESSION_ZSTD);
141        assert!(
142            out.bytes.len() < payload.len(),
143            "zstd should shrink a low-entropy payload"
144        );
145        let restored =
146            decompress_payload(out.compression, &out.bytes, out.uncompressed_len).unwrap();
147        assert_eq!(restored, payload);
148    }
149
150    #[test]
151    fn raw_round_trips() {
152        let payload = compressible(16);
153        let out = compress_payload(&payload).unwrap();
154        let restored =
155            decompress_payload(out.compression, &out.bytes, out.uncompressed_len).unwrap();
156        assert_eq!(restored, payload);
157    }
158
159    #[test]
160    fn incompressible_payload_falls_back_to_raw() {
161        // High-entropy (SHA-derived) bytes zstd cannot shrink -> the codec MUST fall back to raw
162        // (SPEC §1.1). Derived from a hashed seed, never a hard-coded literal (CodeQL).
163        use sha2::{Digest, Sha256};
164        let mut payload = Vec::new();
165        let mut ctr = 0u64;
166        while payload.len() < 1024 {
167            let mut h = Sha256::new();
168            h.update(b"incompressible");
169            h.update(ctr.to_le_bytes());
170            payload.extend_from_slice(&h.finalize());
171            ctr += 1;
172        }
173        let out = compress_payload(&payload).unwrap();
174        assert_eq!(out.compression, COMPRESSION_NONE);
175        assert_eq!(out.bytes, payload);
176    }
177
178    #[test]
179    fn compression_is_deterministic() {
180        let payload = compressible(2048);
181        assert_eq!(
182            compress_payload(&payload).unwrap(),
183            compress_payload(&payload).unwrap()
184        );
185    }
186
187    #[test]
188    fn unknown_compression_id_is_rejected() {
189        let err = decompress_payload(2, &[1, 2, 3], 3).unwrap_err();
190        assert_eq!(err, MessageError::UnsupportedCompression(2));
191    }
192
193    #[test]
194    fn declared_length_over_cap_is_a_bomb() {
195        let declared = (MAX_DECOMPRESSED_BYTES + 1) as u32;
196        // A u32 cannot exceed 64 MiB+1? 64 MiB = 67_108_864 which fits u32. Good.
197        let err = decompress_payload(COMPRESSION_ZSTD, &[], declared).unwrap_err();
198        assert_eq!(
199            err,
200            MessageError::DecompressionBomb {
201                declared: declared as usize,
202                max: MAX_DECOMPRESSED_BYTES
203            }
204        );
205    }
206
207    #[test]
208    fn zstd_bomb_that_underdeclares_output_is_rejected() {
209        // Compress a large payload, then claim a tiny uncompressed_len: the decoder is bounded to the
210        // (small) declared capacity and MUST reject the overrun (SPEC §1.1).
211        let payload = compressible(8192);
212        let out = compress_payload(&payload).unwrap();
213        assert_eq!(out.compression, COMPRESSION_ZSTD);
214        let err = decompress_payload(out.compression, &out.bytes, 16).unwrap_err();
215        // Either an overrun (Codec) or a length mismatch — both are a clean reject, never a panic/OOM.
216        assert!(matches!(
217            err,
218            MessageError::Codec(_) | MessageError::DecompressedLengthMismatch { .. }
219        ));
220    }
221
222    #[test]
223    fn raw_length_mismatch_is_rejected() {
224        let err = decompress_payload(COMPRESSION_NONE, &[1, 2, 3], 4).unwrap_err();
225        assert_eq!(
226            err,
227            MessageError::DecompressedLengthMismatch {
228                expected: 4,
229                actual: 3
230            }
231        );
232    }
233}
dig_message/compression.rs

dig_message/
compression.rs
