Merge commit from fork

Author: freddyheppell

tests/test_helpers.py

@@ -3,6 +3,7 @@
 
 import pytest
 
+from tests.helpers import gzip
 from usp.exceptions import (
     GunzipException,
     SitemapException,
@@ -216,6 +217,19 @@ def test_gunzip():
         # noinspection PyTypeChecker
         gunzip(b"foo")
 
+    assert gunzip(gzip("foo"), max_output_bytes=512 * 1024) == b"foo"
+
+
+def test_gunzip_above_max_output_bytes():
+    # Create a gzip-compressed byte string that exceeds the max_output_bytes limit
+    original_data = b"A" * 1024 * 1024  # 1 MB of data
+    compressed_data = gzip(original_data)
+
+    with pytest.raises(GunzipException) as exc_info:
+        gunzip(compressed_data, max_output_bytes=512 * 1024)  # Limit to 512 KB
+
+    assert "Gunzipped data exceeds maximum output size" in str(exc_info.value)
+
 
 class MockWebClientErrorResponse(WebClientErrorResponse):
     pass

usp/fetch_parse.py

@@ -172,7 +172,11 @@ def sitemap(self) -> AbstractSitemap:
 
         self._url = response_url
 
-        response_content = ungzipped_response_content(url=self._url, response=response)
+        response_content = ungzipped_response_content(
+            url=self._url,
+            response=response,
+            max_uncompressed_bytes=self.__MAX_SITEMAP_SIZE,
+        )
 
         # MIME types returned in Content-Type are unpredictable, so peek into the content instead
         if response_content[:20].strip().startswith("<"):

usp/helpers.py

@@ -3,6 +3,7 @@
 import datetime
 import gzip as gzip_lib
 import html
+import io
 import logging
 import re
 import sys
@@ -215,7 +216,7 @@ def __response_is_gzipped_data(
         return False
 
 
-def gunzip(data: bytes) -> bytes:
+def gunzip(data: bytes, max_output_bytes: int | None = None) -> bytes:
     """
     Gunzip data.
 
@@ -236,7 +237,16 @@ def gunzip(data: bytes) -> bytes:
         )
 
     try:
-        gunzipped_data = gzip_lib.decompress(data)
+        chunks, total = [], 0
+        with gzip_lib.GzipFile(fileobj=io.BytesIO(data)) as gz:
+            while chunk := gz.read(1024 * 1024):
+                total += len(chunk)
+                if max_output_bytes is not None and total > max_output_bytes:
+                    raise GunzipException(
+                        f"Gunzipped data exceeds maximum output size of {max_output_bytes} bytes."
+                    )
+                chunks.append(chunk)
+        gunzipped_data = b"".join(chunks)
     except Exception as ex:
         raise GunzipException(f"Unable to gunzip data: {str(ex)}")
 
@@ -250,7 +260,9 @@ def gunzip(data: bytes) -> bytes:
 
 
 def ungzipped_response_content(
-    url: str, response: AbstractWebClientSuccessResponse
+    url: str,
+    response: AbstractWebClientSuccessResponse,
+    max_uncompressed_bytes: int | None = None,
 ) -> str:
     """
     Return HTTP response's decoded content, gunzip it if necessary.
@@ -264,7 +276,7 @@ def ungzipped_response_content(
 
     if __response_is_gzipped_data(url=url, response=response):
         try:
-            data = gunzip(data)
+            data = gunzip(data, max_output_bytes=max_uncompressed_bytes)
         except GunzipException as ex:
             # In case of an error, just assume that it's one of the non-gzipped sitemaps with ".gz" extension
             log.warning(