From 9cc870387745f45f744ceef6ed226eefcce0e066 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adri=C3=A1n=20Chaves?= <adrian@chaves.io>
Date: Wed, 22 Nov 2023 17:53:00 +0100
Subject: [PATCH] Protect against zstandard bombs

---
 .../downloadermiddlewares/httpcompression.py  |  15 ++++++-----
 scrapy/utils/_compression.py                  |  25 ++++++++++++++++++
 tests/sample_data/compressed/bomb-zstd.bin    | Bin 0 -> 1096 bytes
 ...st_downloadermiddleware_httpcompression.py |   5 +++-
 4 files changed, 37 insertions(+), 8 deletions(-)
 create mode 100644 tests/sample_data/compressed/bomb-zstd.bin

diff --git a/scrapy/downloadermiddlewares/httpcompression.py b/scrapy/downloadermiddlewares/httpcompression.py
index 91748e57e..8ee1d95a6 100644
--- a/scrapy/downloadermiddlewares/httpcompression.py
+++ b/scrapy/downloadermiddlewares/httpcompression.py
@@ -1,10 +1,14 @@
-import io
 import warnings
 
 from scrapy.exceptions import IgnoreRequest, NotConfigured
 from scrapy.http import Response, TextResponse
 from scrapy.responsetypes import responsetypes
-from scrapy.utils._compression import _DecompressionMaxSizeExceeded, _inflate, _unbrotli
+from scrapy.utils._compression import (
+    _DecompressionMaxSizeExceeded,
+    _inflate,
+    _unbrotli,
+    _unzstd,
+)
 from scrapy.utils.deprecate import ScrapyDeprecationWarning
 from scrapy.utils.gz import gunzip
 
@@ -18,7 +22,7 @@ else:
     ACCEPTED_ENCODINGS.append(b"br")
 
 try:
-    import zstandard
+    import zstandard  # noqa: F401
 except ImportError:
     pass
 else:
@@ -102,8 +106,5 @@ class HttpCompressionMiddleware:
         if encoding == b"br" and b"br" in ACCEPTED_ENCODINGS:
             return _unbrotli(body, max_size=self._max_size)
         if encoding == b"zstd" and b"zstd" in ACCEPTED_ENCODINGS:
-            # Using its streaming API since its simple API could handle only cases
-            # where there is content size data embedded in the frame
-            reader = zstandard.ZstdDecompressor().stream_reader(io.BytesIO(body))
-            return reader.read()
+            return _unzstd(body, max_size=self._max_size)
         return body
diff --git a/scrapy/utils/_compression.py b/scrapy/utils/_compression.py
index 9a32ce4f0..93aa254b2 100644
--- a/scrapy/utils/_compression.py
+++ b/scrapy/utils/_compression.py
@@ -7,6 +7,11 @@ try:
 except ImportError:
     pass
 
+try:
+    import zstandard
+except ImportError:
+    pass
+
 
 class _DecompressionMaxSizeExceeded(ValueError):
     pass
@@ -65,3 +70,23 @@ def _unbrotli(data: bytes, *, max_size: int = 0) -> bytes:
             )
         output_list.append(output_chunk)
     return b"".join(output_list)
+
+
+def _unzstd(data: bytes, *, max_size: int = 0) -> bytes:
+    decompressor = zstandard.ZstdDecompressor()
+    stream_reader = decompressor.stream_reader(BytesIO(data))
+    output_list: List[bytes] = []
+    output_chunk = b"."
+    decompressed_size = 0
+    CHUNK_SIZE = 8196
+    while output_chunk:
+        output_chunk = stream_reader.read(CHUNK_SIZE)
+        decompressed_size += len(output_chunk)
+        if max_size and decompressed_size > max_size:
+            raise _DecompressionMaxSizeExceeded(
+                f"The number of bytes decompressed so far "
+                f"({decompressed_size}B) exceed the specified maximum "
+                f"({max_size}B)."
+            )
+        output_list.append(output_chunk)
+    return b"".join(output_list)
diff --git a/tests/sample_data/compressed/bomb-zstd.bin b/tests/sample_data/compressed/bomb-zstd.bin
new file mode 100644
index 0000000000000000000000000000000000000000..4b0efa8a41c88a38dffe7cea9e4a6726bdb137c4
GIT binary patch
literal 1096
zcmdPcs{gko!e;q;1{Fqz2O$}m#R@=_sF0kWTTql*T%4Jor;wDNo219Z$k6h?-fpfB
z0|SQwBg3EnmI6#5b|Mlx7l~br#Lh!vBdZBP5+5~lG&~1uT5@4vU|?kU`+vT_!f29*
VWPM*?)(2)~i{-##pxebr9RRD(6-595

literal 0
HcmV?d00001

diff --git a/tests/test_downloadermiddleware_httpcompression.py b/tests/test_downloadermiddleware_httpcompression.py
index 8858916bc..7babd1318 100644
--- a/tests/test_downloadermiddleware_httpcompression.py
+++ b/tests/test_downloadermiddleware_httpcompression.py
@@ -41,7 +41,7 @@ FORMAT = {
             "br",  # 34 → 11 511 612
             "deflate",  # 27 968 → 11 511 612
             "gzip",  # 27 988 → 11 511 612
-            # "zstd",
+            "zstd",  # 1 096 → 11 511 612
         )
     },
 }
@@ -396,6 +396,9 @@ class HttpCompressionTest(TestCase):
     def test_compression_bomb_gzip(self):
         self._test_compression_bomb("gzip")
 
+    def test_compression_bomb_zstd(self):
+        self._test_compression_bomb("zstd")
+
 
 class HttpCompressionSubclassTest(TestCase):
     def test_init_missing_stats(self):