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);
+    }
 }