diff --git a/joule-code/joule-code-rs/crates/jwp/src/compression.rs b/joule-code/joule-code-rs/crates/jwp/src/compression.rs index 2ee717436..9b5ef859f 100644 --- a/joule-code/joule-code-rs/crates/jwp/src/compression.rs +++ b/joule-code/joule-code-rs/crates/jwp/src/compression.rs @@ -114,11 +114,27 @@ impl CompressionStrategy for ZstdCompression { let mut output = Vec::with_capacity(data.len().min(max_output)); use std::io::Read; + // One byte past the cap, so the cap being hit is detectable. + // + // `take(max_output)` alone yields exactly `max_output` bytes and + // reports success, discarding the rest: the bomb case — the one + // condition `max_output` exists for — came back as a value, and the + // caller could not tell a complete payload from a truncated prefix of + // an arbitrarily larger one. Reading one extra byte distinguishes + // "exactly at the cap" (valid) from "over it" (rejected), and matches + // `Lz4Compression::decompress`, which has always rejected this. + let limit = (max_output as u64).saturating_add(1); decoder - .take(max_output as u64) + .take(limit) .read_to_end(&mut output) .map_err(|e| JwpError::CompressionError(e.to_string()))?; + if output.len() > max_output { + return Err(JwpError::CompressionError(format!( + "decompressed size exceeds max {max_output}" + ))); + } + Ok(output) } @@ -312,4 +328,40 @@ mod tests { let result = c.decompress(&compressed, 10); assert!(result.is_err()); } + + /// The zstd half of the trait contract — "`max_output` caps the + /// decompressed size to prevent decompression bombs". + /// + /// This used to return `Ok` with the first `max_output` bytes, so a + /// payload that expands past the cap was accepted as a complete one and + /// `AdaptiveCodec::decode` went on to rewrite the frame header to the + /// truncated length. Only the lz4 strategy was tested, which is why the + /// two implementations behind one trait had opposite semantics. + #[test] + fn zstd_rejects_oversized_decompression() { + let c = ZstdCompression::default(); + let compressed = c.compress(TEST_DATA).unwrap(); + assert!( + TEST_DATA.len() > 10, + "the fixture has to exceed the cap for this to test anything" + ); + let result = c.decompress(&compressed, 10); + assert!( + result.is_err(), + "a payload expanding past the cap must be rejected, not truncated: {:?}", + result.map(|v| v.len()) + ); + } + + /// The boundary the extra byte must not break: output exactly at the cap + /// is valid and comes back whole. + #[test] + fn zstd_accepts_output_exactly_at_the_cap() { + let c = ZstdCompression::default(); + let compressed = c.compress(TEST_DATA).unwrap(); + let out = c + .decompress(&compressed, TEST_DATA.len()) + .expect("exactly at the cap is within the cap"); + assert_eq!(out, TEST_DATA); + } }