Zstandard: fix crash on frame boundaries

Co-authored-by: Michiel W. Beijen <mb@x14.nl>

Author: Rogdham

httpx/_decoders.py

@@ -179,28 +179,27 @@ def __init__(self) -> None:
             ) from None
 
         self.decompressor = zstd.ZstdDecompressor()
-        self.seen_data = False
+        self.at_valid_eof = True
 
     def decode(self, data: bytes) -> bytes:
         assert zstd is not None
-        self.seen_data = True
         output = io.BytesIO()
         try:
-            output.write(self.decompressor.decompress(data))
-            while self.decompressor.eof and self.decompressor.unused_data:
-                unused_data = self.decompressor.unused_data
-                self.decompressor = zstd.ZstdDecompressor()
-                output.write(self.decompressor.decompress(unused_data))
+            while data:
+                self.at_valid_eof = False
+                output.write(self.decompressor.decompress(data))
+                data = self.decompressor.unused_data
+                if self.decompressor.eof:
+                    self.at_valid_eof = True
+                    self.decompressor = zstd.ZstdDecompressor()
         except zstd.ZstdError as exc:
             raise DecodingError(str(exc)) from exc
         return output.getvalue()
 
     def flush(self) -> bytes:
-        if not self.seen_data:
+        if self.at_valid_eof:
             return b""
-        if not self.decompressor.eof:
-            raise DecodingError("Zstandard data is incomplete")  # pragma: no cover
-        return b""
+        raise DecodingError("Zstandard data is incomplete")  # pragma: no cover
 
 
 class MultiDecoder(ContentDecoder):

tests/test_decoders.py

@@ -146,6 +146,26 @@ def test_zstd_multiframe():
     assert response.content == b"foobar"
 
 
+def test_zstd_streaming_multiple_frames():
+    body1 = b"test 123 "
+    body2 = b"another frame"
+
+    # Create two separate complete frames
+    frame1 = zstd.compress(body1)
+    frame2 = zstd.compress(body2)
+
+    # Create an iterator that yields frames separately
+    def content_iterator() -> typing.Iterator[bytes]:
+        yield frame1
+        yield frame2
+
+    headers = [(b"Content-Encoding", b"zstd")]
+    response = httpx.Response(200, headers=headers, content=content_iterator())
+    response.read()
+
+    assert response.content == body1 + body2
+
+
 def test_multi():
     body = b"test 123"