From 7edced5e7a0d9907b952bcb6aa70b47effca9e8c Mon Sep 17 00:00:00 2001 From: Vlada Dusek Date: Thu, 2 Jul 2026 17:13:24 +0200 Subject: [PATCH] fix: bound sitemap decompression size and nesting depth to prevent DoS --- src/crawlee/_utils/sitemap.py | 74 ++++++++++++---- .../_sitemap_request_loader.py | 25 +++++- tests/unit/_utils/test_sitemap.py | 86 ++++++++++++++++++- .../test_sitemap_request_loader.py | 42 ++++++++- 4 files changed, 205 insertions(+), 22 deletions(-) diff --git a/src/crawlee/_utils/sitemap.py b/src/crawlee/_utils/sitemap.py index 190381c1f2..c3690e2a50 100644 --- a/src/crawlee/_utils/sitemap.py +++ b/src/crawlee/_utils/sitemap.py @@ -37,6 +37,15 @@ SITEMAP_URL_PATTERN = re.compile(r'\/sitemap\.(?:xml|txt)(?:\.gz)?$', re.IGNORECASE) COMMON_SITEMAP_PATHS = ['/sitemap.xml', '/sitemap.txt', '/sitemap_index.xml'] +MAX_SITEMAP_SIZE = 50 * 1024 * 1024 +"""Maximum sitemap size in bytes after decompression, per the sitemap protocol; content beyond it is truncated.""" + +CONTENT_CHUNK_SIZE = 1024 * 1024 +"""Maximum size of a single decompressed content piece fed to the parser, keeping memory usage flat.""" + +DEFAULT_MAX_DEPTH = 10 +"""Default maximum depth of nested sitemaps to follow, guarding against malicious infinite sitemap chains.""" + @dataclass() class SitemapUrl: @@ -368,6 +377,9 @@ async def _fetch_and_process_sitemap( # Create appropriate parser parser = _get_parser(content_type, sitemap_url) decompressor = None + content_bytes = 0 + raw_bytes = 0 + truncated = False try: # Process chunks as they arrive first_chunk = True @@ -377,20 +389,50 @@ async def _fetch_and_process_sitemap( decompressor = zlib.decompressobj(zlib.MAX_WBITS | 16) first_chunk = False - chunk = decompressor.decompress(raw_chunk) if decompressor else raw_chunk - text_chunk = decoder.decode(chunk) - async for item in parser.process_chunk(text_chunk): - async for result in _process_sitemap_item( - item, - source, - depth, - visited_sitemap_urls, - sources, - emit_nested_sitemaps=emit_nested_sitemaps, - enqueue_strategy=enqueue_strategy, - ): - if result: - yield result + raw_bytes += len(raw_chunk) + + # Decompress and parse in bounded pieces, capping the total decompressed size, + # so that e.g. a gzip bomb cannot exhaust memory. Slice via a memoryview to avoid + # repeatedly copying the tail when a single chunk is much larger than a piece. + pending: memoryview | bytes = memoryview(raw_chunk) + while pending and content_bytes < MAX_SITEMAP_SIZE: + max_length = min(MAX_SITEMAP_SIZE - content_bytes, CONTENT_CHUNK_SIZE) + if decompressor: + chunk = decompressor.decompress(pending, max_length) + pending = decompressor.unconsumed_tail + else: + chunk, pending = pending[:max_length], pending[max_length:] + content_bytes += len(chunk) + + text_chunk = decoder.decode(chunk) + async for item in parser.process_chunk(text_chunk): + async for result in _process_sitemap_item( + item, + source, + depth, + visited_sitemap_urls, + sources, + emit_nested_sitemaps=emit_nested_sitemaps, + enqueue_strategy=enqueue_strategy, + ): + if result: + yield result + + # Stop once the decompressed-size cap leaves unprocessed bytes, or the compressed input + # exceeds the cap (e.g. a stream that produces little or no output). Both are truncations. + if pending or raw_bytes > MAX_SITEMAP_SIZE: + truncated = True + break + # Stop once the gzip stream has ended; further bytes would only accumulate in the + # decompressor's `unused_data` (trailing garbage or extra members are not processed). + if decompressor and decompressor.eof: + break + + if truncated: + logger.warning( + f'Sitemap {sitemap_url} exceeded the maximum size of {MAX_SITEMAP_SIZE} bytes, ' + f'processing only the truncated content.' + ) # Process any remaining content async for item in parser.flush(): @@ -472,7 +514,7 @@ async def parse_sitemap( This function coordinates the process of fetching and parsing sitemaps, handling both URL-based and raw content sources. It follows nested sitemaps - up to the specified maximum depth. + up to the specified maximum depth (`ParseSitemapOptions.max_depth`, 10 by default). Default `ParseSitemapOptions.enqueue_strategy` is `same-hostname` which will skip cross-host URLs. Use strategy `all` to process all links. @@ -480,7 +522,7 @@ async def parse_sitemap( # Set default options options = options or {} emit_nested_sitemaps = options.get('emit_nested_sitemaps', False) - max_depth = options.get('max_depth', float('inf')) + max_depth = options.get('max_depth', DEFAULT_MAX_DEPTH) sitemap_retries = options.get('sitemap_retries', 3) timeout = options.get('timeout', timedelta(seconds=30)) enqueue_strategy = options.get('enqueue_strategy', 'same-hostname') diff --git a/src/crawlee/request_loaders/_sitemap_request_loader.py b/src/crawlee/request_loaders/_sitemap_request_loader.py index 368347e708..47b026bbd4 100644 --- a/src/crawlee/request_loaders/_sitemap_request_loader.py +++ b/src/crawlee/request_loaders/_sitemap_request_loader.py @@ -14,7 +14,14 @@ from crawlee._utils.docs import docs_group from crawlee._utils.globs import Glob from crawlee._utils.recoverable_state import RecoverableState -from crawlee._utils.sitemap import NestedSitemap, ParseSitemapOptions, SitemapSource, SitemapUrl, parse_sitemap +from crawlee._utils.sitemap import ( + DEFAULT_MAX_DEPTH, + NestedSitemap, + ParseSitemapOptions, + SitemapSource, + SitemapUrl, + parse_sitemap, +) from crawlee._utils.urls import filter_url from crawlee.request_loaders._request_loader import RequestLoader @@ -80,6 +87,9 @@ class SitemapRequestLoaderState(BaseModel): processed_sitemap_urls: Annotated[set[str], Field(alias='processedSitemapUrls')] = set() """Set of processed sitemap URLs.""" + sitemap_depths: Annotated[dict[str, int], Field(alias='sitemapDepths')] = {} + """Nesting depth of each known sitemap URL, used to bound how far nested sitemap chains are followed.""" + completed: Annotated[bool, Field(alias='sitemapCompleted')] = False """Whether all sitemaps have been fully processed.""" @@ -286,6 +296,8 @@ async def _get_state(self) -> SitemapRequestLoaderState: ) if not has_sitemap_for_processing and not self._state.current_value.completed: self._state.current_value.pending_sitemap_urls.extend(self._sitemap_urls) + for sitemap_url in self._sitemap_urls: + self._state.current_value.sitemap_depths.setdefault(sitemap_url, 0) if self._state.current_value.in_progress: self._state.current_value.url_queue.extendleft(self._state.current_value.in_progress) @@ -352,6 +364,10 @@ async def _load_sitemaps(self) -> None: continue state.in_progress_sitemap_url = sitemap_url + # Depth of the sitemap currently being processed; nested sitemaps found in it are one level deeper. + # This bounds a malicious chain of unique nested sitemap URLs, which exact-URL dedup alone cannot. + current_depth = state.sitemap_depths.get(sitemap_url, 0) + parse_options = ParseSitemapOptions( max_depth=0, emit_nested_sitemaps=True, @@ -375,9 +391,15 @@ async def _load_sitemaps(self) -> None: item.loc not in state.pending_sitemap_urls and item.loc not in state.processed_sitemap_urls ): + if current_depth >= DEFAULT_MAX_DEPTH: + logger.warning( + f'Skipping nested sitemap {item.loc!r}: max depth {DEFAULT_MAX_DEPTH} reached.' + ) + continue if not self._passes_filters(item.loc, parsed_sitemap_url, 'nested sitemap'): continue state.pending_sitemap_urls.append(item.loc) + state.sitemap_depths[item.loc] = current_depth + 1 continue if isinstance(item, SitemapUrl): @@ -413,6 +435,7 @@ async def _load_sitemaps(self) -> None: state.in_progress_sitemap_url = None if current_sitemap_url: state.processed_sitemap_urls.add(current_sitemap_url) + state.sitemap_depths.pop(current_sitemap_url, None) state.current_sitemap_processed_urls.clear() # Mark as completed after processing all sitemap urls diff --git a/tests/unit/_utils/test_sitemap.py b/tests/unit/_utils/test_sitemap.py index e746c314c7..d209a4a3d4 100644 --- a/tests/unit/_utils/test_sitemap.py +++ b/tests/unit/_utils/test_sitemap.py @@ -10,6 +10,7 @@ from yarl import URL from crawlee._utils.sitemap import ( + DEFAULT_MAX_DEPTH, ParseSitemapOptions, Sitemap, SitemapUrl, @@ -22,7 +23,7 @@ from tests.unit.utils import DEFAULT_URL, get_basic_results, get_basic_sitemap if TYPE_CHECKING: - from collections.abc import AsyncIterator + from collections.abc import AsyncIterator, Callable def _make_mock_client(url_map: dict[str, tuple[int, bytes]]) -> AsyncMock: @@ -66,6 +67,27 @@ async def read_stream() -> 'AsyncIterator[bytes]': return client, attempts +def _make_stream_client(body_for_url: 'Callable[[str], bytes]') -> tuple[AsyncMock, list[str]]: + """Create a mock client whose `stream` serves `body_for_url(url)` in a single chunk and records fetched URLs.""" + fetched: list[str] = [] + + @asynccontextmanager + async def stream(url: str, **_kwargs: Any) -> 'AsyncIterator[HttpResponse]': + fetched.append(url) + + async def read_stream() -> 'AsyncIterator[bytes]': + yield body_for_url(url) + + response = MagicMock(spec=HttpResponse) + response.headers = {'content-type': 'application/xml; charset=utf-8'} + response.read_stream = read_stream + yield cast('HttpResponse', response) + + client = AsyncMock(spec=HttpClient) + client.stream = stream + return client, fetched + + def compress_gzip(data: str) -> bytes: """Compress a string using gzip.""" return gzip.compress(data.encode()) @@ -317,6 +339,68 @@ async def test_sitemap_fetch_raises_after_retries_exhausted() -> None: assert len(attempts) == 3 +async def test_gzip_bomb_sitemap_truncated_at_size_cap(monkeypatch: pytest.MonkeyPatch) -> None: + """A gzip sitemap inflating past the size cap is truncated instead of being decompressed without bound.""" + monkeypatch.setattr('crawlee._utils.sitemap.MAX_SITEMAP_SIZE', 64 * 1024) + total_urls = 20_000 + locs = ''.join(f'{DEFAULT_URL}page-{i}' for i in range(total_urls)) + # ~1 MB of XML arriving as a single small compressed chunk. + bomb = compress_gzip(f'{locs}') + + client, _ = _make_stream_client(lambda _url: bomb) + items = [item async for item in parse_sitemap([{'type': 'url', 'url': f'{DEFAULT_URL}sitemap.xml'}], client)] + + assert 0 < len(items) < total_urls + + +async def test_gzip_sitemap_stops_reading_after_member_end() -> None: + """Trailing bytes after a complete gzip member must not be read; otherwise they accumulate unbounded.""" + member = compress_gzip(get_basic_sitemap()) + chunks_read = 0 + + @asynccontextmanager + async def stream(_url: str, **_kwargs: Any) -> 'AsyncIterator[HttpResponse]': + async def read_stream() -> 'AsyncIterator[bytes]': + nonlocal chunks_read + chunks_read += 1 + yield member + # A malicious server would stream trailing junk forever; reading any of it is the bug. + while True: + chunks_read += 1 + yield b'\x00' * 65536 + + response = MagicMock(spec=HttpResponse) + response.headers = {'content-type': 'application/gzip'} + response.read_stream = read_stream + yield cast('HttpResponse', response) + + client = AsyncMock(spec=HttpClient) + client.stream = stream + + items = [item async for item in parse_sitemap([{'type': 'url', 'url': f'{DEFAULT_URL}sitemap.xml.gz'}], client)] + + assert {item.loc for item in items} == get_basic_results() + # The member fits in the first chunk, so at most the first trailing chunk may be pulled before we stop. + assert chunks_read <= 2 + + +async def test_nested_sitemaps_followed_only_to_default_max_depth() -> None: + """A chain of unique nested sitemap URLs is followed only up to the default max depth.""" + chain_length = 3 * DEFAULT_MAX_DEPTH + + def body_for_url(url: str) -> bytes: + index = int(url.removesuffix('.xml').rsplit('-', 1)[-1]) + if index >= chain_length: + return f'{DEFAULT_URL}page'.encode() + next_sitemap = f'{DEFAULT_URL}sitemap-{index + 1}.xml' + return f'{next_sitemap}'.encode() + + client, fetched = _make_stream_client(body_for_url) + _ = [item async for item in parse_sitemap([{'type': 'url', 'url': f'{DEFAULT_URL}sitemap-0.xml'}], client)] + + assert len(fetched) == DEFAULT_MAX_DEPTH + 1 + + async def test_parse_sitemap_with_partial_options() -> None: """Test that missing keys in partial `ParseSitemapOptions` fall back to defaults.""" options = ParseSitemapOptions(timeout=timedelta(seconds=10)) diff --git a/tests/unit/request_loaders/test_sitemap_request_loader.py b/tests/unit/request_loaders/test_sitemap_request_loader.py index abfd83cd62..a70c133afa 100644 --- a/tests/unit/request_loaders/test_sitemap_request_loader.py +++ b/tests/unit/request_loaders/test_sitemap_request_loader.py @@ -1,17 +1,21 @@ import base64 import gzip -from typing import TYPE_CHECKING -from unittest.mock import patch +from contextlib import asynccontextmanager +from typing import TYPE_CHECKING, Any, cast +from unittest.mock import AsyncMock, MagicMock, patch from yarl import URL from crawlee import RequestOptions, RequestTransformAction -from crawlee.http_clients._base import HttpClient +from crawlee._utils.sitemap import DEFAULT_MAX_DEPTH +from crawlee.http_clients._base import HttpClient, HttpResponse from crawlee.request_loaders._sitemap_request_loader import SitemapRequestLoader from crawlee.storages import KeyValueStore -from tests.unit.utils import get_basic_results, get_basic_sitemap, poll_until_condition +from tests.unit.utils import DEFAULT_URL, get_basic_results, get_basic_sitemap, poll_until_condition if TYPE_CHECKING: + from collections.abc import AsyncIterator + from crawlee._types import JsonSerializable @@ -25,6 +29,36 @@ def encode_base64(data: bytes) -> str: return base64.b64encode(data).decode('utf-8') +async def test_nested_sitemap_chain_bounded_by_max_depth() -> None: + """A malicious endless chain of unique nested sitemaps is followed only up to the default max depth.""" + fetched: list[str] = [] + + @asynccontextmanager + async def stream(url: str, **_kwargs: Any) -> 'AsyncIterator[HttpResponse]': + fetched.append(url) + index = int(url.removesuffix('.xml').rsplit('-', 1)[-1]) + # Every sitemap is an index pointing at the next unique URL, so only the depth cap can stop the chain. + next_sitemap = f'{DEFAULT_URL}sitemap-{index + 1}.xml' + body = f'{next_sitemap}'.encode() + + async def read_stream() -> 'AsyncIterator[bytes]': + yield body + + response = MagicMock(spec=HttpResponse) + response.status_code = 200 + response.headers = {'content-type': 'application/xml; charset=utf-8'} + response.read_stream = read_stream + yield cast('HttpResponse', response) + + client = AsyncMock(spec=HttpClient) + client.stream = stream + + loader = SitemapRequestLoader([f'{DEFAULT_URL}sitemap-0.xml'], http_client=client) + assert await poll_until_condition(loader.is_finished) + + assert len(fetched) == DEFAULT_MAX_DEPTH + 1 + + async def test_sitemap_traversal(server_url: URL, http_client: HttpClient) -> None: sitemap_url = (server_url / 'sitemap.xml').with_query( base64=encode_base64(get_basic_sitemap(url=server_url).encode())