diff --git a/crates/common/src/compression/zlib.rs b/crates/common/src/compression/zlib.rs index f955f9be7d0..6c34d7cbe55 100644 --- a/crates/common/src/compression/zlib.rs +++ b/crates/common/src/compression/zlib.rs @@ -670,17 +670,16 @@ impl Decompressor { }; let consumed = (d.total_in() - prev_in) as usize; - // save unused input - let unconsumed = &data[consumed..]; - if !unconsumed.is_empty() { - if stream_end { - unused_data.extend_from_slice(unconsumed); - } else { - *unconsumed_tail = unconsumed.to_vec(); - } - } else if !unconsumed_tail.is_empty() { - unconsumed_tail.clear(); - } + // save unused input: past the end of the stream it is unused_data and + // unconsumed_tail is emptied, otherwise it is the new unconsumed_tail + let unconsumed: &[u8] = if stream_end { + unused_data.extend_from_slice(&data[consumed..]); + &[] + } else { + &data[consumed..] + }; + unconsumed_tail.clear(); + unconsumed_tail.extend_from_slice(unconsumed); Ok((ret, stream_end)) } @@ -1005,4 +1004,22 @@ mod tests { assert_eq!(d.decompress(&encoded, None).unwrap(), b"later input"); assert!(d.eof()); } + + #[test] + fn stream_end_clears_unconsumed_tail() { + let source = b"x".repeat(20000); + let data = [compress(&source, -1, MAX_WBITS).unwrap(), b"NEXT".to_vec()].concat(); + let mut d = Decompressor::new(MAX_WBITS, None).unwrap(); + let first = d.decompress(&data, Some(100)).unwrap(); + assert!(!d.eof()); + let tail = d.unconsumed_tail().to_vec(); + let rest = d.decompress(&tail, Some(100_000)).unwrap(); + assert!(d.eof()); + assert_eq!([first, rest].concat(), source); + assert_eq!(d.unconsumed_tail(), b""); + assert_eq!(d.unused_data(), b"NEXT"); + assert_eq!(d.flush(DEF_BUF_SIZE).unwrap(), b""); + assert_eq!(d.unconsumed_tail(), b""); + assert_eq!(d.unused_data(), b"NEXT"); + } } diff --git a/extra_tests/snippets/stdlib_zlib.py b/extra_tests/snippets/stdlib_zlib.py index ac2b525b0a7..dfadf7eebdc 100644 --- a/extra_tests/snippets/stdlib_zlib.py +++ b/extra_tests/snippets/stdlib_zlib.py @@ -62,3 +62,20 @@ assert_raises(zlib.error, lambda: zlib.compress(b"123", -40)) assert_raises(zlib.error, lambda: zlib.compress(b"123", 10)) + +# A max_length-bounded call that reaches the end of the stream keeps no +# stale input in unconsumed_tail (gh-8906). The tail is empty, or, as in +# CPython before gh-158169, the bytes after the stream. +source = b"x" * 20000 +d = zlib.decompressobj() +out = d.decompress(zlib.compress(source) + b"NEXT", 100) +assert not d.eof +out += d.decompress(d.unconsumed_tail, 100000) +assert d.eof +assert out == source +assert d.unused_data == b"NEXT" +tail = d.unconsumed_tail +assert tail in (b"", b"NEXT") +assert d.flush() == b"" +assert d.unconsumed_tail == tail +assert d.unused_data == b"NEXT" + tail