diff --git a/assert_ai/core/security.py b/assert_ai/core/security.py index fb38c9a0..0fccaf14 100644 --- a/assert_ai/core/security.py +++ b/assert_ai/core/security.py @@ -132,6 +132,13 @@ def validate_sys_path_addition(path: Path, *, config_path: Path | None = None) - ipaddress.ip_network("fe80::/10"), # IPv6 link-local ] +_NAT64_WELL_KNOWN_PREFIX = ipaddress.ip_network("64:ff9b::/96") +_IPV4_COMPATIBLE_PREFIX = ipaddress.ip_network("::/96") +_IPV4_COMPATIBLE_RESERVED = { + ipaddress.ip_address("::"), + ipaddress.ip_address("::1"), +} + _BLOCKED_HOSTNAMES = { "metadata.google.internal", "metadata.google.com", @@ -189,20 +196,15 @@ def validate_endpoint_url(url: str, *, allow_private: bool = False) -> None: f"URL hostname '{hostname}' is blocked (potential metadata endpoint)" ) - # Try to parse as IP address + # Validate IP literals directly; hostnames are checked after DNS resolution. try: - ip = ipaddress.ip_address(hostname) - for network in _BLOCKED_IP_RANGES: - if ip in network: - raise ValueError( - f"URL resolves to blocked IP range ({network}): {hostname}" - ) - except ValueError as e: - if "blocked" in str(e).lower(): - raise - # Not an IP literal — resolve hostname and check resulting IPs + ipaddress.ip_address(hostname) + except ValueError: + # Not an IP literal — resolve hostname and check resulting IPs. if hostname.lower() not in _LOCAL_DEV_HOSTNAMES: _validate_resolved_ips(hostname) + else: + validate_resolved_endpoint_ip(hostname, hostname) def _validate_resolved_ips(hostname: str) -> None: @@ -219,23 +221,81 @@ def _validate_resolved_ips(hostname: str) -> None: log.debug("DNS resolution failed for '%s'; skipping IP validation", hostname) return - for family, _type, _proto, _canonname, sockaddr in addrinfo: - ip_str = sockaddr[0] - try: - ip = ipaddress.ip_address(ip_str) - except ValueError: - continue - for network in _BLOCKED_IP_RANGES: - if ip in network: - log.warning( - "SSRF protection: hostname '%s' resolves to blocked IP %s (range %s)", - hostname, - ip_str, - network, - ) - raise ValueError( - f"URL hostname '{hostname}' resolves to blocked IP range ({network}): {ip_str}" - ) + for _family, _type, _proto, _canonname, sockaddr in addrinfo: + validate_resolved_endpoint_ip(hostname, sockaddr[0]) + + +def _canonicalize_endpoint_ip( + ip: ipaddress.IPv4Address | ipaddress.IPv6Address, +) -> ipaddress.IPv4Address | ipaddress.IPv6Address: + """Return the address an IPv4-in-IPv6 endpoint ultimately reaches.""" + if not isinstance(ip, ipaddress.IPv6Address): + return ip + + if ip.ipv4_mapped is not None: + return ip.ipv4_mapped + + if ip in _NAT64_WELL_KNOWN_PREFIX or ( + ip in _IPV4_COMPATIBLE_PREFIX and ip not in _IPV4_COMPATIBLE_RESERVED + ): + return ipaddress.IPv4Address(int(ip) & 0xFFFFFFFF) + + return ip + + +def validate_resolved_endpoint_ip(hostname: str, ip_str: str) -> None: + """Validate one DNS answer immediately before an endpoint connection. + + ``validate_endpoint_url`` performs an eager DNS check for fast feedback, but + HTTP clients resolve again when they connect. Call this helper on the + resolver results that will actually be used for the socket so a hostname + cannot pass validation with a public address and later rebind to a private + one. + """ + if os.environ.get("ASSERT_ALLOW_PRIVATE_ENDPOINTS", "").lower() in ( + "1", + "true", + "yes", + ): + return + if hostname.lower() in _LOCAL_DEV_HOSTNAMES: + return + + try: + ip = ipaddress.ip_address(ip_str) + except ValueError as exc: + raise ValueError( + f"Resolver returned a non-IP address for endpoint hostname '{hostname}': {ip_str}" + ) from exc + + # Mapped, NAT64, and deprecated IPv4-compatible IPv6 addresses inherit the + # security properties of the IPv4 endpoint they ultimately reach. + checked_ip = _canonicalize_endpoint_ip(ip) + + for network in _BLOCKED_IP_RANGES: + if checked_ip.version == network.version and checked_ip in network: + log.warning( + "SSRF protection: hostname '%s' resolves to blocked IP %s (range %s)", + hostname, + ip_str, + network, + ) + raise ValueError( + f"URL hostname '{hostname}' resolves to blocked IP range ({network}): {ip_str}" + ) + + # Block unspecified, reserved, documentation, shared, and other special-use + # addresses as well as multicast. SSRF targets must resolve to a publicly + # routable unicast address unless the explicit development override is set. + if not checked_ip.is_global or checked_ip.is_multicast: + log.warning( + "SSRF protection: hostname '%s' resolves to non-public IP %s", + hostname, + ip_str, + ) + raise ValueError( + f"URL hostname '{hostname}' resolves to a non-public IP address: {ip_str}" + ) # ── Credential sanitization ──────────────────────────────────── @@ -278,4 +338,3 @@ def sanitize_payload(payload: Any, *, depth: int = 0, max_depth: int = 10) -> An return _REDACTED return payload return payload - diff --git a/assert_ai/core/session.py b/assert_ai/core/session.py index 5fb3d8fc..5c21651f 100644 --- a/assert_ai/core/session.py +++ b/assert_ai/core/session.py @@ -6,6 +6,7 @@ from __future__ import annotations import asyncio +import contextlib import inspect import json import logging @@ -660,6 +661,7 @@ def __init__( self._system_prompt = system_prompt self._timeout_s = message_timeout_s self._session = None # aiohttp.ClientSession + self._resolver = None @property def runtime_mode(self) -> str: @@ -675,12 +677,29 @@ async def open(self) -> None: ) self._aiohttp = aiohttp timeout = aiohttp.ClientTimeout(total=self._timeout_s or 60) - self._session = aiohttp.ClientSession(timeout=timeout) + resolver = _ValidatingResolver(aiohttp.DefaultResolver()) + connector = None + try: + connector = aiohttp.TCPConnector(resolver=resolver) + self._session = aiohttp.ClientSession(timeout=timeout, connector=connector) + except Exception: + if connector is not None: + with contextlib.suppress(Exception): + await connector.close() + with contextlib.suppress(Exception): + await resolver.close() + raise + self._resolver = resolver async def close(self) -> None: - if self._session: - await self._session.close() - self._session = None + try: + if self._session: + await self._session.close() + self._session = None + finally: + if self._resolver: + await self._resolver.close() + self._resolver = None async def run_turn(self, messages: list[Message]) -> TurnResult: aiohttp = self._aiohttp @@ -704,7 +723,12 @@ async def run_turn(self, messages: list[Message]) -> TurnResult: self._endpoint, json=payload, headers=self._headers, + allow_redirects=False, ) as resp: + if 300 <= resp.status < 400: + raise RuntimeError( + f"HTTP endpoint {self._endpoint} returned a redirect, which is not allowed" + ) resp.raise_for_status() data = await resp.json() response_text = data.get("response", "") @@ -733,6 +757,29 @@ async def run_turn(self, messages: list[Message]) -> TurnResult: ) +class _ValidatingResolver: + """Validate the exact DNS answers aiohttp will use for a connection.""" + + def __init__(self, resolver: Any) -> None: + self._resolver = resolver + + async def resolve(self, host: str, port: int = 0, family: int = 0) -> list[Any]: + from assert_ai.core.security import validate_resolved_endpoint_ip + + results = await self._resolver.resolve(host, port, family) + try: + for result in results: + validate_resolved_endpoint_ip(host, str(result["host"])) + except ValueError as exc: + # aiohttp converts resolver OSErrors into ClientConnectorError, + # which the session already exposes as a scoped RuntimeError. + raise OSError(str(exc)) from exc + return results + + async def close(self) -> None: + await self._resolver.close() + + class ExternalSession: def __init__( self, diff --git a/assert_ai/stages/inference.py b/assert_ai/stages/inference.py index be40211a..63dc71e6 100644 --- a/assert_ai/stages/inference.py +++ b/assert_ai/stages/inference.py @@ -809,9 +809,8 @@ async def _run_tester_target_loop( assert action_message is not None target_messages.append(Message(role="user", content=action_message)) - tester_call_id = transcript.append_llm_call( - **build_llm_call_trace(tester_response, source="tester") - ) + tester_call_trace = build_llm_call_trace(tester_response, source="tester") + tester_call_id = transcript.append_llm_call(**tester_call_trace) message_id = f"event:{len(transcript.events)}" transcript.add_event(TranscriptEvent( view=["target", "combined"], @@ -819,7 +818,7 @@ async def _run_tester_target_loop( edit=AddMessageEdit(message=TranscriptMessage(role="user", content=action_message)), raw={ "call": "tester", - "request": to_jsonable(tester_response.request_payload or {}), + "request": tester_call_trace["request"], "response": serialize_response(tester_response), }, )) @@ -1140,7 +1139,7 @@ async def _worker(test_case: tuple[int, dict[str, Any]]) -> dict[str, Any]: # LLMInputError) so we emit a quieter debug log instead of # the warning the generic input-error handler below would # emit for every adversarial case. - test_case_id = seed_row.get("seed_id", "?") + test_case_id = test_case_row.get("test_case_id", "?") log.debug( "Inference worker hit provider content filter for test case %s: %s", test_case_id, exc, diff --git a/tests/test_http_endpoint_security.py b/tests/test_http_endpoint_security.py new file mode 100644 index 00000000..d5d4c0d1 --- /dev/null +++ b/tests/test_http_endpoint_security.py @@ -0,0 +1,168 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Regression tests for HTTP endpoint SSRF protections.""" + +import socket +import unittest +from unittest.mock import patch + +try: + import aiohttp + from aiohttp import web +except ImportError: # pragma: no cover - optional dependency + aiohttp = None + web = None + +from assert_ai.core.model_client import Message +from assert_ai.core.session import HTTPEndpointSession + + +@unittest.skipIf(web is None, "aiohttp not installed") +class HTTPEndpointSecurityTest(unittest.IsolatedAsyncioTestCase): + async def test_open_closes_resolver_if_connector_setup_fails(self) -> None: + class FakeResolver: + def __init__(self) -> None: + self.closed = False + + async def close(self) -> None: + self.closed = True + + resolver = FakeResolver() + session = HTTPEndpointSession(endpoint="http://localhost:8080/target") + + with ( + patch.object(aiohttp, "DefaultResolver", return_value=resolver), + patch.object( + aiohttp, + "TCPConnector", + side_effect=RuntimeError("setup failed"), + ), + ): + with self.assertRaisesRegex(RuntimeError, "setup failed"): + await session.open() + + self.assertTrue(resolver.closed) + + async def test_open_closes_connector_and_resolver_if_session_setup_fails(self) -> None: + class FakeResolver: + def __init__(self) -> None: + self.closed = False + + async def close(self) -> None: + self.closed = True + + class FakeConnector: + def __init__(self) -> None: + self.closed = False + + async def close(self) -> None: + self.closed = True + + resolver = FakeResolver() + connector = FakeConnector() + session = HTTPEndpointSession(endpoint="http://localhost:8080/target") + + with ( + patch.object(aiohttp, "DefaultResolver", return_value=resolver), + patch.object(aiohttp, "TCPConnector", return_value=connector), + patch.object( + aiohttp, + "ClientSession", + side_effect=RuntimeError("setup failed"), + ), + ): + with self.assertRaisesRegex(RuntimeError, "setup failed"): + await session.open() + + self.assertTrue(connector.closed) + self.assertTrue(resolver.closed) + + async def test_redirect_is_rejected_before_destination_request(self) -> None: + destination_reached = False + + async def redirect(_request): + raise web.HTTPTemporaryRedirect(location="/private") + + async def private(_request): + nonlocal destination_reached + destination_reached = True + return web.json_response({"response": "should not be reached"}) + + app = web.Application() + app.router.add_post("/start", redirect) + app.router.add_post("/private", private) + runner = web.AppRunner(app) + await runner.setup() + site = web.TCPSite(runner, "127.0.0.1", 0) + await site.start() + port = site._server.sockets[0].getsockname()[1] + + session = HTTPEndpointSession(endpoint=f"http://localhost:{port}/start") + await session.open() + try: + with self.assertRaisesRegex(RuntimeError, "redirect"): + await session.run_turn([Message(role="user", content="probe")]) + self.assertFalse(destination_reached) + finally: + await session.close() + await runner.cleanup() + + async def test_connection_time_private_dns_answer_is_rejected(self) -> None: + destination_reached = False + + async def private(_request): + nonlocal destination_reached + destination_reached = True + return web.json_response({"response": "should not be reached"}) + + app = web.Application() + app.router.add_post("/target", private) + runner = web.AppRunner(app) + await runner.setup() + site = web.TCPSite(runner, "127.0.0.1", 0) + await site.start() + port = site._server.sockets[0].getsockname()[1] + resolution_count = 0 + + def fake_getaddrinfo(host, requested_port, *args, **kwargs): + nonlocal resolution_count + resolution_count += 1 + if resolution_count == 1: + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + socket.IPPROTO_TCP, + "", + ("93.184.216.34", 0), + ) + ] + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + socket.IPPROTO_TCP, + "", + ("127.0.0.1", requested_port), + ) + ] + + try: + with patch("socket.getaddrinfo", side_effect=fake_getaddrinfo): + session = HTTPEndpointSession( + endpoint=f"http://rebind.test:{port}/target" + ) + await session.open() + try: + with self.assertRaisesRegex(RuntimeError, "Connection error"): + await session.run_turn([Message(role="user", content="probe")]) + self.assertFalse(destination_reached) + finally: + await session.close() + finally: + await runner.cleanup() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_inference_stage.py b/tests/test_inference_stage.py index bc1d1d2e..c197bc15 100644 --- a/tests/test_inference_stage.py +++ b/tests/test_inference_stage.py @@ -11,7 +11,7 @@ from assert_ai.core.config_model import TesterConfig, EvaluationConfig, JudgeConfig, InferenceConfig, TargetConfig, ToolsConfig from assert_ai.core.io import load_test_cases -from assert_ai.core.model_client import LLMInputError, LLMProviderError, Message, ModelResponse +from assert_ai.core.model_client import LLMContentFilterError, LLMInputError, LLMProviderError, Message, ModelResponse from assert_ai.core.session import TurnResult from assert_ai.stages.inference import _prepare_test_cases, _inference_config_fingerprint, _run_prompt_test_case, run_inference from assert_ai.viewer_read_model import ViewerReadModelBuildError @@ -1184,7 +1184,7 @@ async def fake_run_prompt_test_case(**kwargs): test_case_id = str(kwargs["test_case"]["test_case_id"]) # Fail exactly one of 20 (5% — under 10% threshold). if test_case_id == "test_case_000010": - raise RuntimeError("content_filter_blocked") + raise LLMContentFilterError("content_filter_blocked") class FakeTranscript: def to_dict(self_inner) -> dict[str, str]: @@ -1217,6 +1217,7 @@ def to_dict(self_inner) -> dict[str, str]: ) self.assertEqual(result["count"], 20) + self.assertEqual(result["errored_count"], 1) inference_rows = [ json.loads(line) for line in (out_dir / "inference_set.jsonl") diff --git a/tests/test_security.py b/tests/test_security.py index 0388ae3d..12cfd77f 100644 --- a/tests/test_security.py +++ b/tests/test_security.py @@ -145,6 +145,42 @@ def test_azure_wireserver_ip_blocked(self) -> None: with self.assertRaises(ValueError, msg="blocked"): validate_endpoint_url("http://168.63.129.16/metadata") + def test_unspecified_ipv4_blocked(self) -> None: + with self.assertRaises(ValueError, msg="non-public"): + validate_endpoint_url("http://0.0.0.0/api") + + def test_unspecified_ipv6_blocked(self) -> None: + with self.assertRaises(ValueError, msg="non-public"): + validate_endpoint_url("http://[::]/api") + + def test_ipv6_loopback_blocked(self) -> None: + with self.assertRaises(ValueError, msg="blocked"): + validate_endpoint_url("http://[::1]/api") + + def test_ipv4_mapped_private_ip_blocked(self) -> None: + with self.assertRaises(ValueError, msg="blocked"): + validate_endpoint_url("http://[::ffff:127.0.0.1]/api") + + def test_ipv4_mapped_azure_wireserver_blocked(self) -> None: + with self.assertRaises(ValueError, msg="blocked"): + validate_endpoint_url("http://[::ffff:168.63.129.16]/metadata") + + def test_nat64_loopback_blocked(self) -> None: + with self.assertRaises(ValueError, msg="blocked"): + validate_endpoint_url("http://[64:ff9b::7f00:1]/api") + + def test_nat64_metadata_ip_blocked(self) -> None: + with self.assertRaises(ValueError, msg="blocked"): + validate_endpoint_url("http://[64:ff9b::a9fe:a9fe]/metadata") + + def test_nat64_azure_wireserver_blocked(self) -> None: + with self.assertRaises(ValueError, msg="blocked"): + validate_endpoint_url("http://[64:ff9b::a83f:8110]/metadata") + + def test_ipv4_compatible_loopback_blocked(self) -> None: + with self.assertRaises(ValueError, msg="blocked"): + validate_endpoint_url("http://[::127.0.0.1]/api") + # ── Blocked hostnames ── def test_gcp_metadata_hostname_blocked(self) -> None: @@ -183,6 +219,15 @@ def test_hostname_resolving_to_private_ip_blocked(self) -> None: with self.assertRaises(ValueError, msg="blocked IP range"): validate_endpoint_url("http://evil-rebind.attacker.com/api") + def test_hostname_with_mixed_public_and_private_answers_blocked(self) -> None: + fake_addrinfo = [ + (socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 0)), + (socket.AF_INET, socket.SOCK_STREAM, 6, "", ("127.0.0.1", 0)), + ] + with patch("assert_ai.core.security.socket.getaddrinfo", return_value=fake_addrinfo): + with self.assertRaises(ValueError, msg="blocked IP range"): + validate_endpoint_url("http://mixed-answers.attacker.com/api") + def test_dns_failure_allows_passthrough(self) -> None: with patch( "assert_ai.core.security.socket.getaddrinfo", @@ -196,6 +241,15 @@ def test_dns_failure_allows_passthrough(self) -> None: def test_public_ip_allowed(self) -> None: validate_endpoint_url("http://93.184.216.34/api") + def test_ipv4_mapped_public_ip_allowed(self) -> None: + validate_endpoint_url("http://[::ffff:93.184.216.34]/api") + + def test_nat64_public_ip_allowed(self) -> None: + validate_endpoint_url("http://[64:ff9b::5db8:d822]/api") + + def test_ipv4_compatible_public_ip_allowed(self) -> None: + validate_endpoint_url("http://[::93.184.216.34]/api") + def test_public_hostname_with_public_ip_allowed(self) -> None: fake_addrinfo = [ (socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 0)), diff --git a/tests/test_tester_target_loop.py b/tests/test_tester_target_loop.py index bdd5de30..080b3a45 100644 --- a/tests/test_tester_target_loop.py +++ b/tests/test_tester_target_loop.py @@ -8,6 +8,7 @@ recording, and stop-reason behaviour. """ +import json import unittest from typing import Any from unittest.mock import patch @@ -101,6 +102,7 @@ async def fake_generate(model, messages, options): response.request_payload = { "model": model, "messages": [message.to_openai_dict() for message in messages], + **(response.request_payload or {}), } return response @@ -348,6 +350,27 @@ async def test_transcript_tester_event_omits_raw(self) -> None: # New code records raw call data on tester events self.assertIsNotNone(tester_events[0].raw) + async def test_transcript_redacts_tester_request_credentials(self) -> None: + response = _tester_response("Hello") + response.request_payload = { + "api_key": "sentinel-azure-foundry-token", + "model": "azure_ai/tester", + } + + result = await self._run_loop( + tester_responses=[response], + target_replies=["World"], + max_turns=1, + ) + + transcript: Transcript = result["transcript"] + tester_event = next(event for event in transcript.events if event.actor == "tester") + self.assertEqual(tester_event.raw["request"]["api_key"], "[REDACTED]") + self.assertNotIn( + "sentinel-azure-foundry-token", + json.dumps(transcript.to_dict()), + ) + async def test_transcript_target_event_omits_raw(self) -> None: result = await self._run_loop( tester_responses=[_tester_response("Hello")],