Repository: urllib3/urllib3

Issue:
Partial responses when using chunked encoding + zstd compression
### Subject

I am testing using zstd compression with [bazel-remote](https://github.com/buchgr/bazel-remote/). As background, bazel-remote provides content-addressible storage under the /cas/ folder and storage of encoded build action data under the /ac/ folder. For the purposes of this bug report, I am only doing /cas/ downloads. So bazel-remote is used for basic PUT and GET of binary data.

Without providing the server headers to say that I support zstd compression, I don't receive chunked encoding and response.data contains the full binary data that I would expect. When I provide the server headers to say that I support zstd compression, the response from the server indicates chunked encoding and len(response.data) is always set to 1048576 bytes for files that are larger than that.

### Environment

OS Windows-10-10.0.22621-SP0
Python 3.10.11
urllib3 2.0.1

pip freeze results:

appdirs==1.4.4
certifi==2022.9.24
charset-normalizer==2.1.1
colorama==0.4.6
flake8==5.0.4
idna==3.4
mccabe==0.7.0
numpy==1.23.4
pycodestyle==2.9.1
pydiffx==1.1
pyflakes==2.5.0
python-gitlab==3.12.0
RBTools==4.0
requests==2.28.1
requests-toolbelt==0.10.1
six==1.16.0
texttable==1.6.7
tqdm==4.64.1
typing_extensions==4.4.0
urllib3==2.0.1
zstandard==0.21.0

### Steps to Reproduce

1. I deployed bazel-remote and did a PUT of a file larger than 1048576 bytes into /cas/<sha_256_hash_of_the_file_contents>
2. I ran the following code:
```
import urllib3

server_url = 'http://<server_name_redacted>:8000'
cas_url = '/cas/d1dcece4b328f7a12d86dd77cb0373496da2dd3971db2f3e3a29612f831a9738'
headers = urllib3.make_headers(accept_encoding=True)

response = urllib3.request('GET', server_url + cas_url, headers=headers)

print('Status: %d' % response.status)
print('Info: %s' % response.info())
print('Got %d back' % len(response.data))
```
### Expected Behavior

Script should print: Got 1308160 back.

### Actual Behavior

Script prints: Got 1048576 back

### Notes

If I set headers to None, the script outputs the following:
```
Status: 200
Info: HTTPHeaderDict({'Content-Length': '1308160', 'Content-Type': 'application/octet-stream', 'Date': 'Tue, 02 May 2023 12:47:40 GMT'})
Got 1308160 back
```
If I initialize headers to support compression, the script outputs the following:
```
Status: 200
Info: HTTPHeaderDict({'Content-Encoding': 'zstd', 'Content-Type': 'application/octet-stream', 'Date': 'Tue, 02 May 2023 12:48:24 GMT', 'Transfer-Encoding': 'chunked'})
Got 1048576 back
```


Discussion:
@mamico @indygreg Any idea what is going on here? (@grossag will soon report more data that seems to show that with the way we use python-zstandard, we don't decompress everything.)
Some updates from my current investigation:

curl 8.0.1 works and gives me a valid output .dll:
```
C:\Utils\temp>c:\utils\libcurl\bin\curl -w "We downloaded %{size_download} bytes\n" -D - --compressed <server_address_redacted>/cas/d1dcece4b328f7a12d86dd77cb0373496da2dd3971db2f3e3a29612f831a9738 -o C:\utils\temp\a.dll
  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                 Dload  Upload   Total   Spent    Left  Speed
  0     0    0     0    0     0      0      0 --:--:-- --:--:-- --:--:--     0HTTP/1.1 200 OK
Content-Encoding: zstd
Content-Type: application/octet-stream
Date: Tue, 02 May 2023 13:36:04 GMT
Transfer-Encoding: chunked

100  303k    0  303k    0     0   555k      0 --:--:-- --:--:-- --:--:--  559k
We downloaded 311215 bytes
```

Both urllib3 and curl are downloading 311215 bytes, which seems right. But if I add this to response.py:
```
                print('Decoding %d bytes' % len(data))
                data = self._decoder.decompress(data)
                print('Decoded into %d bytes' % len(data))
```
I see:
Decoding 311215 bytes
Decoded into 1048576 bytes

So the source data is correctly 311215 bytes but zstd's decompress API is only giving a chunk. This data size of 1048576 seems to match some magic numbers in python-zstandard's cffi backend.
It does seem to be something related to the usage of python-zstandard or a bug inside of python-zstandard. Inside of urllib3's ZstdDecoder.decompress(), I added this code:
```
            with open(r'C:\utils\temp\a.dll.zstd', 'wb') as f:
                f.write(data)
```
and then I ran:
```
C:\Utils\temp>win64_vc140\bin\zstd.exe -d a.dll.zstd -o manual.dll
a.dll.zstd          : 1308160 bytes
```
and ended up with a correct file manual.dll.

I also added the following code inside of decompress()
```

            print('approach 1: decompressed into %d bytes' %
                  len(zstd.decompress(data, max_output_size=9999999999)))
            print('approach 2: decompressed into %d bytes' %
                  len(zstd.ZstdDecompressor().decompress(data, max_output_size=9999999999)))
```
and got this output every time:
```
approach 1: decompressed into 1048576 bytes
approach 2: decompressed into 1048576 bytes
```
(which notably is less than the expected 1308160 which zstd.exe produces)
[text.txt.zip](https://github.com/urllib3/urllib3/files/11374676/text.txt.zip)

I have attached a zip of text.txt.zstd because GitHub doesn't support attaching .zstd files. This file text.txt.zstd contains the raw data passed to ZstdDecoder.decompress(). Extract that text.txt.zstd file out of the zip then run:
```
import zstandard as zstd

with open(r'C:\utils\temp\text.txt.zstd', 'rb') as f:
    data = f.read()

print('approach 1: decompressed into %d bytes' %
        len(zstd.decompress(data)))
print('approach 2: decompressed into %d bytes' %
        len(zstd.ZstdDecompressor().decompress(data)))
print('approach 3: decompressed into %d bytes' %
        len(zstd.ZstdDecompressor().decompressobj().decompress(data)))
```
This produces:
```
approach 1: decompressed into 1048576 bytes
approach 2: decompressed into 1048576 bytes
approach 3: decompressed into 1048576 bytes
```
Whereas `zstd.exe text.txt.zstd -d -o hi.txt` produces 8724469 bytes and a valid text file.
For context, we had a similar issue with all compressed HTTP responses being truncated when using the `.stream()` method: https://github.com/urllib3/urllib3/pull/3012 Perhaps this is a similar, something to be aware of.
Ref for downstream issue: https://github.com/indygreg/python-zstandard/issues/196
@grossag  Could you try again with the newest v2.0.2 urllib3 release?
I tried with 2.0.2 and could reproduce (which makes sense as we have a python-zstandard reproducer).

This is still an issue.
> I tried with 2.0.2 and could reproduce (which makes sense as we have a python-zstandard reproducer).
> 
> This is still an issue.

+1. Tried and confirmed that this is still an issue.

Resolving patch:
diff --git a/src/urllib3/response.py b/src/urllib3/response.py
--- a/src/urllib3/response.py
+++ b/src/urllib3/response.py
@@ -169,10 +169,15 @@ def __init__(self) -> None:
         def decompress(self, data: bytes) -> bytes:
             if not data:
                 return b""
-            return self._obj.decompress(data)  # type: ignore[no-any-return]
+            data_parts = [self._obj.decompress(data)]
+            while self._obj.eof and self._obj.unused_data:
+                unused_data = self._obj.unused_data
+                self._obj = zstd.ZstdDecompressor().decompressobj()
+                data_parts.append(self._obj.decompress(unused_data))
+            return b"".join(data_parts)
 
         def flush(self) -> bytes:
-            ret = self._obj.flush()
+            ret = self._obj.flush()  # note: this is a no-op
             if not self._obj.eof:
                 raise DecodeError("Zstandard data is incomplete")
             return ret  # type: ignore[no-any-return]


Test patch:
diff --git a/test/test_response.py b/test/test_response.py
--- a/test/test_response.py
+++ b/test/test_response.py
@@ -332,6 +332,25 @@ def test_decode_zstd(self) -> None:
         r = HTTPResponse(fp, headers={"content-encoding": "zstd"})
         assert r.data == b"foo"
 
+    @onlyZstd()
+    def test_decode_multiframe_zstd(self) -> None:
+        data = (
+            # Zstandard frame
+            zstd.compress(b"foo")
+            # skippable frame (must be ignored)
+            + bytes.fromhex(
+                "50 2A 4D 18"  # Magic_Number (little-endian)
+                "07 00 00 00"  # Frame_Size (little-endian)
+                "00 00 00 00 00 00 00"  # User_Data
+            )
+            # Zstandard frame
+            + zstd.compress(b"bar")
+        )
+
+        fp = BytesIO(data)
+        r = HTTPResponse(fp, headers={"content-encoding": "zstd"})
+        assert r.data == b"foobar"
+
     @onlyZstd()
     def test_chunked_decoding_zstd(self) -> None:
         data = zstd.compress(b"foobarbaz")