Fix data loss in ZstdDecompressor.decompress(max_length=0)

The zero-cap probe iteration handed libzstd a 1-byte buffer and the post-loop truncation discarded whatever it emitted, while the input counter had already advanced past the bytes that produced it — every decompress(..., max_length=0) call with pending output silently dropped one byte. The suite never caught it because max_length=0 was only exercised against zero-output (skippable) frames. Match CPython's mechanism instead: probe with a zero-size output buffer, so libzstd consumes input without emitting and never reports frame completion with un-emitted content. Add a regression snippet covering lossless zero-cap probes (content, truncated, and skippable frames), verified byte-exact against CPython.
JamesClarke7283 · May 23, 2026 · May 23, 2026 · May 23, 2026 · May 23, 2026 · May 23, 2026
commit 8fe59b8e930a524231d3260909469e029866868d
diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs
@@ -1403,17 +1403,18 @@ mod _zstd {
             let outcome = loop {
                 iteration += 1;
                 // Honor `max_length`: stop growing the output buffer once
-                // we have produced enough. Special-case the first iteration
-                // when the cap is zero so a zero-output frame (skippable
-                // frame, empty content frame) can still complete; we hand
-                // libzstd a 1-byte scratch and discard the byte if it ends
-                // up writing one.
+                // we have produced enough. When the cap is zero, hand
+                // libzstd a zero-size buffer instead: it then consumes
+                // input without emitting, so zero-output frames (skippable
+                // frame, empty content frame) still complete while real
+                // output stays inside libzstd until a later call has room.
+                // CPython's decompressor uses the same zero-size mechanism.
                 let grow = match max_length {
                     Some(maxl) if output.len() >= maxl && iteration > 1 => {
                         hit_max = true;
                         break Ok(());
                     }
-                    Some(maxl) if output.len() >= maxl => 1,
+                    Some(maxl) if output.len() >= maxl => 0,
                     Some(maxl) => (maxl - output.len()).min(chunk_size),
                     None => chunk_size,
                 };
@@ -1441,7 +1442,10 @@ mod _zstd {
                     state.eof = true;
                     break Ok(());
                 }
-                let output_was_full = written == grow;
+                // Only meaningful with a non-zero buffer: libzstd had more
+                // to emit but ran out of room. A zero-size call (the
+                // max_length == 0 probe) is not "full".
+                let output_was_full = grow > 0 && written == grow;
                 let input_consumed = input.pos == input.size;
 
                 if let Some(maxl) = max_length
@@ -1462,18 +1466,7 @@ mod _zstd {
             outcome.map(|()| (output, hit_max, input.pos))
         });
 
-        let (mut output, mut hit_max, consumed) =
-            loop_result.map_err(|c| catch_zstd_error(c, vm))?;
-
-        // If `max_length == 0` opened a courtesy iteration that produced more
-        // bytes than the caller asked for, truncate. Should not happen with
-        // the scratch slicing above, but keep the safety net.
-        if let Some(maxl) = max_length
-            && output.len() > maxl
-        {
-            output.truncate(maxl);
-            hit_max = true;
-        }
+        let (output, hit_max, consumed) = loop_result.map_err(|c| catch_zstd_error(c, vm))?;
 
         let remaining = &work_data[consumed..];
 

diff --git a/extra_tests/snippets/stdlib_zstd.py b/extra_tests/snippets/stdlib_zstd.py
@@ -0,0 +1,51 @@
+from compression.zstd import ZstdCompressor, ZstdDecompressor
+
+# basic roundtrip
+payload = b"hello world" * 100
+c = ZstdCompressor()
+frame = c.compress(payload) + c.flush()
+d = ZstdDecompressor()
+out = d.decompress(frame)
+assert out == payload
+assert d.eof
+assert not d.needs_input
+
+# max_length=0 must not drop output bytes: probing with a zero cap has to
+# leave every byte available to later drain calls (regression: the probe
+# used to emit one byte per call into the void).
+d = ZstdDecompressor()
+out = d.decompress(frame, 0)
+assert out == b""
+assert not d.eof
+assert not d.needs_input
+rest = d.decompress(b"", len(payload))
+assert rest == payload
+assert d.eof
+
+# two zero-cap probes in a row, then drain; still lossless
+d = ZstdDecompressor()
+assert d.decompress(frame, 0) == b""
+assert d.decompress(b"", 0) == b""
+assert d.decompress(b"", len(payload)) == payload
+assert d.eof
+
+# zero-cap probe on a truncated frame: no output, not at end
+d = ZstdDecompressor()
+out = d.decompress(frame[:-5], 0)
+assert out == b""
+assert not d.eof
+assert not d.needs_input
+
+# zero-cap probe on a skippable frame completes it
+skippable = b"\x50\x2a\x4d\x18\x04\x00\x00\x00abcd"
+d = ZstdDecompressor()
+assert d.decompress(skippable, 0) == b""
+assert d.eof
+assert not d.needs_input
+assert d.unused_data == b""
+
+# skippable frame followed by trailing bytes lands in unused_data
+d = ZstdDecompressor()
+assert d.decompress(skippable + b"trail", 0) == b""
+assert d.eof
+assert d.unused_data == b"trail"