From 0c0619731e892cdf7fa945571b77753cd10861c6 Mon Sep 17 00:00:00 2001 From: simpleqt <89645338+simpleqt@users.noreply.github.com> Date: Mon, 24 Aug 2026 19:12:03 +0800 Subject: [PATCH] fix: harden LinkContentFetcher against SSRF Validate every request target before the request is made, including each redirect hop: redirects are now followed manually with per-hop re-validation and a max_redirects cap (default 5). Add an optional allowed_hosts domain suffix whitelist and always reject hosts resolving to private, loopback, link-local, multicast, reserved, unspecified, unique-local, or shared address space ranges. Stream response bodies with a max_response_bytes cap (default 50 MB) and raise a clear error when it is exceeded. --- haystack/components/fetchers/link_content.py | 275 +++++++++- ...fetcher-against-ssrf-029fcd8d99523a95.yaml | 10 + .../fetchers/test_link_content_fetcher.py | 515 +++++++++++++++--- test/dataclasses/test_file_content.py | 48 +- test/dataclasses/test_image_content.py | 52 +- 5 files changed, 782 insertions(+), 118 deletions(-) create mode 100644 releasenotes/notes/harden-link-content-fetcher-against-ssrf-029fcd8d99523a95.yaml diff --git a/haystack/components/fetchers/link_content.py b/haystack/components/fetchers/link_content.py index 5ff092377c0..b15fb6e2e1d 100644 --- a/haystack/components/fetchers/link_content.py +++ b/haystack/components/fetchers/link_content.py @@ -3,6 +3,8 @@ # SPDX-License-Identifier: Apache-2.0 import asyncio +import ipaddress +import socket from collections import defaultdict from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor @@ -26,6 +28,13 @@ DEFAULT_USER_AGENT = f"haystack/LinkContentFetcher/{__version__}" +DEFAULT_MAX_RESPONSE_BYTES = 50 * 1024 * 1024 +DEFAULT_MAX_REDIRECTS = 5 + +# IPv4 shared address space (RFC 6598), commonly used for internal networks and VPN overlays like Tailscale. +# `ipaddress` does not classify it as private, so it's checked explicitly. +_SHARED_ADDRESS_SPACE = ipaddress.ip_network("100.64.0.0/10") + REQUEST_HEADERS = { "accept": "*/*", "User-Agent": DEFAULT_USER_AGENT, @@ -34,6 +43,82 @@ } +class UnsafeTargetError(ValueError): + """ + Raised when a URL or one of its redirect targets points to a host that must not be fetched from. + + This includes hosts outside the ``allowed_hosts`` whitelist and hosts that resolve to a private, loopback, + link-local, multicast, or otherwise internal IP address. + """ + + +class ResponseTooLargeError(ValueError): + """ + Raised when a response body exceeds the fetcher's ``max_response_bytes`` limit. + """ + + +def _resolve_host(host: str, port: int) -> list[str]: + """ + Resolves a host name to all of its IPv4/IPv6 addresses (A/AAAA records). + + :param host: The host name or IP literal to resolve. + :param port: The port to resolve for. It does not affect the returned addresses. + :returns: The resolved addresses as strings. An empty list if resolution fails, so that the underlying + HTTP client can surface its own (retryable) connection error. + """ + try: + addr_infos = socket.getaddrinfo(host, port, family=socket.AF_UNSPEC, type=socket.SOCK_STREAM) + except (socket.gaierror, OSError): + return [] + return [addr_info[4][0] for addr_info in addr_infos] + + +def _is_forbidden_ip(address: str) -> bool: + """ + Checks whether an IP address belongs to a range the fetcher must never connect to. + + This covers private, loopback, link-local, multicast, reserved, unspecified, and unique-local addresses, as + well as the RFC 6598 shared address space. Unparseable addresses are rejected as well. + + :param address: The IP address to check, as a string. + :returns: `True` if the address is forbidden. + """ + try: + ip = ipaddress.ip_address(address) + except ValueError: + return True + return ( + ip.is_private + or ip.is_loopback + or ip.is_link_local + or ip.is_multicast + or ip.is_reserved + or ip.is_unspecified + or ip in _SHARED_ADDRESS_SPACE + ) + + +def _host_matches_allowlist(host: str, allowed_hosts: list[str]) -> bool: + """ + Checks whether a host matches any entry of a domain suffix whitelist. + + A host matches an entry if it is equal to it or a subdomain of it. For example, `api.example.com` matches + `example.com` but `notexample.com` does not. Comparison is case-insensitive and ignores a trailing dot + (the DNS root). + + :param host: The host name to check. + :param allowed_hosts: The whitelist entries. + :returns: `True` if the host matches at least one entry. + """ + normalized_host = host.lower().rstrip(".") + for allowed_host in allowed_hosts: + normalized_allowed = allowed_host.lower().rstrip(".") + if normalized_host == normalized_allowed or normalized_host.endswith(f".{normalized_allowed}"): + return True + return False + + def _merge_headers(*args: dict[str, str]) -> dict[str, str]: """ Merge a list of dict using case-insensitively @@ -81,6 +166,11 @@ class LinkContentFetcher: It supports various content types, retries on failures, and automatic user-agent rotation for failed web requests. Use it as the data-fetching step in your pipelines. + For security, every request target is validated before the request is made, including each redirect hop: + hosts can be restricted to an `allowed_hosts` domain suffix whitelist, and hosts resolving to private, + loopback, link-local, multicast, or otherwise internal IP addresses are rejected. Response bodies are + streamed and capped at `max_response_bytes`. + You may need to convert LinkContentFetcher's output into a list of documents. Use HTMLToDocument converter to do this. @@ -121,6 +211,9 @@ def __init__( http2: bool = False, client_kwargs: dict | None = None, request_headers: dict[str, str] | None = None, + allowed_hosts: list[str] | None = None, + max_response_bytes: int = DEFAULT_MAX_RESPONSE_BYTES, + max_redirects: int = DEFAULT_MAX_REDIRECTS, ) -> None: """ Initializes the component. @@ -135,8 +228,20 @@ def __init__( Requires the 'h2' package to be installed (via `pip install httpx[http2]`). :param client_kwargs: Additional keyword arguments to pass to the httpx client. If `None`, default values are used. + `follow_redirects` is always overridden to `False`: redirects are followed manually, + one hop at a time, so that every hop can be validated. To disable redirect following + entirely, pass `{"follow_redirects": False}`. :param request_headers: Additional headers to send with every request. These take precedence over the component's default headers but not over the rotating `User-Agent`. + :param allowed_hosts: Optional whitelist of allowed domain suffixes, for example + `["example.com", "cdn.example.org"]`. A host is allowed if it equals one of the entries + or is a subdomain of one (so `api.example.com` matches `example.com`, but + `notexample.com` does not). If `None`, any host is allowed. Hosts are additionally + always checked against forbidden IP ranges, regardless of this whitelist. + :param max_response_bytes: Maximum size in bytes of a response body. Responses are streamed and + bodies exceeding this limit raise a `ResponseTooLargeError`. + :param max_redirects: Maximum number of redirects to follow. Each hop is validated before the + request is made. """ self.raise_on_failure = raise_on_failure self.user_agents = user_agents or [DEFAULT_USER_AGENT] @@ -145,10 +250,16 @@ def __init__( self.http2 = http2 self.client_kwargs = client_kwargs or {} self.request_headers = request_headers or {} + self.allowed_hosts = allowed_hosts + self.max_response_bytes = max_response_bytes + self.max_redirects = max_redirects # Configure default client settings self.client_kwargs.setdefault("timeout", timeout) - self.client_kwargs.setdefault("follow_redirects", True) + # Redirects are followed manually (see `_request_following_redirects`) so that every hop can be + # validated before the request is made; the underlying client has `follow_redirects` forced to `False` + # in `_build_client_kwargs`. Here we only remember whether the user wants redirects followed at all. + self._follow_redirects = self.client_kwargs.get("follow_redirects", True) # httpx clients are built lazily in warm_up / warm_up_async (resource lifecycle) self._client: httpx.Client | None = None @@ -191,12 +302,103 @@ def rotate_user_agent(retry_state: RetryCallState) -> None: # noqa: ARG001 ) def get_response(url: str) -> httpx.Response: assert self._client is not None # mypy: client is built by warm_up before run - response = self._client.get(url, headers=self._get_headers(self.user_agents[user_agent_idx])) - response.raise_for_status() - return response + headers = self._get_headers(self.user_agents[user_agent_idx]) + return self._request_following_redirects(self._client, url, headers) return get_response(url) + def _validate_target(self, url: str) -> None: + """ + Validates that a URL may be fetched before the request is made. + + The URL's host must match the `allowed_hosts` whitelist (if set) and must not resolve to any forbidden + IP address (private, loopback, link-local, multicast, reserved, unspecified, unique-local, or shared + address space ranges). Note that this mitigates, but cannot fully eliminate, DNS rebinding: the HTTP + client performs its own name resolution when connecting. + + :param url: The URL to validate. + :raises UnsafeTargetError: If the URL's host is not allowed or resolves to a forbidden address. + """ + parsed = httpx.URL(url) + host = parsed.host + if not host: + raise UnsafeTargetError(f"URL '{url}' has no host") + + if self.allowed_hosts is not None and not _host_matches_allowlist(host, self.allowed_hosts): + raise UnsafeTargetError( + f"Host '{host}' (URL '{url}') is not in the allowed_hosts whitelist: {self.allowed_hosts}" + ) + + port = parsed.port or (443 if parsed.scheme == "https" else 80) + for address in _resolve_host(host, port): + if _is_forbidden_ip(address): + raise UnsafeTargetError( + f"Host '{host}' (URL '{url}') resolves to the forbidden IP address '{address}'. " + "Requests to private, loopback, link-local, multicast, or other internal addresses " + "are not allowed." + ) + + def _request_following_redirects(self, client: httpx.Client, url: str, headers: dict[str, str]) -> httpx.Response: + """ + Performs the request for a URL, following redirects manually one hop at a time. + + The HTTP client itself never follows redirects (see `_build_client_kwargs`): every hop, starting from + the original URL, is validated before its request is made. This prevents a redirect from a trusted host + to an internal address from ever being fetched. + + :param client: The httpx client to make the requests with. + :param url: The URL to fetch. + :param headers: Headers to send with every hop. + :returns: The final httpx Response object, with its body already read. + :raises httpx.TooManyRedirects: If more than `max_redirects` redirects are encountered. + """ + current_url = url + for _hop in range(self.max_redirects + 1): + self._validate_target(current_url) + response = self._request_hop(client, current_url, headers) + if not (self._follow_redirects and response.is_redirect): + response.raise_for_status() + return response + current_url = str(httpx.URL(current_url).join(response.headers["location"])) + raise httpx.TooManyRedirects(f"Exceeded {self.max_redirects} redirects.", request=httpx.Request("GET", url)) + + def _request_hop(self, client: httpx.Client, url: str, headers: dict[str, str]) -> httpx.Response: + """ + Performs a single GET request, streaming the response body with a size cap. + + The body is read while the connection is still open so that the transfer is aborted as soon as it + exceeds `max_response_bytes`. A plain, fully-read response is returned so that content handlers can + access it like a regular, non-streaming response. + + Redirect responses are returned without reading their body: it is never used. + + :param client: The httpx client to make the request with. + :param url: The URL to fetch. + :param headers: Headers to send with the request. + :returns: The httpx Response object for this hop, with its body read unless it is a redirect. + :raises ResponseTooLargeError: If the response body exceeds `max_response_bytes`. + """ + with client.stream("GET", url, headers=headers) as response: + body = b"" + if not (self._follow_redirects and response.is_redirect): + chunks = bytearray() + for chunk in response.iter_bytes(): + chunks.extend(chunk) + if len(chunks) > self.max_response_bytes: + raise ResponseTooLargeError( + f"Response from '{url}' exceeds max_response_bytes={self.max_response_bytes} " + f"(received at least {len(chunks)} bytes)." + ) + body = bytes(chunks) + # Rebuild a fully-read response so handlers can use `response.text`/`response.content` as usual. + return httpx.Response( + status_code=response.status_code, + headers=response.headers, + content=body, + request=response.request, + default_encoding=response.default_encoding, + ) + def _build_client_kwargs(self) -> dict[str, Any]: """ Build the keyword arguments used to construct the httpx clients. @@ -205,6 +407,10 @@ def _build_client_kwargs(self) -> dict[str, Any]: """ client_kwargs = {**self.client_kwargs} + # Redirects are followed manually, one validated hop at a time (see `_request_following_redirects`), + # so the underlying client must never follow them on its own. + client_kwargs["follow_redirects"] = False + # Optional HTTP/2 support if self.http2: try: @@ -444,7 +650,8 @@ async def _get_response_async(self, url: str, client: httpx.AsyncClient) -> http while attempt <= self.retry_attempts: try: - response = await client.get(url, headers=self._get_headers(self.user_agents[user_agent_idx])) + headers = self._get_headers(self.user_agents[user_agent_idx]) + response = await self._request_following_redirects_async(client, url, headers) response.raise_for_status() return response except (httpx.HTTPStatusError, httpx.RequestError) as e: @@ -465,6 +672,64 @@ async def _get_response_async(self, url: str, client: httpx.AsyncClient) -> http # This should never happen, but just in case raise httpx.RequestError("Failed to get response after retries", request=None) + async def _request_following_redirects_async( + self, client: httpx.AsyncClient, url: str, headers: dict[str, str] + ) -> httpx.Response: + """ + Asynchronously performs the request for a URL, following redirects manually one hop at a time. + + Every hop, starting from the original URL, is validated before its request is made. This is the + asynchronous counterpart of `_request_following_redirects`. + + :param client: The async httpx client to make the requests with. + :param url: The URL to fetch. + :param headers: Headers to send with every hop. + :returns: The final httpx Response object, with its body already read. + :raises httpx.TooManyRedirects: If more than `max_redirects` redirects are encountered. + """ + current_url = url + for _hop in range(self.max_redirects + 1): + self._validate_target(current_url) + response = await self._request_hop_async(client, current_url, headers) + if not (self._follow_redirects and response.is_redirect): + response.raise_for_status() + return response + current_url = str(httpx.URL(current_url).join(response.headers["location"])) + raise httpx.TooManyRedirects(f"Exceeded {self.max_redirects} redirects.", request=httpx.Request("GET", url)) + + async def _request_hop_async(self, client: httpx.AsyncClient, url: str, headers: dict[str, str]) -> httpx.Response: + """ + Asynchronously performs a single GET request, streaming the response body with a size cap. + + This is the asynchronous counterpart of `_request_hop`. + + :param client: The async httpx client to make the request with. + :param url: The URL to fetch. + :param headers: Headers to send with the request. + :returns: The httpx Response object for this hop, with its body read unless it is a redirect. + :raises ResponseTooLargeError: If the response body exceeds `max_response_bytes`. + """ + async with client.stream("GET", url, headers=headers) as response: + body = b"" + if not (self._follow_redirects and response.is_redirect): + chunks = bytearray() + async for chunk in response.aiter_bytes(): + chunks.extend(chunk) + if len(chunks) > self.max_response_bytes: + raise ResponseTooLargeError( + f"Response from '{url}' exceeds max_response_bytes={self.max_response_bytes} " + f"(received at least {len(chunks)} bytes)." + ) + body = bytes(chunks) + # Rebuild a fully-read response so handlers can use `response.text`/`response.content` as usual. + return httpx.Response( + status_code=response.status_code, + headers=response.headers, + content=body, + request=response.request, + default_encoding=response.default_encoding, + ) + def _get_content_type(self, response: httpx.Response) -> str: """ Get the content type of the response. diff --git a/releasenotes/notes/harden-link-content-fetcher-against-ssrf-029fcd8d99523a95.yaml b/releasenotes/notes/harden-link-content-fetcher-against-ssrf-029fcd8d99523a95.yaml new file mode 100644 index 00000000000..49df30fda7d --- /dev/null +++ b/releasenotes/notes/harden-link-content-fetcher-against-ssrf-029fcd8d99523a95.yaml @@ -0,0 +1,10 @@ +fixes: + - | + Hardened ``LinkContentFetcher`` against server-side request forgery (SSRF). Every request target is now + validated before the request is made, including each redirect hop: redirects are followed manually with a + per-hop re-validation and a cap of ``max_redirects`` (default 5) hops, so a redirect from a trusted host to + an internal address is never fetched. Hosts can be restricted to a new ``allowed_hosts`` domain suffix + whitelist, and hosts resolving to private, loopback, link-local, multicast, reserved, or unique-local IP + addresses are always rejected. In addition, response bodies are now streamed and capped at + ``max_response_bytes`` (default 50 MB): larger responses raise an error instead of being fully loaded into + memory. diff --git a/test/components/fetchers/test_link_content_fetcher.py b/test/components/fetchers/test_link_content_fetcher.py index 1869dbbf73a..950697ae743 100644 --- a/test/components/fetchers/test_link_content_fetcher.py +++ b/test/components/fetchers/test_link_content_fetcher.py @@ -3,40 +3,97 @@ # SPDX-License-Identifier: Apache-2.0 import threading -from unittest.mock import AsyncMock, Mock, patch +from contextlib import nullcontext +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest from tenacity import wait_none from haystack.components.fetchers.link_content import ( + DEFAULT_MAX_REDIRECTS, + DEFAULT_MAX_RESPONSE_BYTES, DEFAULT_USER_AGENT, LinkContentFetcher, + ResponseTooLargeError, + UnsafeTargetError, _binary_content_handler, _text_content_handler, ) +from haystack.core.serialization import component_from_dict, component_to_dict HTML_URL = "https://docs.haystack.deepset.ai/docs/intro" TEXT_URL = "https://raw.githubusercontent.com/deepset-ai/haystack/main/README.md" PDF_URL = "https://raw.githubusercontent.com/deepset-ai/haystack/b5987a6d8d0714eb2f3011183ab40093d2e4a41a/e2e/samples/pipelines/sample_pdf_1.pdf" +PUBLIC_IP = "93.184.216.34" + + +def make_response( + url: str, + status_code: int = 200, + text: str | None = None, + content: bytes | None = None, + headers: dict[str, str] | None = None, +) -> httpx.Response: + """ + Builds a real `httpx.Response` suitable for replaying from a mocked `Client.stream` call. + """ + return httpx.Response(status_code, text=text, content=content, headers=headers, request=httpx.Request("GET", url)) + + +def stream_side_effect(responses: list) -> "MagicMock": + """ + Creates a `side_effect` for a mocked `httpx.Client.stream` that replays the given responses in order. + + Entries may be `httpx.Response` objects (replayed through a context manager) or exceptions (raised). + The last entry is repeated once the list is exhausted. + """ + state = {"idx": 0} + + def fake_stream(method: str, url: str, headers: dict | None = None, **kwargs): # noqa: ARG001 + idx = state["idx"] + state["idx"] = min(idx + 1, len(responses) - 1) + item = responses[idx] + if isinstance(item, Exception): + raise item + return nullcontext(item) + + return fake_stream + + +@pytest.fixture(autouse=True) +def mock_dns_resolution(request): + """ + Points all name resolution performed by the fetcher at a public IP so unit tests never touch the network. + + Tests that exercise the IP validation install their own patch on top of this one. + """ + if request.node.get_closest_marker("integration"): + yield + return + with patch("haystack.components.fetchers.link_content._resolve_host", return_value=[PUBLIC_IP]) as mocked: + yield mocked + @pytest.fixture def mock_get_link_text_content(): - with patch("haystack.components.fetchers.link_content.httpx.Client.get") as mock_get: - mock_response = Mock(status_code=200, text="Example test response", headers={"Content-Type": "text/plain"}) - mock_get.return_value = mock_response - yield mock_get + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect( + [make_response("https://www.example.com", text="Example test response")] + ) + yield mock_stream @pytest.fixture def mock_get_link_content(test_files_path): - with patch("haystack.components.fetchers.link_content.httpx.Client.get") as mock_get: + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: with open(test_files_path / "pdf" / "sample_pdf_1.pdf", "rb") as f1: file_bytes = f1.read() - mock_response = Mock(status_code=200, content=file_bytes, headers={"Content-Type": "application/pdf"}) - mock_get.return_value = mock_response - yield mock_get + mock_stream.side_effect = stream_side_effect( + [make_response("https://www.example.com", content=file_bytes, headers={"Content-Type": "application/pdf"})] + ) + yield mock_stream class TestLinkContentFetcher: @@ -49,6 +106,9 @@ def test_init(self): assert fetcher.timeout == 3 assert fetcher.http2 is False assert isinstance(fetcher.client_kwargs, dict) + assert fetcher.allowed_hosts is None + assert fetcher.max_response_bytes == DEFAULT_MAX_RESPONSE_BYTES + assert fetcher.max_redirects == DEFAULT_MAX_REDIRECTS assert fetcher.handlers == { "text/*": _text_content_handler, "text/html": _binary_content_handler, @@ -83,9 +143,10 @@ def test_init_with_params(self): def test_run_text(self): """Test fetching text content""" correct_response = b"Example test response" - with patch("haystack.components.fetchers.link_content.httpx.Client.get") as mock_get: - mock_response = Mock(status_code=200, text="Example test response", headers={"Content-Type": "text/plain"}) - mock_get.return_value = mock_response + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect( + [make_response("https://www.example.com", text="Example test response")] + ) fetcher = LinkContentFetcher() streams = fetcher.run(urls=["https://www.example.com"])["streams"] first_stream = streams[0] @@ -96,11 +157,16 @@ def test_run_text(self): def test_run_html(self): """Test fetching HTML content""" correct_response = b"

