[WiP] Rewrite zstd decoder to use an API that supports multiple frames (fix issue #3008)

src/urllib3/response.py

@@ -164,18 +164,37 @@ def flush(self) -> bytes:
 
     class ZstdDecoder(ContentDecoder):
         def __init__(self) -> None:
-            self._obj = zstd.ZstdDecompressor().decompressobj()
+            self._zstd_stream = io.BytesIO()
+            self._sr = zstd.ZstdDecompressor().stream_reader(
+                self._zstd_stream, read_across_frames=True)
 
         def decompress(self, data: bytes) -> bytes:
             if not data:
                 return b""
-            return self._obj.decompress(data)  # type: ignore[no-any-return]
+
+            # Push more data to the end, then go back to the previous position.
+            current_position = self._zstd_stream.tell()
+            self._zstd_stream.seek(0, io.SEEK_END)
+            self._zstd_stream.write(data)
+            self._zstd_stream.seek(current_position)
+            bytes_available = (len(self._zstd_stream.getvalue()) -
+                current_position)
+
+            if bytes_available >= zstd.DECOMPRESSION_RECOMMENDED_INPUT_SIZE:
+                return self._sr.read(zstd.DECOMPRESSION_RECOMMENDED_INPUT_SIZE)
+
+            # Not enough data available for a decompress operation.
+            return b""
+
+        def _is_stream_at_end(self) -> bool:
+            return (len(self._zstd_stream.getvalue()) ==
+                    self._zstd_stream.tell())
 
         def flush(self) -> bytes:
-            ret = self._obj.flush()
-            if not self._obj.eof:
-                raise DecodeError("Zstandard data is incomplete")
-            return ret  # type: ignore[no-any-return]
+            if self._is_stream_at_end():
+                return b""
+
+            return self._sr.readall()
 
 
 class MultiDecoder(ContentDecoder):
@@ -439,15 +458,16 @@ def _decode(
             if self._decoder:
                 data = self._decoder.decompress(data)
                 self._has_decoded_content = True
+
+                if flush_decoder:
+                    data += self._flush_decoder()
         except self.DECODER_ERROR_CLASSES as e:
             content_encoding = self.headers.get("content-encoding", "").lower()
             raise DecodeError(
                 "Received response with content-encoding: %s, but "
                 "failed to decode it." % content_encoding,
                 e,
             ) from e
-        if flush_decoder:
-            data += self._flush_decoder()
 
         return data
 

test/test_response.py

@@ -2,6 +2,7 @@
 
 import contextlib
 import http.client as httplib
+import os
 import socket
 import ssl
 import sys
@@ -332,6 +333,17 @@ def test_decode_zstd(self) -> None:
         r = HTTPResponse(fp, headers={"content-encoding": "zstd"})
         assert r.data == b"foo"
 
+    @onlyZstd()
+    def test_decode_zstd_multiple_frames(self) -> None:
+        # TODO: Can we dynamically generate data that caused the bug too?
+        with open(os.path.join(os.path.dirname(__file__), "text.txt.zstd"), "rb") as f:
+            data = f.read()
+
+        fp = BytesIO(data)
+        r = HTTPResponse(fp, headers={"content-encoding": "zstd"})
+        # Each frame size is 1048576 so correct data should be longer.
+        assert len(r.data) > 1048576
+
     @onlyZstd()
     def test_chunked_decoding_zstd(self) -> None:
         data = zstd.compress(b"foobarbaz")