Protect against zstandard bombs

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

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)

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):