From bf5026666f8f25993c834c0ad313d197b03ddbd3 Mon Sep 17 00:00:00 2001 From: Codex Date: Sat, 3 Oct 2026 02:42:36 +0900 Subject: [PATCH] test: use real HTTPX objects in SSRF proxy tests --- .../unit_tests/core/helper/test_ssrf_proxy.py | 193 +++++++++--------- 1 file changed, 92 insertions(+), 101 deletions(-) diff --git a/api/tests/unit_tests/core/helper/test_ssrf_proxy.py b/api/tests/unit_tests/core/helper/test_ssrf_proxy.py index 617fffd618c..d522367f28e 100644 --- a/api/tests/unit_tests/core/helper/test_ssrf_proxy.py +++ b/api/tests/unit_tests/core/helper/test_ssrf_proxy.py @@ -1,8 +1,8 @@ import gzip -from collections.abc import Callable -from types import SimpleNamespace -from typing import override -from unittest.mock import ANY, MagicMock, call, patch +from collections.abc import Callable, Iterator +from contextlib import ExitStack +from typing import TypedDict, Unpack +from unittest.mock import ANY, call, patch import httpx import pytest @@ -25,18 +25,44 @@ from core.helper.ssrf_proxy import ( from core.tools.errors import ToolSSRFError +class _SendOptions(TypedDict, total=False): + stream: bool + follow_redirects: bool + auth: httpx.Auth | tuple[str, str] | Callable[[httpx.Request], httpx.Request] | None + + +@pytest.fixture +def client_factory( + monkeypatch: pytest.MonkeyPatch, +) -> Iterator[Callable[[httpx.Response], tuple[httpx.Client, list[_SendOptions]]]]: + with ExitStack() as stack: + + def create_client(response: httpx.Response) -> tuple[httpx.Client, list[_SendOptions]]: + client = stack.enter_context(httpx.Client(transport=httpx.MockTransport(lambda _request: response))) + send = client.send + calls: list[_SendOptions] = [] + + def record_send(request: httpx.Request, **kwargs: Unpack[_SendOptions]) -> httpx.Response: + calls.append(kwargs) + return send(request, **kwargs) + + monkeypatch.setattr(client, "send", record_send) + return client, calls + + yield create_client + + @patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True) -def test_successful_request(mock_get_client): - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_client.send.return_value = mock_response +def test_successful_request(mock_get_client, client_factory): + mock_response = httpx.Response(200) + mock_client, send_calls = client_factory(mock_response) mock_get_client.return_value = mock_client response = make_request("GET", "http://example.com") assert response.status_code == 200 - mock_client.build_request.assert_called_once() - mock_client.send.assert_called_once() + assert response.request.method == "GET" + assert str(response.request.url) == "http://example.com" + assert len(send_calls) == 1 def test_buffer_response_rejects_encoded_response_before_decoding() -> None: @@ -85,15 +111,10 @@ def test_buffer_response_rejects_identity_response_exceeding_byte_limit() -> Non def test_request_can_return_an_open_stream_the_caller_closes() -> None: - class EventStream(httpx.SyncByteStream): - @override - def __iter__(self): - yield b"event: delta\ndata: first\n\n" - transport = httpx.MockTransport( lambda request: httpx.Response( 200, - stream=EventStream(), + stream=httpx.ByteStream(b"event: delta\ndata: first\n\n"), request=request, ) ) @@ -118,36 +139,30 @@ def test_request_can_return_an_open_stream_the_caller_closes() -> None: @patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True) -def test_retry_exceed_max_retries(mock_get_client, monkeypatch: pytest.MonkeyPatch): - sleep = MagicMock() - monkeypatch.setattr("core.helper.ssrf_proxy.time", SimpleNamespace(sleep=sleep)) - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 500 - mock_client.send.return_value = mock_response +def test_retry_exceed_max_retries(mock_get_client, client_factory, monkeypatch: pytest.MonkeyPatch): + sleep_calls: list[float] = [] + monkeypatch.setattr("core.helper.ssrf_proxy.time.sleep", sleep_calls.append) + mock_response = httpx.Response(500) + mock_client, send_calls = client_factory(mock_response) mock_get_client.return_value = mock_client with pytest.raises(Exception) as e: make_request("GET", "http://example.com", max_retries=SSRF_DEFAULT_MAX_RETRIES - 1) assert str(e.value) == f"Reached maximum retries ({SSRF_DEFAULT_MAX_RETRIES - 1}) for URL http://example.com" - assert mock_client.send.call_count == SSRF_DEFAULT_MAX_RETRIES - assert sleep.call_args_list == [ - call(BACKOFF_FACTOR * 2**attempt) for attempt in range(SSRF_DEFAULT_MAX_RETRIES - 1) - ] + assert len(send_calls) == SSRF_DEFAULT_MAX_RETRIES + assert sleep_calls == [BACKOFF_FACTOR * 2**attempt for attempt in range(SSRF_DEFAULT_MAX_RETRIES - 1)] @patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True) -def test_force_list_response_returns_when_retries_disabled(mock_get_client): - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 500 - mock_client.send.return_value = mock_response +def test_force_list_response_returns_when_retries_disabled(mock_get_client, client_factory): + mock_response = httpx.Response(500) + mock_client, send_calls = client_factory(mock_response) mock_get_client.return_value = mock_client response = make_request("GET", "http://example.com", max_retries=0) assert response is mock_response - mock_client.send.assert_called_once() + assert len(send_calls) == 1 def test_build_ssrf_client_passes_ssl_verify_to_proxy_mount_transports( @@ -158,17 +173,18 @@ def test_build_ssrf_client_passes_ssl_verify_to_proxy_mount_transports( SSRF_PROXY_HTTP_URL="http://proxy.example.com:8080", SSRF_PROXY_HTTPS_URL="http://proxy.example.com:8443", ) - mock_client = MagicMock() - http_transport = MagicMock() - https_transport = MagicMock() + http_transport = httpx.HTTPTransport(proxy="http://proxy.example.com:8080", verify=False) + https_transport = httpx.HTTPTransport(proxy="http://proxy.example.com:8443", verify=False) + client_class = httpx.Client with ( patch("core.helper.ssrf_proxy.httpx.HTTPTransport", side_effect=[http_transport, https_transport]) as transport, - patch("core.helper.ssrf_proxy.httpx.Client", return_value=mock_client) as client, + patch("core.helper.ssrf_proxy.httpx.Client", wraps=client_class) as client, ): ssrf_client = _build_ssrf_client(verify=False) - assert ssrf_client is mock_client + assert isinstance(ssrf_client, client_class) + ssrf_client.close() transport.assert_has_calls( [ call(proxy="http://proxy.example.com:8080", verify=False), @@ -219,12 +235,10 @@ class TestGetUserProvidedHostHeader: @patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True) -def test_host_header_preservation_with_user_header(mock_get_client): +def test_host_header_preservation_with_user_header(mock_get_client, client_factory): """Test that user-provided Host header is preserved in the request.""" - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_client.send.return_value = mock_response + mock_response = httpx.Response(200) + mock_client, send_calls = client_factory(mock_response) mock_get_client.return_value = mock_client custom_host = "custom.example.com:8080" @@ -232,26 +246,22 @@ def test_host_header_preservation_with_user_header(mock_get_client): assert response.status_code == 200 # Verify the request was built with the host header preserved (lowercase) - call_kwargs = mock_client.build_request.call_args.kwargs - assert call_kwargs["headers"]["host"] == custom_host + assert response.request.headers["host"] == custom_host @patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True) @pytest.mark.parametrize("host_key", ["host", "HOST", "Host"]) -def test_host_header_preservation_case_insensitive(mock_get_client, host_key): +def test_host_header_preservation_case_insensitive(mock_get_client, client_factory, host_key): """Test that Host header is preserved regardless of case.""" - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_client.send.return_value = mock_response + mock_response = httpx.Response(200) + mock_client, send_calls = client_factory(mock_response) mock_get_client.return_value = mock_client response = make_request("GET", "http://example.com", headers={host_key: "api.example.com"}) assert response.status_code == 200 # Host header should be normalized to lowercase "host" - call_kwargs = mock_client.build_request.call_args.kwargs - assert call_kwargs["headers"]["host"] == "api.example.com" + assert response.request.headers["host"] == "api.example.com" class TestFollowRedirectsParameter: @@ -261,65 +271,57 @@ class TestFollowRedirectsParameter: """ @patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True) - def test_follow_redirects_passed_to_request(self, mock_get_client): + def test_follow_redirects_passed_to_request(self, mock_get_client, client_factory): """Verify follow_redirects IS passed to client.send().""" - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_client.send.return_value = mock_response + mock_response = httpx.Response(200) + mock_client, send_calls = client_factory(mock_response) mock_get_client.return_value = mock_client make_request("GET", "http://example.com", follow_redirects=True) # Verify follow_redirects was passed to send - call_kwargs = mock_client.send.call_args.kwargs + call_kwargs = send_calls[-1] assert call_kwargs.get("follow_redirects") is True @patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True) - def test_allow_redirects_converted_to_follow_redirects(self, mock_get_client): + def test_allow_redirects_converted_to_follow_redirects(self, mock_get_client, client_factory): """Verify allow_redirects (requests-style) is converted to follow_redirects (httpx-style).""" - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_client.send.return_value = mock_response + mock_response = httpx.Response(200) + mock_client, send_calls = client_factory(mock_response) mock_get_client.return_value = mock_client # Use allow_redirects (requests-style parameter) make_request("GET", "http://example.com", allow_redirects=True) # Verify it was converted to follow_redirects - call_kwargs = mock_client.send.call_args.kwargs + call_kwargs = send_calls[-1] assert call_kwargs.get("follow_redirects") is True assert "allow_redirects" not in call_kwargs @patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True) - def test_follow_redirects_not_set_when_not_specified(self, mock_get_client): + def test_follow_redirects_not_set_when_not_specified(self, mock_get_client, client_factory): """Verify follow_redirects is not in kwargs when not specified (httpx default behavior).""" - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_client.send.return_value = mock_response + mock_response = httpx.Response(200) + mock_client, send_calls = client_factory(mock_response) mock_get_client.return_value = mock_client make_request("GET", "http://example.com") # follow_redirects should not be in kwargs, letting httpx use its default - call_kwargs = mock_client.send.call_args.kwargs + call_kwargs = send_calls[-1] assert "follow_redirects" not in call_kwargs @patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True) - def test_follow_redirects_takes_precedence_over_allow_redirects(self, mock_get_client): + def test_follow_redirects_takes_precedence_over_allow_redirects(self, mock_get_client, client_factory): """Verify follow_redirects takes precedence when both are specified.""" - mock_client = MagicMock() - mock_response = MagicMock() - mock_response.status_code = 200 - mock_client.send.return_value = mock_response + mock_response = httpx.Response(200) + mock_client, send_calls = client_factory(mock_response) mock_get_client.return_value = mock_client # Both specified - follow_redirects should take precedence make_request("GET", "http://example.com", allow_redirects=False, follow_redirects=True) - call_kwargs = mock_client.send.call_args.kwargs + call_kwargs = send_calls[-1] assert call_kwargs.get("follow_redirects") is True @@ -386,21 +388,17 @@ def test_graphon_ssrf_proxy_wraps_module_requests(method_name: str) -> None: # protection" (the pre-#38443 message gave no actionable guidance). -def _build_squid_blocked_response(status_code: int = 403) -> MagicMock: - """Construct a mock httpx.Response that looks like Squid's ACL deny.""" - response = MagicMock() - response.status_code = status_code - response.headers = {"server": "squid/4.10", "via": "1.1 squid (squid/4.10)"} - return response +def _build_squid_blocked_response(status_code: int = 403) -> httpx.Response: + """Construct an HTTP response with Squid's ACL denial headers.""" + return httpx.Response(status_code, headers={"server": "squid/4.10", "via": "1.1 squid (squid/4.10)"}) @patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True) -def test_squid_block_raises_actionable_tool_ssrf_error(mock_get_client) -> None: +def test_squid_block_raises_actionable_tool_ssrf_error(mock_get_client, client_factory) -> None: """A 403 from Squid must raise ToolSSRFError whose message tells the user exactly which env var to set. Pre-#38443 the message had no remediation hint, so users hit dead ends when their internal API was blocked.""" - mock_client = MagicMock() - mock_client.send.return_value = _build_squid_blocked_response(status_code=403) + mock_client, _ = client_factory(_build_squid_blocked_response(status_code=403)) mock_get_client.return_value = mock_client with pytest.raises(ToolSSRFError) as exc_info: @@ -418,15 +416,12 @@ def test_squid_block_raises_actionable_tool_ssrf_error(mock_get_client) -> None: @patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True) -def test_squid_401_via_header_also_triggers_actionable_error(mock_get_client) -> None: +def test_squid_401_via_header_also_triggers_actionable_error(mock_get_client, client_factory) -> None: """Squid can return 401 with only the Via header set (no Server header) on some configurations. The detection must work for both.""" - mock_client = MagicMock() - response = MagicMock() - response.status_code = 401 # Server header absent — only Via identifies Squid. - response.headers = {"server": "", "via": "1.1 squid (squid/4.10)"} - mock_client.send.return_value = response + response = httpx.Response(401, headers={"server": "", "via": "1.1 squid (squid/4.10)"}) + mock_client, _ = client_factory(response) mock_get_client.return_value = mock_client with pytest.raises(ToolSSRFError) as exc_info: @@ -436,17 +431,14 @@ def test_squid_401_via_header_also_triggers_actionable_error(mock_get_client) -> @patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True) -def test_non_squid_403_is_not_treated_as_ssrf_block(mock_get_client) -> None: +def test_non_squid_403_is_not_treated_as_ssrf_block(mock_get_client, client_factory) -> None: """A 403 from the *target server* (not Squid) must NOT be re-raised as a ToolSSRFError — that would mislead the user into editing SSRF config when the real problem is application-level authorization on the target. Pre-#38443 we didn't have this guard at all; the new wording only changes the Squid path, so verify we don't accidentally widen it.""" - mock_client = MagicMock() - response = MagicMock() - response.status_code = 403 - response.headers = {"server": "nginx/1.21", "via": "1.1 varnish"} - mock_client.send.return_value = response + response = httpx.Response(403, headers={"server": "nginx/1.21", "via": "1.1 varnish"}) + mock_client, _ = client_factory(response) mock_get_client.return_value = mock_client # Should return the response, not raise. @@ -455,12 +447,11 @@ def test_non_squid_403_is_not_treated_as_ssrf_block(mock_get_client) -> None: @patch("core.helper.ssrf_proxy._get_ssrf_client", autospec=True) -def test_squid_block_with_internal_10_x_url_mentions_allowlist(mock_get_client) -> None: +def test_squid_block_with_internal_10_x_url_mentions_allowlist(mock_get_client, client_factory) -> None: """Real-world repro from #38443: 10.x.x.x internal API blocked. The error message must still point at SSRF_PROXY_ALLOW_PRIVATE_IPS, not just say "private address" without telling the user what to do.""" - mock_client = MagicMock() - mock_client.send.return_value = _build_squid_blocked_response(status_code=403) + mock_client, _ = client_factory(_build_squid_blocked_response(status_code=403)) mock_get_client.return_value = mock_client with pytest.raises(ToolSSRFError) as exc_info: