fix zstd decompression (#6921)

* fix zstd decompression (issue #6914)

* add our fix to CHANGELOG

* add explicit read_across_frames=True + move zstd test to test_encoding.py

---------

Co-authored-by: Maximilian Hils <git@maximilianhils.com>

CHANGELOG.md

@@ -7,7 +7,8 @@
 
 ## Unreleased: mitmproxy next
 
-
+* Fix zstd decompression to read across frames.
+  ([#6921](https://github.com/mitmproxy/mitmproxy/pull/6921), @zendai)
 
 ## 12 June 2024: mitmproxy 10.3.1
 

mitmproxy/net/encoding.py

@@ -174,12 +174,7 @@ def decode_zstd(content: bytes) -> bytes:
     if not content:
         return b""
     zstd_ctx = zstd.ZstdDecompressor()
-    try:
-        return zstd_ctx.decompress(content)
-    except zstd.ZstdError:
-        # If the zstd stream is streamed without a size header,
-        # try decoding with a 10MiB output buffer
-        return zstd_ctx.decompress(content, max_output_size=10 * 2**20)
+    return zstd_ctx.stream_reader(BytesIO(content), read_across_frames=True).read()
 
 
 def encode_zstd(content: bytes) -> bytes:

test/mitmproxy/net/test_encoding.py

@@ -93,3 +93,22 @@ def test_cache():
             # This is not in the cache anymore
             assert encoding.encode(b"decoded", "gzip") == b"encoded"
             assert encode_gzip.call_count == 1
+
+
+def test_zstd():
+    FRAME_SIZE = 1024
+
+    # Create payload of 1024b
+    test_content = "a" * FRAME_SIZE
+
+    # Compress it, will result a single frame
+    single_frame = encoding.encode_zstd(test_content.encode())
+
+    # Concat compressed frame, it'll result two frames, total size of 2048b payload
+    two_frames = single_frame + single_frame
+
+    # Uncompressed single frame should have the size of FRAME_SIZE
+    assert len(encoding.decode_zstd(single_frame)) == FRAME_SIZE
+
+    # Uncompressed two frames should have the size of FRAME_SIZE * 2
+    assert len(encoding.decode_zstd(two_frames)) == FRAME_SIZE * 2