Example test response

" - with patch("haystack.components.fetchers.link_content.httpx.Client.get") as mock_get: - mock_response = Mock( - status_code=200, content=b"

Example test response

", headers={"Content-Type": "text/html"} + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect( + [ + make_response( + "https://www.example.com", + content=b"

Example test response

", + headers={"Content-Type": "text/html"}, + ) + ] ) - mock_get.return_value = mock_response fetcher = LinkContentFetcher() streams = fetcher.run(urls=["https://www.example.com"])["streams"] first_stream = streams[0] @@ -112,9 +178,11 @@ def test_run_binary(self, test_files_path): """Test fetching binary content""" with open(test_files_path / "pdf" / "sample_pdf_1.pdf", "rb") as f1: file_bytes = f1.read() - with patch("haystack.components.fetchers.link_content.httpx.Client.get") as mock_get: - mock_response = Mock(status_code=200, content=file_bytes, headers={"Content-Type": "application/pdf"}) - mock_get.return_value = mock_response + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: + response = make_response( + "https://www.example.com", content=file_bytes, headers={"Content-Type": "application/pdf"} + ) + mock_stream.side_effect = stream_side_effect([response]) fetcher = LinkContentFetcher() streams = fetcher.run(urls=["https://www.example.com"])["streams"] first_stream = streams[0] @@ -126,13 +194,10 @@ def test_run_bad_request_no_exception(self): """Test behavior when a request results in an error status code""" empty_byte_stream = b"" fetcher = LinkContentFetcher(raise_on_failure=False, retry_attempts=0) - mock_response = Mock(status_code=403) - mock_response.raise_for_status.side_effect = httpx.HTTPStatusError( - "403 Client Error", request=Mock(), response=mock_response - ) + mock_response = make_response("https://www.example.com", status_code=403) - with patch("haystack.components.fetchers.link_content.httpx.Client.get") as mock_get: - mock_get.return_value = mock_response + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect([mock_response]) streams = fetcher.run(urls=["https://www.example.com"])["streams"] # empty byte stream is returned because raise_on_failure is False @@ -149,42 +214,38 @@ def test_bad_request_exception_raised(self): """ fetcher = LinkContentFetcher(raise_on_failure=True, retry_attempts=0) - mock_response = Mock(status_code=403) - mock_response.raise_for_status.side_effect = httpx.HTTPStatusError( - "403 Client Error", request=Mock(), response=mock_response - ) + mock_response = make_response("https://non_existent_website_dot.com/", status_code=403) - with patch("haystack.components.fetchers.link_content.httpx.Client.get") as mock_get: - mock_get.return_value = mock_response + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect([mock_response]) with pytest.raises(httpx.HTTPStatusError): fetcher.run(["https://non_existent_website_dot.com/"]) def test_run_retries_once_when_retry_attempts_is_one(self): url = "https://www.example.com" - successful_response = Mock(status_code=200, text="Success", headers={"Content-Type": "text/plain"}) + successful_response = make_response(url, text="Success") with patch("haystack.components.fetchers.link_content.httpx.Client") as client_mock: client = client_mock.return_value client.headers = {} - client.get.side_effect = [ - httpx.RequestError("transient failure", request=httpx.Request("GET", url)), - successful_response, - ] + client.stream.side_effect = stream_side_effect( + [httpx.RequestError("transient failure", request=httpx.Request("GET", url)), successful_response] + ) fetcher = LinkContentFetcher(retry_attempts=1) with patch("haystack.components.fetchers.link_content.wait_exponential", return_value=wait_none()): streams = fetcher.run(urls=[url])["streams"] assert streams[0].data == successful_response.text.encode() - assert client.get.call_count == 2 + assert client.stream.call_count == 2 def test_request_headers_merging_and_ua_override(self): # Patch the Client class to control the instance created by LinkContentFetcher with patch("haystack.components.fetchers.link_content.httpx.Client") as ClientMock: client = ClientMock.return_value client.headers = {} # base headers used in the merge - mock_response = Mock(status_code=200, text="OK", headers={"Content-Type": "text/plain"}) - client.get.return_value = mock_response + mock_response = make_response("https://example.com", text="OK") + client.stream.side_effect = stream_side_effect([mock_response]) fetcher = LinkContentFetcher( user_agents=["ua-sync-1", "ua-sync-2"], @@ -197,8 +258,8 @@ def test_request_headers_merging_and_ua_override(self): _ = fetcher.run(urls=["https://example.com"])["streams"] - client.get.assert_called_once() - sent_headers = client.get.call_args.kwargs["headers"] + client.stream.assert_called_once() + sent_headers = client.stream.call_args.kwargs["headers"] assert sent_headers["X-Test"] == "1" assert sent_headers["Accept-Language"] == "fr-FR" assert sent_headers["User-Agent"] == "ua-sync-1" # rotating UA wins @@ -217,7 +278,7 @@ def test_user_agent_rotation_is_independent_per_url(self): user_agent_on_success: dict[str, str] = {} lock = threading.Lock() - def fake_get(url, headers=None, **kwargs): + def fake_stream(method, url, headers=None, **kwargs): # noqa: ARG001 with lock: attempt = attempts.get(url, 0) attempts[url] = attempt + 1 @@ -226,12 +287,12 @@ def fake_get(url, headers=None, **kwargs): raise httpx.RequestError("simulated transient failure", request=httpx.Request("GET", url)) with lock: user_agent_on_success[url] = headers["User-Agent"] - return Mock(status_code=200, text="OK", headers={"Content-Type": "text/plain"}) + return nullcontext(make_response(url, text="OK")) with patch("haystack.components.fetchers.link_content.httpx.Client") as ClientMock: client = ClientMock.return_value client.headers = {} - client.get.side_effect = fake_get + client.stream.side_effect = fake_stream fetcher = LinkContentFetcher(user_agents=user_agents, retry_attempts=3, raise_on_failure=False) with patch("haystack.components.fetchers.link_content.wait_exponential", return_value=wait_none()): @@ -324,18 +385,16 @@ async def test_close_and_close_async_are_independent(self): client_instance.close.assert_called_once() def test_run_self_heals(self): - with patch("haystack.components.fetchers.link_content.httpx.Client.get") as mock_get: - mock_response = Mock(status_code=200, text="ok", headers={"Content-Type": "text/plain"}) - mock_get.return_value = mock_response + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect([make_response("https://www.example.com", text="ok")]) fetcher = LinkContentFetcher() fetcher.run(urls=["https://www.example.com"]) assert fetcher._client is not None @pytest.mark.asyncio async def test_run_async_self_heals(self): - with patch("haystack.components.fetchers.link_content.httpx.AsyncClient.get") as mock_get: - mock_response = Mock(status_code=200, text="ok", headers={"Content-Type": "text/plain"}) - mock_get.return_value = mock_response + with patch("haystack.components.fetchers.link_content.httpx.AsyncClient.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect([make_response("https://www.example.com", text="ok")]) fetcher = LinkContentFetcher() await fetcher.run_async(urls=["https://www.example.com"]) assert fetcher._async_client is not None @@ -420,9 +479,10 @@ def test_mix_of_good_and_failed_requests(self): class TestLinkContentFetcherAsync: async def test_run_async(self): """Test basic async fetching with a mocked response""" - with patch("haystack.components.fetchers.link_content.httpx.AsyncClient.get") as mock_get: - mock_response = Mock(status_code=200, text="Example test response", headers={"Content-Type": "text/plain"}) - mock_get.return_value = mock_response + with patch("haystack.components.fetchers.link_content.httpx.AsyncClient.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect( + [make_response("https://www.example.com", text="Example test response")] + ) fetcher = LinkContentFetcher() streams = (await fetcher.run_async(urls=["https://www.example.com"]))["streams"] @@ -435,9 +495,10 @@ async def test_run_async(self): async def test_run_async_multiple(self): """Test async fetching of multiple URLs with mocked responses""" - with patch("haystack.components.fetchers.link_content.httpx.AsyncClient.get") as mock_get: - mock_response = Mock(status_code=200, text="Example test response", headers={"Content-Type": "text/plain"}) - mock_get.return_value = mock_response + with patch("haystack.components.fetchers.link_content.httpx.AsyncClient.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect( + [make_response("https://www.example.com", text="Example test response")] + ) fetcher = LinkContentFetcher() streams = (await fetcher.run_async(urls=["https://www.example1.com", "https://www.example2.com"]))[ @@ -459,12 +520,8 @@ async def test_run_async_empty_urls(self): async def test_run_async_error_handling(self): """Test error handling for async fetching""" - with patch("haystack.components.fetchers.link_content.httpx.AsyncClient.get") as mock_get: - mock_response = Mock(status_code=404) - mock_response.raise_for_status.side_effect = httpx.HTTPStatusError( - "404 Not Found", request=Mock(), response=mock_response - ) - mock_get.return_value = mock_response + with patch("haystack.components.fetchers.link_content.httpx.AsyncClient.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect([make_response("https://www.example.com", status_code=404)]) # With raise_on_failure=False fetcher = LinkContentFetcher(raise_on_failure=False, retry_attempts=0) @@ -479,23 +536,22 @@ async def test_run_async_error_handling(self): async def test_run_async_user_agent_rotation(self): """Test user agent rotation in async fetching""" with ( - patch("haystack.components.fetchers.link_content.httpx.AsyncClient.get") as mock_get, + patch("haystack.components.fetchers.link_content.httpx.AsyncClient.stream") as mock_stream, patch("asyncio.sleep") as mock_sleep, ): # Mock asyncio.sleep used by tenacity to keep this test fast mock_sleep.return_value = None # First call raises an error to trigger user agent rotation - first_response = Mock(status_code=403) - first_response.raise_for_status.side_effect = httpx.HTTPStatusError( - "403 Forbidden", request=Mock(), response=first_response - ) + first_response = make_response("https://www.example.com", status_code=403) # Second call succeeds - second_response = Mock(status_code=200, text="Success", headers={"Content-Type": "text/plain"}) + second_response = make_response( + "https://www.example.com", text="Success", headers={"Content-Type": "text/plain"} + ) # Use side_effect to return different responses on consecutive calls - mock_get.side_effect = [first_response, second_response] + mock_stream.side_effect = stream_side_effect([first_response, second_response]) # Create fetcher with custom user agents fetcher = LinkContentFetcher(user_agents=["agent1", "agent2"], retry_attempts=1) @@ -514,8 +570,9 @@ async def test_request_headers_merging_and_ua_override(self): aclient = AsyncClientMock.return_value aclient.headers = {} # base headers used in the merge - mock_response = Mock(status_code=200, text="OK", headers={"Content-Type": "text/plain"}) - aclient.get = AsyncMock(return_value=mock_response) + mock_response = make_response("https://example.com", text="OK") + aclient.stream = MagicMock() + aclient.stream.side_effect = stream_side_effect([mock_response]) fetcher = LinkContentFetcher( user_agents=["ua-async-1", "ua-async-2"], @@ -524,8 +581,8 @@ async def test_request_headers_merging_and_ua_override(self): _ = (await fetcher.run_async(urls=["https://example.com"]))["streams"] - assert aclient.get.await_count == 1 - sent_headers = aclient.get.call_args.kwargs["headers"] + assert aclient.stream.call_count == 1 + sent_headers = aclient.stream.call_args.kwargs["headers"] assert sent_headers["X-Async"] == "true" assert sent_headers["Accept-Language"] == "de-DE" assert sent_headers["User-Agent"] == "ua-async-1" # rotating UA wins @@ -536,8 +593,9 @@ async def test_duplicated_request_headers_merging(self): aclient = AsyncClientMock.return_value aclient.headers = {} # base headers used in the merge - mock_response = Mock(status_code=200, text="OK", headers={"Content-Type": "text/plain"}) - aclient.get = AsyncMock(return_value=mock_response) + mock_response = make_response("https://example.com", text="OK") + aclient.stream = MagicMock() + aclient.stream.side_effect = stream_side_effect([mock_response]) fetcher = LinkContentFetcher( request_headers={ @@ -550,8 +608,8 @@ async def test_duplicated_request_headers_merging(self): _ = (await fetcher.run_async(urls=["https://example.com"]))["streams"] - assert aclient.get.await_count == 1 - sent_headers = aclient.get.call_args.kwargs["headers"] + assert aclient.stream.call_count == 1 + sent_headers = aclient.stream.call_args.kwargs["headers"] existing_keys = {} for key, value in sent_headers.items(): lower_key = key.lower() @@ -565,6 +623,309 @@ async def test_duplicated_request_headers_merging(self): assert existing_keys["x-test-header"] == "X-TeSt-HeAdEr" +class TestLinkContentFetcherSecurity: + """ + Tests for the SSRF protections: the `allowed_hosts` whitelist, rejection of forbidden IP ranges, + per-hop redirect validation, and the response size cap. + """ + + @pytest.mark.parametrize( + "forbidden_ip", + [ + "10.0.0.5", + "192.168.1.10", + "172.16.0.1", + "127.0.0.1", + "169.254.169.254", + "224.0.0.1", + "::1", + "fd00::1", + "fe80::1", + "100.64.0.1", + ], + ) + def test_forbidden_ip_targets_are_rejected(self, mock_dns_resolution, forbidden_ip): + """Hosts resolving to private/internal ranges must be rejected before any request is made.""" + mock_dns_resolution.return_value = [forbidden_ip] + fetcher = LinkContentFetcher(retry_attempts=0) + + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: + with pytest.raises(UnsafeTargetError, match="forbidden IP address"): + fetcher.run(["https://internal.example.com/data"]) + + mock_stream.assert_not_called() + + @pytest.mark.parametrize("url", ["http://127.0.0.1:8080/admin", "http://10.0.0.5/", "http://[::1]/"]) + def test_ip_literal_urls_are_rejected(self, mock_dns_resolution, url): + """URLs pointing directly at an internal IP literal must be rejected. + + Name resolution echoes IP literals back unchanged, like a real resolver does. + """ + mock_dns_resolution.side_effect = lambda host, port: [host] + fetcher = LinkContentFetcher(retry_attempts=0) + + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: + with pytest.raises(UnsafeTargetError, match="forbidden IP address"): + fetcher.run([url]) + + mock_stream.assert_not_called() + + def test_host_resolving_to_mixed_addresses_is_rejected(self, mock_dns_resolution): + """If any of the resolved addresses is forbidden, the whole host is rejected.""" + mock_dns_resolution.return_value = ["93.184.216.34", "192.168.0.1"] + fetcher = LinkContentFetcher(retry_attempts=0) + + with pytest.raises(UnsafeTargetError, match="192.168.0.1"): + fetcher.run(["https://example.com"]) + + def test_allowed_hosts_whitelist(self, mock_dns_resolution): + """Hosts outside the whitelist are rejected; subdomains of allowed entries are accepted.""" + mock_dns_resolution.return_value = [PUBLIC_IP] + fetcher = LinkContentFetcher(allowed_hosts=["example.com"], retry_attempts=0) + expected_data = b"ok" + + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect([make_response("https://api.example.com", text="ok")]) + streams = fetcher.run(["https://api.example.com"])["streams"] + assert streams[0].data == expected_data + + with pytest.raises(UnsafeTargetError, match="whitelist"): + fetcher.run(["https://api.other.com"]) + + # the suffix match must respect domain boundaries + with pytest.raises(UnsafeTargetError, match="whitelist"): + fetcher.run(["https://notexample.com"]) + + def test_no_allowlist_keeps_previous_behavior(self, mock_dns_resolution): + """With the default configuration (no allowlist), any host resolving to a public IP can be fetched.""" + mock_dns_resolution.return_value = [PUBLIC_IP] + fetcher = LinkContentFetcher() + expected_data = b"ok" + + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect([make_response("https://anything.example.org", text="ok")]) + streams = fetcher.run(["https://anything.example.org"])["streams"] + assert streams[0].data == expected_data + + def test_redirect_to_forbidden_target_is_rejected(self, mock_dns_resolution): + """A redirect from a trusted host to an internal address must never be fetched.""" + mock_dns_resolution.side_effect = lambda host, port: ( + ["10.0.0.5"] if host == "internal.example.net" else [PUBLIC_IP] + ) + fetcher = LinkContentFetcher(retry_attempts=0) + + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect( + [ + make_response( + "https://www.example.com", + status_code=302, + headers={"location": "http://internal.example.net/secret"}, + ) + ] + ) + with pytest.raises(UnsafeTargetError, match="forbidden IP address"): + fetcher.run(["https://www.example.com"]) + + # only the first, legitimate hop was requested + assert mock_stream.call_count == 1 + + def test_redirect_to_ip_literal_is_rejected(self, mock_dns_resolution): + """A redirect pointing directly at an internal IP literal must never be fetched.""" + mock_dns_resolution.side_effect = lambda host, port: [host] if host == "127.0.0.1" else [PUBLIC_IP] + fetcher = LinkContentFetcher(retry_attempts=0) + + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect( + [ + make_response( + "https://www.example.com", status_code=301, headers={"location": "http://127.0.0.1:8080/"} + ) + ] + ) + with pytest.raises(UnsafeTargetError, match="forbidden IP address"): + fetcher.run(["https://www.example.com"]) + + assert mock_stream.call_count == 1 + + def test_redirects_are_followed_and_validated(self, mock_dns_resolution): + """Legitimate redirects, including relative ones, are followed hop by hop.""" + mock_dns_resolution.return_value = [PUBLIC_IP] + fetcher = LinkContentFetcher(retry_attempts=0) + expected_data = b"final content" + + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect( + [ + make_response("https://www.example.com/a", status_code=302, headers={"location": "/b"}), + make_response( + "https://www.example.com/b", status_code=301, headers={"location": "https://cdn.example.org/c"} + ), + make_response("https://cdn.example.org/c", text="final content"), + ] + ) + streams = fetcher.run(["https://www.example.com/a"])["streams"] + + assert streams[0].data == expected_data + assert streams[0].meta["url"] == "https://www.example.com/a" + assert mock_stream.call_count == 3 + # every hop went through the same validated request path + requested_urls = [call.args[1] for call in mock_stream.call_args_list] + assert requested_urls == [ + "https://www.example.com/a", + "https://www.example.com/b", + "https://cdn.example.org/c", + ] + + def test_too_many_redirects(self, mock_dns_resolution): + """Redirect chains longer than `max_redirects` raise `httpx.TooManyRedirects`.""" + mock_dns_resolution.return_value = [PUBLIC_IP] + fetcher = LinkContentFetcher(max_redirects=2, retry_attempts=0) + + redirects = [ + make_response( + f"https://www.example.com/{i}", + status_code=302, + headers={"location": f"https://www.example.com/{i + 1}"}, + ) + for i in range(10) + ] + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect(redirects) + with pytest.raises(httpx.TooManyRedirects, match="Exceeded 2 redirects"): + fetcher.run(["https://www.example.com/0"]) + + # the original request plus `max_redirects` followed hops + assert mock_stream.call_count == 3 + + def test_redirects_not_followed_when_disabled_in_client_kwargs(self, mock_dns_resolution): + """`follow_redirects=False` in `client_kwargs` keeps the pre-existing behavior: the redirect is not + followed and its response surfaces like an error, exactly as httpx did before.""" + mock_dns_resolution.return_value = [PUBLIC_IP] + fetcher = LinkContentFetcher(client_kwargs={"follow_redirects": False}, retry_attempts=0) + + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect( + [ + make_response( + "https://www.example.com", status_code=302, content=b"moved", headers={"location": "/b"} + ) + ] + ) + with pytest.raises(httpx.HTTPStatusError, match="302"): + fetcher.run(["https://www.example.com"]) + + assert mock_stream.call_count == 1 + + @pytest.mark.parametrize("raise_on_failure", [True, False]) + def test_response_size_limit(self, mock_dns_resolution, raise_on_failure): + """Response bodies larger than `max_response_bytes` raise `ResponseTooLargeError`.""" + mock_dns_resolution.return_value = [PUBLIC_IP] + fetcher = LinkContentFetcher(max_response_bytes=10, retry_attempts=0, raise_on_failure=raise_on_failure) + expected_empty = b"" + + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect([make_response("https://www.example.com", content=b"a" * 11)]) + if raise_on_failure: + with pytest.raises(ResponseTooLargeError, match="exceeds max_response_bytes=10"): + fetcher.run(["https://www.example.com"]) + else: + # failures are swallowed and an empty stream is returned, like any other fetch error + streams = fetcher.run(["https://www.example.com"])["streams"] + assert len(streams) == 1 + assert streams[0].data == expected_empty + + def test_response_at_size_limit_is_fetched(self, mock_dns_resolution): + """A body exactly of `max_response_bytes` is fetched successfully.""" + mock_dns_resolution.return_value = [PUBLIC_IP] + fetcher = LinkContentFetcher(max_response_bytes=10, retry_attempts=0) + + with patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect([make_response("https://www.example.com", content=b"a" * 10)]) + streams = fetcher.run(["https://www.example.com"])["streams"] + assert streams[0].data == b"a" * 10 + + def test_serialization_includes_new_init_parameters(self): + fetcher = LinkContentFetcher() + data = component_to_dict(fetcher, "fetcher") + assert data["init_parameters"]["allowed_hosts"] is None + assert data["init_parameters"]["max_response_bytes"] == DEFAULT_MAX_RESPONSE_BYTES + assert data["init_parameters"]["max_redirects"] == DEFAULT_MAX_REDIRECTS + # the manual redirect following is an internal detail: it must not leak into the serialized state + assert data["init_parameters"]["client_kwargs"] == {"timeout": 3} + + fetcher = LinkContentFetcher(allowed_hosts=["example.com"], max_response_bytes=1024, max_redirects=2) + data = component_to_dict(fetcher, "fetcher") + assert data["init_parameters"]["allowed_hosts"] == ["example.com"] + assert data["init_parameters"]["max_response_bytes"] == 1024 + assert data["init_parameters"]["max_redirects"] == 2 + + def test_client_kwargs_serialization_round_trip(self, mock_dns_resolution): + """A fetcher serialized and loaded back keeps following redirects.""" + fetcher = LinkContentFetcher() + data = component_to_dict(fetcher, "fetcher") + restored = component_from_dict(LinkContentFetcher, data, "fetcher") + + assert restored._follow_redirects is True + assert restored.client_kwargs.get("timeout") == 3 + + +@pytest.mark.asyncio +class TestLinkContentFetcherAsyncSecurity: + async def test_run_async_rejects_forbidden_ip(self, mock_dns_resolution): + mock_dns_resolution.return_value = ["10.0.0.5"] + fetcher = LinkContentFetcher(retry_attempts=0) + + with patch("haystack.components.fetchers.link_content.httpx.AsyncClient.stream") as mock_stream: + with pytest.raises(UnsafeTargetError, match="forbidden IP address"): + await fetcher.run_async(["https://internal.example.com/data"]) + + mock_stream.assert_not_called() + + async def test_run_async_redirect_to_forbidden_target_is_rejected(self, mock_dns_resolution): + mock_dns_resolution.side_effect = lambda host, port: ["192.168.0.1"] if host == "192.168.0.1" else [PUBLIC_IP] + fetcher = LinkContentFetcher(retry_attempts=0) + + with patch("haystack.components.fetchers.link_content.httpx.AsyncClient.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect( + [ + make_response( + "https://www.example.com", status_code=302, headers={"location": "http://192.168.0.1/admin"} + ) + ] + ) + with pytest.raises(UnsafeTargetError, match="forbidden IP address"): + await fetcher.run_async(["https://www.example.com"]) + + assert mock_stream.call_count == 1 + + async def test_run_async_response_size_limit(self, mock_dns_resolution): + mock_dns_resolution.return_value = [PUBLIC_IP] + fetcher = LinkContentFetcher(max_response_bytes=10, retry_attempts=0) + + with patch("haystack.components.fetchers.link_content.httpx.AsyncClient.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect([make_response("https://www.example.com", content=b"a" * 20)]) + with pytest.raises(ResponseTooLargeError, match="exceeds max_response_bytes=10"): + await fetcher.run_async(["https://www.example.com"]) + + async def test_run_async_follows_validated_redirects(self, mock_dns_resolution): + mock_dns_resolution.return_value = [PUBLIC_IP] + fetcher = LinkContentFetcher(retry_attempts=0) + expected_data = b"final content" + + with patch("haystack.components.fetchers.link_content.httpx.AsyncClient.stream") as mock_stream: + mock_stream.side_effect = stream_side_effect( + [ + make_response("https://www.example.com/a", status_code=302, headers={"location": "/b"}), + make_response("https://www.example.com/b", text="final content"), + ] + ) + streams = (await fetcher.run_async(["https://www.example.com/a"]))["streams"] + + assert streams[0].data == expected_data + assert mock_stream.call_count == 2 + + @pytest.mark.flaky(reruns=3, reruns_delay=5) @pytest.mark.integration @pytest.mark.asyncio diff --git a/test/dataclasses/test_file_content.py b/test/dataclasses/test_file_content.py index a2bb13dbad5..dcf85492d6b 100644 --- a/test/dataclasses/test_file_content.py +++ b/test/dataclasses/test_file_content.py @@ -5,8 +5,10 @@ import base64 import logging import warnings +from collections.abc import Iterator +from contextlib import contextmanager, nullcontext from pathlib import Path -from unittest.mock import Mock, patch +from unittest.mock import patch import httpx import pytest @@ -14,6 +16,28 @@ from haystack.dataclasses.file_content import FileContent +@contextmanager +def mock_fetch_response( + url: str, status_code: int = 200, content: bytes | None = None, headers: dict[str, str] | None = None +) -> Iterator[None]: + """ + Not a test: makes the LinkContentFetcher used by `FileContent.from_url` replay a single, real httpx + response without touching the network. + """ + response = httpx.Response( + status_code, + content=content, + headers=headers or {"Content-Type": "application/pdf"}, + request=httpx.Request("GET", url), + ) + with ( + patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream, + patch("haystack.components.fetchers.link_content._resolve_host", return_value=["93.184.216.34"]), + ): + mock_stream.side_effect = lambda method, url, headers=None, **kwargs: nullcontext(response) + yield + + def test_file_content_init(base64_pdf_string): file_content = FileContent( base64_data=base64_pdf_string, mime_type="application/pdf", filename="test.pdf", extra={"key": "value"} @@ -118,12 +142,9 @@ def test_file_content_from_file_path_default_filename(test_files_path): def test_file_content_from_url(test_files_path): - with patch("haystack.components.fetchers.link_content.httpx.Client.get") as mock_get: - with open(test_files_path / "pdf" / "sample_pdf_3.pdf", "rb") as f: - pdf_bytes = f.read() - mock_response = Mock(status_code=200, content=pdf_bytes, headers={"Content-Type": "application/pdf"}) - mock_get.return_value = mock_response - + with open(test_files_path / "pdf" / "sample_pdf_3.pdf", "rb") as f: + pdf_bytes = f.read() + with mock_fetch_response("https://example.com/sample.pdf", content=pdf_bytes): file_content = FileContent.from_url( url="https://example.com/sample.pdf", filename="custom.pdf", extra={"test": "test"} ) @@ -135,21 +156,16 @@ def test_file_content_from_url(test_files_path): def test_file_content_from_url_default_filename(test_files_path): - with patch("haystack.components.fetchers.link_content.httpx.Client.get") as mock_get: - with open(test_files_path / "pdf" / "sample_pdf_3.pdf", "rb") as f: - pdf_bytes = f.read() - mock_response = Mock(status_code=200, content=pdf_bytes, headers={"Content-Type": "application/pdf"}) - mock_get.return_value = mock_response - + with open(test_files_path / "pdf" / "sample_pdf_3.pdf", "rb") as f: + pdf_bytes = f.read() + with mock_fetch_response("https://example.com/documents/sample.pdf", content=pdf_bytes): file_content = FileContent.from_url(url="https://example.com/documents/sample.pdf") assert file_content.filename == "sample.pdf" def test_file_content_from_url_bad_request(): - with patch("haystack.components.fetchers.link_content.httpx.Client.get") as mock_get: - mock_get.side_effect = httpx.HTTPStatusError("403 Client Error", request=Mock(), response=Mock()) - + with mock_fetch_response("https://non_existent_website_dot.com/file.pdf", status_code=403): with pytest.raises(httpx.HTTPStatusError): FileContent.from_url(url="https://non_existent_website_dot.com/file.pdf", retry_attempts=0, timeout=1) diff --git a/test/dataclasses/test_image_content.py b/test/dataclasses/test_image_content.py index ce035794fa2..23de3938581 100644 --- a/test/dataclasses/test_image_content.py +++ b/test/dataclasses/test_image_content.py @@ -5,7 +5,9 @@ import base64 import logging import warnings -from unittest.mock import Mock, patch +from collections.abc import Iterator +from contextlib import contextmanager, nullcontext +from unittest.mock import patch import httpx import pytest @@ -14,6 +16,23 @@ from haystack.dataclasses.image_content import ImageContent +@contextmanager +def mock_fetch_response( + url: str, status_code: int = 200, content: bytes | None = None, headers: dict[str, str] | None = None +) -> Iterator[None]: + """ + Not a test: makes the LinkContentFetcher used by `ImageContent.from_url` replay a single, real httpx + response without touching the network. + """ + response = httpx.Response(status_code, content=content, headers=headers, request=httpx.Request("GET", url)) + with ( + patch("haystack.components.fetchers.link_content.httpx.Client.stream") as mock_stream, + patch("haystack.components.fetchers.link_content._resolve_host", return_value=["93.184.216.34"]), + ): + mock_stream.side_effect = lambda method, url, headers=None, **kwargs: nullcontext(response) + yield + + def test_image_content_init(base64_image_string): image_content = ImageContent( base64_image=base64_image_string, mime_type="image/png", detail="auto", meta={"key": "value"} @@ -171,12 +190,11 @@ def test_image_content_from_file_path_non_existing(test_files_path, caplog): def test_image_content_from_url(test_files_path): - with patch("haystack.components.fetchers.link_content.httpx.Client.get") as mock_get: - with open(test_files_path / "images" / "apple.jpg", "rb") as image_file: - image_bytes = image_file.read() - mock_response = Mock(status_code=200, content=image_bytes, headers={"Content-Type": "image/jpeg"}) - mock_get.return_value = mock_response - + with open(test_files_path / "images" / "apple.jpg", "rb") as image_file: + image_bytes = image_file.read() + with mock_fetch_response( + "https://example.com/apple.jpg", content=image_bytes, headers={"Content-Type": "image/jpeg"} + ): image_content = ImageContent.from_url( url="https://example.com/apple.jpg", size=(100, 100), detail="high", meta={"test": "test"} ) @@ -188,18 +206,13 @@ def test_image_content_from_url(test_files_path): def test_image_content_from_url_bad_request(): - with patch("haystack.components.fetchers.link_content.httpx.Client.get") as mock_get: - mock_get.side_effect = httpx.HTTPStatusError("403 Client Error", request=Mock(), response=Mock()) - + with mock_fetch_response("https://non_existent_website_dot.com/image.jpg", status_code=403): with pytest.raises(httpx.HTTPStatusError): ImageContent.from_url(url="https://non_existent_website_dot.com/image.jpg", retry_attempts=0, timeout=1) def test_image_content_from_url_wrong_mime_type_text(): - with patch("haystack.components.fetchers.link_content.httpx.Client.get") as mock_get: - mock_response = Mock(status_code=200, text="a text", headers={"Content-Type": "text/plain"}) - mock_get.return_value = mock_response - + with mock_fetch_response("https://example.com/text.txt", content=b"a text", headers={"Content-Type": "text/plain"}): with pytest.raises(ValueError): ImageContent.from_url( url="https://example.com/text.txt", size=(100, 100), detail="high", meta={"test": "test"} @@ -207,12 +220,11 @@ def test_image_content_from_url_wrong_mime_type_text(): def test_image_content_from_url_wrong_mime_type_pdf(test_files_path): - with patch("haystack.components.fetchers.link_content.httpx.Client.get") as mock_get: - with open(test_files_path / "pdf" / "sample_pdf_1.pdf", "rb") as pdf_file: - pdf_bytes = pdf_file.read() - mock_response = Mock(status_code=200, content=pdf_bytes, headers={"Content-Type": "application/pdf"}) - mock_get.return_value = mock_response - + with open(test_files_path / "pdf" / "sample_pdf_1.pdf", "rb") as pdf_file: + pdf_bytes = pdf_file.read() + with mock_fetch_response( + "https://example.com/sample_pdf_1.pdf", content=pdf_bytes, headers={"Content-Type": "application/pdf"} + ): with pytest.raises(ValueError): ImageContent.from_url( url="https://example.com/sample_pdf_1.pdf", size=(100, 100), detail="high", meta={"test": "test"}