diff --git a/litellm/__init__.py b/litellm/__init__.py index c868ae55b4fe..075fbdb309ff 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -456,6 +456,11 @@ False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6. ) network_mock: bool = False # When True, use mock transport — no real network calls +allow_requests_to_internal_ips: bool = ( + False # When True, disables SSRF protection for private/reserved IPs. + # Only set this for self-hosted or on-prem deployments where api_base + # intentionally points at an internal address (e.g. local Ollama, vLLM). +) ####### STOP SEQUENCE LIMIT ####### disable_stop_sequence_limit: bool = False # when True, stop sequence limit is disabled diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index 93b6c563dc1e..fbc09052af0a 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -1,8 +1,14 @@ +import asyncio +import ipaddress +import os +import socket from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple, Union, cast +from urllib.parse import urlparse import aiohttp import httpx # type: ignore from aiohttp import ClientSession, FormData +from aiohttp.abc import AbstractResolver import litellm import litellm.litellm_core_utils @@ -12,10 +18,11 @@ from litellm.llms.base_llm.image_variations.transformation import ( BaseImageVariationConfig, ) +from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, - _get_httpx_client, + get_ssl_configuration, ) from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport from litellm.types.llms.openai import FileTypes @@ -31,6 +38,281 @@ DEFAULT_TIMEOUT = 600 +# CGNAT (100.64.0.0/10) is misclassified as global on Python < 3.11 +_CGNAT = ipaddress.ip_network("100.64.0.0/10") + + +def _is_blocked_address(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool: + """Return True if addr is not globally routable (private, loopback, reserved, etc.). + + Uses Python's built-in is_global to cover all IANA special-use ranges — + including RFC 5737 documentation (192.0.2.0/24), RFC 2544 benchmarking + (198.18.0.0/15), and class E (240.0.0.0/4) — without maintaining a manual list. + CGNAT (100.64.0.0/10) is explicitly checked for Python < 3.11 compat. + """ + # Unwrap IPv4-mapped IPv6 (::ffff:10.0.0.1 → 10.0.0.1) + if isinstance(addr, ipaddress.IPv6Address) and addr.ipv4_mapped is not None: + addr = addr.ipv4_mapped + # Python < 3.11 incorrectly marks CGNAT as global + if isinstance(addr, ipaddress.IPv4Address) and addr in _CGNAT: + return True + return not addr.is_global or addr.is_multicast + + +def _assert_not_private_url(url: str) -> None: + """Raise ValueError if url resolves to any private/reserved IP (SSRF protection). + + Validates all DNS answers, not just the first, to prevent A-record rotation attacks. + + On the sync path this is defence-in-depth for externally-supplied httpx clients. + Internally-created clients use ``_SSRFGuardTransport``, which resolves, validates, + and pins the IP at TCP-connect time — eliminating the DNS-rebinding TOCTOU window + that a preflight-only check cannot close. + + Set ``litellm.allow_requests_to_internal_ips = True`` to disable this check + for self-hosted / on-prem deployments where api_base is an internal address. + """ + if litellm.allow_requests_to_internal_ips: + return + parsed = urlparse(url) + hostname = parsed.hostname + if not hostname: + return + try: + answers = socket.getaddrinfo(hostname, None) + except socket.gaierror: + return # DNS failure — request will fail naturally + for answer in answers: + raw_ip = answer[4][0] + try: + addr = ipaddress.ip_address(raw_ip) + except ValueError: + continue + if _is_blocked_address(addr): + raise ValueError( + f"api_base '{url}' resolves to a private/reserved IP address " + f"({raw_ip}) which is not allowed (SSRF protection)" + ) + + +def _assert_not_private_ip_literal(url: str) -> None: + """Raise ValueError if the url host is a private/reserved IP address literal. + + aiohttp's TCPConnector skips _SSRFGuardResolver when the host is already an + IP address (no DNS lookup needed), so the resolver cannot block such URLs. + This function fills that gap with a fast, non-blocking check (no DNS I/O). + Hostname-based URLs are handled by _SSRFGuardResolver at connect time. + + Set ``litellm.allow_requests_to_internal_ips = True`` to disable this check. + """ + if litellm.allow_requests_to_internal_ips: + return + parsed = urlparse(url) + hostname = parsed.hostname + if not hostname: + return + try: + addr = ipaddress.ip_address(hostname) + except ValueError: + return # Not an IP literal — resolver will validate at connect time + if _is_blocked_address(addr): + raise ValueError( + f"api_base '{url}' contains a private/reserved IP address " + f"({hostname}) which is not allowed (SSRF protection)" + ) + + +class _SSRFGuardResolver(AbstractResolver): + """Custom aiohttp resolver that validates IPs at TCP-connection time. + + By hooking into aiohttp's resolver — used for every connection including + redirect targets — this eliminates the DNS-rebinding TOCTOU window that + a separate preflight check cannot close. All DNS answers are validated, + not just the first, to defend against A-record rotation. + """ + + async def resolve( + self, host: str, port: int = 0, family: int = socket.AF_INET + ) -> list: + loop = asyncio.get_running_loop() + try: + infos = await loop.getaddrinfo( + host, port, family=family, type=socket.SOCK_STREAM + ) + except socket.gaierror: + raise # Propagate so aiohttp wraps it in a ClientConnectorError + if not litellm.allow_requests_to_internal_ips: + for info in infos: + raw_ip = info[4][0] + try: + addr = ipaddress.ip_address(raw_ip) + except ValueError: + continue + if _is_blocked_address(addr): + raise ValueError( + f"Host '{host}' resolves to a private/reserved IP address " + f"({raw_ip}) which is not allowed (SSRF protection)" + ) + return [ + { + "hostname": host, + "host": info[4][0], + "port": info[4][1] if len(info[4]) > 1 else port, + "family": info[0], + "proto": info[2], + "flags": 0, + } + for info in infos + ] + + async def close(self) -> None: + pass + + +class _SSRFGuardTransport(httpx.HTTPTransport): + """Sync httpx transport that eliminates the DNS-rebinding TOCTOU gap. + + Mirrors ``_SSRFGuardResolver`` on the async path: resolves the hostname + once, validates every returned IP, then pins the TCP connection to the + first safe address by rewriting the request URL to an IP literal. + httpcore skips DNS for IP-literal hosts, so the second resolution that + creates the TOCTOU window never occurs. + + The original hostname is preserved in the ``Host`` header and the + ``sni_hostname`` extension so that TLS certificate validation and + virtual-host routing are unaffected. + + Set ``litellm.allow_requests_to_internal_ips = True`` to disable all + SSRF checking. + """ + + def handle_request(self, request: httpx.Request) -> httpx.Response: # type: ignore[override] + if not litellm.allow_requests_to_internal_ips: + host = request.url.host + try: + addr: Optional[ipaddress.IPv4Address | ipaddress.IPv6Address] = ( + ipaddress.ip_address(host) + ) + except ValueError: + addr = None + + if addr is not None: + # IP-literal host: validate inline — no DNS needed, no TOCTOU window. + if _is_blocked_address(addr): + raise ValueError( + f"Host '{host}' is a private/reserved IP address " + f"which is not allowed (SSRF protection)" + ) + else: + # Hostname path: resolve, validate all answers, then pin to prevent + # a second DNS lookup (closes the DNS-rebinding TOCTOU window). + port = request.url.port or ( + 443 if request.url.scheme == "https" else 80 + ) + try: + infos = socket.getaddrinfo(host, port, type=socket.SOCK_STREAM) + except socket.gaierror: + # DNS failed — let httpcore propagate the error naturally. + return super().handle_request(request) + + for info in infos: + raw_ip = info[4][0] + try: + resolved = ipaddress.ip_address(raw_ip) + except ValueError: + continue + if _is_blocked_address(resolved): + raise ValueError( + f"Host '{host}' resolves to a private/reserved IP " + f"address ({raw_ip}) which is not allowed (SSRF protection)" + ) + + if infos: + # Rewrite the URL to the first validated IP so that httpcore + # connects directly without a second DNS resolution — this closes + # the TOCTOU window entirely. + pinned_ip = infos[0][4][0] + pinned_url = request.url.copy_with(host=pinned_ip) + + # Preserve the original hostname in the Host header (virtual + # hosting) and sni_hostname extension (TLS cert validation). + headers = [ + (k, v) for k, v in request.headers.raw if k.lower() != b"host" + ] + headers.append((b"host", host.encode("utf-8"))) + extensions = {**request.extensions, "sni_hostname": host} + request = httpx.Request( + method=request.method, + url=pinned_url, + headers=headers, + content=request.stream, + extensions=extensions, + ) + + return super().handle_request(request) + + +async def _on_ssrf_request_start( + session: Any, + trace_config_ctx: Any, + params: Any, +) -> None: + """Block IP-literal redirect targets before aiohttp makes the connection. + + _SSRFGuardResolver handles all hostname-based requests (including redirect + hops) at DNS-resolution time. This hook closes the gap where a server + returns a Location header containing a bare IP literal — aiohttp skips DNS + for those, so the resolver never fires. Hostname-based URLs pass through + untouched here. + """ + _assert_not_private_ip_literal(str(params.url)) + + +def _make_ssrf_trace_config() -> aiohttp.TraceConfig: + """Return a TraceConfig that blocks IP-literal redirect targets.""" + cfg = aiohttp.TraceConfig() + cfg.on_request_start.append(_on_ssrf_request_start) + return cfg + + +_SSRF_SAFE_CLIENT_CACHE_KEY = "ssrf_safe_httpx_client" + + +def _get_ssrf_safe_sync_client() -> HTTPHandler: + """Return a cached HTTPHandler whose transport validates and pins IPs at connect time. + + Mirrors the caching behaviour of ``_get_httpx_client`` (1-hour TTL via + ``in_memory_llm_clients_cache``) so that persistent TCP/TLS connections are + reused across requests — fixing the per-call pool regression that a naive + ``HTTPHandler(transport=_SSRFGuardTransport())`` would introduce. + + SSL settings (``litellm.ssl_verify``, ``SSL_CERTIFICATE`` env-var) are + forwarded to ``_SSRFGuardTransport`` so that callers using self-signed certs + or mTLS are not silently broken. + """ + cache = getattr(litellm, "in_memory_llm_clients_cache", None) + if cache is None: + from litellm.caching.llm_caching_handler import LLMClientCache + + cache = LLMClientCache() + setattr(litellm, "in_memory_llm_clients_cache", cache) + + _cached_client = cache.get_cache(_SSRF_SAFE_CLIENT_CACHE_KEY) + if _cached_client: + return _cached_client + + ssl_config = get_ssl_configuration(None) + cert = os.getenv("SSL_CERTIFICATE", litellm.ssl_certificate) + _new_client = HTTPHandler( + transport=_SSRFGuardTransport(verify=ssl_config, cert=cert) + ) + cache.set_cache( + key=_SSRF_SAFE_CLIENT_CACHE_KEY, + value=_new_client, + ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS, + ) + return _new_client + class BaseLLMAIOHTTPHandler: def __init__( @@ -95,8 +377,15 @@ def _create_client_session_with_transport(self) -> ClientSession: session = aiohttp.ClientSession(connector=connector) return session else: - # Default session creation - session = aiohttp.ClientSession() + # Default session creation — attach SSRF guard resolver so every + # TCP connection (including redirect targets) is validated at the + # network layer, eliminating the DNS-rebinding TOCTOU window. + # The trace config adds a second layer: it blocks IP-literal + # redirect targets that bypass the resolver (no DNS lookup needed). + session = aiohttp.ClientSession( + connector=aiohttp.TCPConnector(resolver=_SSRFGuardResolver()), + trace_configs=[_make_ssrf_trace_config()], + ) return session def _get_async_client_session( @@ -147,26 +436,17 @@ def __del__(self): and self._owns_session ): try: - import asyncio - try: - loop = asyncio.get_event_loop() - if loop.is_running(): - # Event loop is running - schedule cleanup task - asyncio.create_task(self.close()) - else: - # Event loop exists but not running - run cleanup - loop.run_until_complete(self.close()) + loop = asyncio.get_running_loop() + loop.create_task(self.close()) except RuntimeError: - # No event loop available - create one for cleanup + # No running loop — run cleanup in a temporary one. loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) try: loop.run_until_complete(self.close()) finally: loop.close() except Exception: - # Silently ignore errors during __del__ to avoid issues pass async def _make_common_async_call( @@ -191,6 +471,13 @@ async def _make_common_async_call( dynamic_client_session=async_client_session ) + # IP-literal URLs bypass _SSRFGuardResolver because aiohttp's TCPConnector + # skips DNS resolution for hosts that are already IP addresses. Check them + # here with a fast, non-blocking parse (no socket I/O). Hostname-based URLs + # are handled by _SSRFGuardResolver at TCP-connect time, which also covers + # redirect targets, eliminating the DNS-rebinding TOCTOU window. + _assert_not_private_ip_literal(api_base) + for i in range(max(max_retry_on_unprocessable_entity_error, 1)): try: response = await async_client_session.post( @@ -235,6 +522,8 @@ def _make_common_sync_call( provider_config.max_retry_on_unprocessable_entity_error ) + _assert_not_private_url(api_base) + response: Optional[httpx.Response] = None for i in range(max(max_retry_on_unprocessable_entity_error, 1)): @@ -434,7 +723,7 @@ def completion( ) if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client() + sync_httpx_client = _get_ssrf_safe_sync_client() else: sync_httpx_client = client @@ -475,7 +764,7 @@ def make_sync_call( client: Optional[HTTPHandler] = None, ) -> Tuple[Any, dict]: if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client() + sync_httpx_client = _get_ssrf_safe_sync_client() else: sync_httpx_client = client stream = True @@ -660,7 +949,7 @@ def image_variations( ) # type: ignore if client is None or not isinstance(client, HTTPHandler): - sync_httpx_client = _get_httpx_client() + sync_httpx_client = _get_ssrf_safe_sync_client() else: sync_httpx_client = client diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index e11d8532dbf3..47236e9cd8f9 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -1076,6 +1076,7 @@ def __init__( disable_default_headers: Optional[ bool ] = False, # arize phoenix returns different API responses when user agent header in request + transport: Optional[httpx.BaseTransport] = None, ): if timeout is None: timeout = _DEFAULT_TIMEOUT @@ -1091,11 +1092,11 @@ def __init__( default_headers = get_default_headers() if not disable_default_headers else None if client is None: - transport = self._create_sync_transport() + _transport = transport if transport is not None else self._create_sync_transport() # Create a client with a connection pool self.client = httpx.Client( - transport=transport, + transport=_transport, timeout=timeout, verify=ssl_config, cert=cert, diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_handler.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_handler.py index 789c88d66f88..f4c9a77833fa 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_handler.py @@ -54,12 +54,16 @@ def test_get_async_client_session_with_instance_session(self): assert result is instance_session - @patch("aiohttp.ClientSession") - def test_get_async_client_session_create_new(self, mock_client_session): + @patch("litellm.llms.custom_httpx.aiohttp_handler.aiohttp.TCPConnector") + @patch("litellm.llms.custom_httpx.aiohttp_handler.aiohttp.ClientSession") + def test_get_async_client_session_create_new( + self, mock_client_session, mock_tcp_connector + ): """Test _get_async_client_session creates new session when none provided""" handler = BaseLLMAIOHTTPHandler() mock_session_instance = Mock() mock_client_session.return_value = mock_session_instance + mock_tcp_connector.return_value = Mock() result = handler._get_async_client_session() @@ -151,12 +155,15 @@ def test_session_ownership_tracking(self): handler2 = BaseLLMAIOHTTPHandler() assert handler2._owns_session - with patch("aiohttp.ClientSession") as mock_client_session: - mock_session_instance = Mock() - mock_client_session.return_value = mock_session_instance + with patch("litellm.llms.custom_httpx.aiohttp_handler.aiohttp.TCPConnector"): + with patch( + "litellm.llms.custom_httpx.aiohttp_handler.aiohttp.ClientSession" + ) as mock_client_session: + mock_session_instance = Mock() + mock_client_session.return_value = mock_session_instance - handler2._get_async_client_session() - assert handler2._owns_session + handler2._get_async_client_session() + assert handler2._owns_session @pytest.mark.asyncio async def test_context_manager_pattern_compatibility(self): @@ -180,8 +187,9 @@ async def test_context_manager_pattern_compatibility(self): # Verify cleanup happened mock_session.close.assert_called_once() + @patch("litellm.llms.custom_httpx.aiohttp_handler.aiohttp.TCPConnector") @patch("litellm.llms.custom_httpx.aiohttp_handler.aiohttp.ClientSession") - def test_lazy_session_creation(self, mock_client_session): + def test_lazy_session_creation(self, mock_client_session, mock_tcp_connector): """Test that session is created lazily only when needed""" handler = BaseLLMAIOHTTPHandler() @@ -192,6 +200,7 @@ def test_lazy_session_creation(self, mock_client_session): # Session should be created when requested mock_session_instance = Mock() mock_client_session.return_value = mock_session_instance + mock_tcp_connector.return_value = Mock() session = handler._get_async_client_session() @@ -323,18 +332,37 @@ def test_create_client_session_with_connector(self, mock_client_session): mock_client_session.assert_called_once_with(connector=mock_connector) assert result is mock_session_instance - @patch("aiohttp.ClientSession") - def test_create_client_session_default(self, mock_client_session): - """Test default session creation when no transport/connector provided""" + @patch("litellm.llms.custom_httpx.aiohttp_handler.aiohttp.TCPConnector") + @patch("litellm.llms.custom_httpx.aiohttp_handler.aiohttp.ClientSession") + def test_create_client_session_default( + self, mock_client_session, mock_tcp_connector + ): + """Test default session creation attaches SSRFGuardResolver via TCPConnector.""" mock_session_instance = Mock() mock_client_session.return_value = mock_session_instance + mock_connector_instance = Mock() + mock_tcp_connector.return_value = mock_connector_instance handler = BaseLLMAIOHTTPHandler() result = handler._create_client_session_with_transport() - # Should create default session - mock_client_session.assert_called_once_with() + # Verify TCPConnector was created with an SSRFGuardResolver + mock_tcp_connector.assert_called_once() + _, kwargs = mock_tcp_connector.call_args + from litellm.llms.custom_httpx.aiohttp_handler import _SSRFGuardResolver + + assert isinstance(kwargs.get("resolver"), _SSRFGuardResolver) + + # Verify ClientSession received the connector and SSRF trace config + mock_client_session.assert_called_once() + _, call_kwargs = mock_client_session.call_args + assert call_kwargs.get("connector") is mock_connector_instance + trace_configs = call_kwargs.get("trace_configs", []) + assert len(trace_configs) == 1 + import aiohttp as _aiohttp + + assert isinstance(trace_configs[0], _aiohttp.TraceConfig) assert result is mock_session_instance def test_get_or_create_transport(self): @@ -408,6 +436,32 @@ async def test_close_transport_without_aclose_method(self): # Should not raise any exceptions await handler.close() + def test_del_without_running_loop_uses_new_event_loop(self): + """__del__ falls back to asyncio.new_event_loop() when no loop is running. + + This covers the RuntimeError branch added to replace the deprecated + asyncio.get_event_loop() call — the fallback is needed during interpreter + shutdown or in synchronous contexts where no event loop is active. + """ + mock_session = Mock() + mock_session.closed = False + + handler = BaseLLMAIOHTTPHandler() + handler.client_session = mock_session + handler._owns_session = True + + with patch( + "asyncio.get_running_loop", side_effect=RuntimeError("no running loop") + ): + with patch("asyncio.new_event_loop") as mock_new_loop: + mock_loop = Mock() + mock_new_loop.return_value = mock_loop + handler.__del__() + + mock_new_loop.assert_called_once() + mock_loop.run_until_complete.assert_called_once() + mock_loop.close.assert_called_once() + def test_transport_priority_hierarchy(self): """Test that session creation follows the right priority: transport > connector > default""" # Test with transport having _get_valid_client_session diff --git a/tests/test_litellm/llms/test_aiohttp_ssrf_protection.py b/tests/test_litellm/llms/test_aiohttp_ssrf_protection.py new file mode 100644 index 000000000000..ae2b716dc366 --- /dev/null +++ b/tests/test_litellm/llms/test_aiohttp_ssrf_protection.py @@ -0,0 +1,839 @@ +import asyncio +import ipaddress +import unittest.mock +import pytest +from unittest.mock import Mock, patch + +import httpx + +import litellm +from litellm.llms.custom_httpx.aiohttp_handler import ( + _SSRFGuardResolver, + _SSRFGuardTransport, + _assert_not_private_ip_literal, + _assert_not_private_url, + _get_ssrf_safe_sync_client, + _is_blocked_address, +) + + +class TestBlockedAddress: + def test_ipv4_mapped_ipv6_private_blocked(self): + addr = ipaddress.ip_address("::ffff:10.0.0.1") + assert _is_blocked_address(addr) + + def test_ipv4_mapped_ipv6_public_allowed(self): + addr = ipaddress.ip_address("::ffff:104.18.7.8") + assert not _is_blocked_address(addr) + + def test_ipv6_link_local_blocked(self): + assert _is_blocked_address(ipaddress.ip_address("fe80::1")) + + def test_ipv6_ula_blocked(self): + assert _is_blocked_address(ipaddress.ip_address("fc00::1")) + + def test_ipv6_unspecified_blocked(self): + assert _is_blocked_address(ipaddress.ip_address("::")) + + def test_0_0_0_0_blocked(self): + assert _is_blocked_address(ipaddress.ip_address("0.0.0.0")) + + def test_benchmarking_range_blocked(self): + assert _is_blocked_address(ipaddress.ip_address("198.18.0.1")) + + def test_class_e_reserved_blocked(self): + assert _is_blocked_address(ipaddress.ip_address("240.0.0.1")) + + def test_documentation_range_blocked(self): + assert _is_blocked_address(ipaddress.ip_address("192.0.2.1")) + + +class TestAiohttpSSRFProtection: + def test_aws_metadata_endpoint_blocked(self): + with pytest.raises(ValueError, match="private/reserved"): + _assert_not_private_url("http://169.254.169.254/latest/meta-data/") + + def test_localhost_blocked(self): + with pytest.raises(ValueError, match="private/reserved"): + _assert_not_private_url("http://127.0.0.1/admin") + + def test_private_10_network_blocked(self): + with pytest.raises(ValueError, match="private/reserved"): + _assert_not_private_url("http://10.0.0.1/internal") + + def test_private_172_16_network_blocked(self): + with pytest.raises(ValueError, match="private/reserved"): + _assert_not_private_url("http://172.16.0.1/internal") + + def test_private_192_168_network_blocked(self): + with pytest.raises(ValueError, match="private/reserved"): + _assert_not_private_url("http://192.168.1.1/internal") + + def test_cgnat_blocked(self): + with pytest.raises(ValueError, match="private/reserved"): + _assert_not_private_url("http://100.64.0.1/internal") + + def test_all_dns_answers_checked(self): + with patch( + "socket.getaddrinfo", + return_value=[ + (None, None, None, None, ("104.18.7.8", None)), + (None, None, None, None, ("10.0.0.1", None)), + ], + ): + with pytest.raises(ValueError, match="private/reserved"): + _assert_not_private_url("https://evil-rebinding.example.com/") + + def test_public_ip_allowed(self): + with patch( + "socket.getaddrinfo", + return_value=[(None, None, None, None, ("104.18.7.8", None))], + ): + _assert_not_private_url("https://api.openai.com/v1/chat/completions") + + def test_empty_hostname_allowed(self): + _assert_not_private_url("not-a-url") + + def test_dns_failure_does_not_block(self): + import socket as _socket + + with patch("socket.getaddrinfo", side_effect=_socket.gaierror("DNS fail")): + _assert_not_private_url("https://nonexistent.invalid/path") + + def test_unparseable_ip_in_dns_answer_skipped(self): + # If getaddrinfo returns a non-IP string (edge case), it should be skipped + with patch( + "socket.getaddrinfo", + return_value=[(None, None, None, None, ("not-an-ip", None))], + ): + _assert_not_private_url("https://example.com/") # should not raise + + +class TestAllowInternalIpsOptOut: + """allow_requests_to_internal_ips=True disables SSRF protection for on-prem use.""" + + def setup_method(self): + self._original = litellm.allow_requests_to_internal_ips + + def teardown_method(self): + litellm.allow_requests_to_internal_ips = self._original + + def test_private_ip_allowed_when_flag_set(self): + litellm.allow_requests_to_internal_ips = True + _assert_not_private_url("http://10.0.0.1/internal") # must not raise + + def test_localhost_allowed_when_flag_set(self): + litellm.allow_requests_to_internal_ips = True + _assert_not_private_url("http://127.0.0.1:11434/api/chat") # Ollama local + + def test_aws_metadata_allowed_when_flag_set(self): + litellm.allow_requests_to_internal_ips = True + _assert_not_private_url("http://169.254.169.254/latest/meta-data/") + + def test_flag_false_still_blocks(self): + litellm.allow_requests_to_internal_ips = False + with pytest.raises(ValueError, match="private/reserved"): + _assert_not_private_url("http://10.0.0.1/internal") + + @pytest.mark.asyncio + async def test_resolver_allows_private_when_flag_set(self): + litellm.allow_requests_to_internal_ips = True + resolver = _SSRFGuardResolver() + mock_infos = [(2, 1, 6, "", ("10.0.0.1", 443))] + + async def run(): + loop = asyncio.get_running_loop() + with patch.object(loop, "getaddrinfo", return_value=mock_infos): + result = await resolver.resolve("internal.corp", 443) + assert result[0]["host"] == "10.0.0.1" + + await run() + + @pytest.mark.asyncio + async def test_resolver_blocks_private_when_flag_false(self): + litellm.allow_requests_to_internal_ips = False + resolver = _SSRFGuardResolver() + mock_infos = [(2, 1, 6, "", ("10.0.0.1", 443))] + + async def run(): + loop = asyncio.get_running_loop() + with patch.object(loop, "getaddrinfo", return_value=mock_infos): + with pytest.raises(ValueError, match="private/reserved"): + await resolver.resolve("evil.internal", 443) + + await run() + + +class TestAssertNotPrivateIpLiteral: + """Tests for the IP-literal bypass guard on the async path.""" + + def test_aws_metadata_ip_literal_blocked(self): + with pytest.raises(ValueError, match="private/reserved"): + _assert_not_private_ip_literal("http://169.254.169.254/latest/meta-data/") + + def test_localhost_ip_literal_blocked(self): + with pytest.raises(ValueError, match="private/reserved"): + _assert_not_private_ip_literal("http://127.0.0.1/admin") + + def test_private_10_network_ip_literal_blocked(self): + with pytest.raises(ValueError, match="private/reserved"): + _assert_not_private_ip_literal("http://10.0.0.1/internal") + + def test_hostname_not_blocked(self): + # Hostnames are not IP literals — resolver handles them + _assert_not_private_ip_literal("https://api.openai.com/v1/chat/completions") + + def test_public_ip_literal_allowed(self): + _assert_not_private_ip_literal("https://104.18.7.8/v1/chat/completions") + + def test_ipv6_loopback_literal_blocked(self): + with pytest.raises(ValueError, match="private/reserved"): + _assert_not_private_ip_literal("http://[::1]/admin") + + def test_ipv4_mapped_ipv6_private_literal_blocked(self): + with pytest.raises(ValueError, match="private/reserved"): + _assert_not_private_ip_literal("http://[::ffff:10.0.0.1]/internal") + + def test_flag_disables_check(self): + litellm.allow_requests_to_internal_ips = True + try: + _assert_not_private_ip_literal("http://169.254.169.254/latest/meta-data/") + finally: + litellm.allow_requests_to_internal_ips = False + + +class TestSSRFGuardOnRequestMethods: + """Verify SSRF protection is enforced on both sync and async request paths.""" + + @pytest.mark.asyncio + async def test_make_common_async_call_blocks_ip_literal_without_dns(self): + """Async path blocks IP-literal api_base at preflight (no DNS I/O needed). + This closes the TCPConnector bypass where aiohttp skips the custom resolver + when the host is already an IP address.""" + from unittest.mock import AsyncMock, Mock + + from litellm.llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler + + handler = BaseLLMAIOHTTPHandler() + mock_config = Mock() + mock_config.max_retry_on_unprocessable_entity_error = 1 + mock_session = AsyncMock() + + with pytest.raises(ValueError, match="private/reserved"): + await handler._make_common_async_call( + async_client_session=mock_session, + provider_config=mock_config, + api_base="http://169.254.169.254/latest/meta-data/", + headers={}, + data={}, + timeout=30, + litellm_params={}, + ) + + @pytest.mark.asyncio + async def test_make_common_async_call_hostname_defers_to_resolver(self): + """Async path does NOT do a blocking DNS preflight for hostname-based URLs. + SSRF protection for those is handled by _SSRFGuardResolver at TCP-connect time. + """ + from unittest.mock import AsyncMock, Mock + + from litellm.llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler + + handler = BaseLLMAIOHTTPHandler() + mock_config = Mock() + mock_config.max_retry_on_unprocessable_entity_error = 1 + mock_session = AsyncMock() + mock_response = AsyncMock() + mock_response.ok = True + mock_session.post.return_value.__aenter__ = AsyncMock( + return_value=mock_response + ) + mock_session.post.return_value.__aexit__ = AsyncMock(return_value=False) + + # Hostname-based private URL must NOT raise at preflight — resolver handles it + try: + await handler._make_common_async_call( + async_client_session=mock_session, + provider_config=mock_config, + api_base="http://internal.corp/api", + headers={}, + data={}, + timeout=30, + litellm_params={}, + ) + except ValueError as e: + pytest.fail(f"Async path must not do a blocking DNS preflight: {e}") + except Exception: + pass # Other errors (e.g. from mock) are fine + + def test_make_common_sync_call_blocks_private_ip(self): + """Sync path _assert_not_private_url preflight still blocks private IPs + for externally-supplied clients (defence-in-depth). Internally created + clients use _SSRFGuardTransport for connect-time validation.""" + from litellm.llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler + + handler = BaseLLMAIOHTTPHandler() + mock_config = Mock() + mock_config.max_retry_on_unprocessable_entity_error = 1 + mock_sync_client = Mock() + + with pytest.raises(ValueError, match="private/reserved"): + handler._make_common_sync_call( + sync_httpx_client=mock_sync_client, + provider_config=mock_config, + api_base="http://10.0.0.1/internal", + headers={}, + data={}, + timeout=30, + litellm_params={}, + ) + + def test_make_common_sync_call_succeeds_for_public_url(self): + """Sync code path still works for legitimate (public-IP) API bases after + introducing _SSRFGuardTransport. Uses _get_ssrf_safe_sync_client() as the + internally-created client — the same path taken by BaseLLMAIOHTTPHandler + when no explicit client is supplied. + + Verifies both that: + - the transport's validation passes (public IP, no block), and + - the request is forwarded to the pinned IP (no second DNS lookup). + """ + from litellm.llms.custom_httpx.aiohttp_handler import ( + BaseLLMAIOHTTPHandler, + _get_ssrf_safe_sync_client, + ) + + handler = BaseLLMAIOHTTPHandler() + mock_config = Mock() + mock_config.max_retry_on_unprocessable_entity_error = 1 + mock_config.should_retry_llm_api_inside_llm_translation_on_http_error = Mock( + return_value=False + ) + + # Simulate DNS resolving to a public IP (legitimate provider) + public_dns = [(2, 1, 6, "", ("104.18.7.8", 443))] + + # httpx.Client._send_single_request asserts isinstance(response.stream, + # SyncByteStream) after calling transport.handle_request, so we must + # return a real httpx.Response (not a plain Mock). + fake_response = httpx.Response( + status_code=200, + content=b'{"id": "chatcmpl-test", "choices": []}', + headers={"content-type": "application/json"}, + ) + + ssrf_safe_client = _get_ssrf_safe_sync_client() + forwarded_hosts: list = [] + + def spy_handle(self_transport, request): + forwarded_hosts.append(request.url.host) + return fake_response + + with patch("socket.getaddrinfo", return_value=public_dns): + with patch.object(httpx.HTTPTransport, "handle_request", spy_handle): + result = handler._make_common_sync_call( + sync_httpx_client=ssrf_safe_client, + provider_config=mock_config, + api_base="https://api.openai.com/v1/chat/completions", + headers={"Authorization": "Bearer sk-test"}, + data={"model": "gpt-4", "messages": []}, + timeout=30, + litellm_params={}, + ) + + assert result is fake_response + # Transport received the pinned IP, proving no second DNS resolution occurs + assert forwarded_hosts == ["104.18.7.8"] + + +class TestSSRFGuardResolver: + """Tests for the async resolver that eliminates TOCTOU DNS rebinding.""" + + def _run(self, coro): + return asyncio.run(coro) + + def test_private_ip_blocked_at_connection_time(self): + resolver = _SSRFGuardResolver() + mock_infos = [ + (2, 1, 6, "", ("10.0.0.1", 443)), + ] + + async def run(): + loop = asyncio.get_running_loop() + with patch.object(loop, "getaddrinfo", return_value=mock_infos): + with pytest.raises(ValueError, match="private/reserved"): + await resolver.resolve("evil.internal", 443) + + self._run(run()) + + def test_public_ip_passes_resolver(self): + resolver = _SSRFGuardResolver() + mock_infos = [ + (2, 1, 6, "", ("104.18.7.8", 443)), + ] + + async def run(): + loop = asyncio.get_running_loop() + with patch.object(loop, "getaddrinfo", return_value=mock_infos): + result = await resolver.resolve("api.openai.com", 443) + assert result[0]["host"] == "104.18.7.8" + + self._run(run()) + + def test_all_answers_checked_by_resolver(self): + resolver = _SSRFGuardResolver() + mock_infos = [ + (2, 1, 6, "", ("104.18.7.8", 443)), + (2, 1, 6, "", ("169.254.169.254", 443)), + ] + + async def run(): + loop = asyncio.get_running_loop() + with patch.object(loop, "getaddrinfo", return_value=mock_infos): + with pytest.raises(ValueError, match="private/reserved"): + await resolver.resolve("rebinding.example.com", 443) + + self._run(run()) + + def test_resolver_dns_failure_propagates(self): + import socket as _socket + + resolver = _SSRFGuardResolver() + + async def run(): + loop = asyncio.get_running_loop() + with patch.object( + loop, "getaddrinfo", side_effect=_socket.gaierror("DNS fail") + ): + with pytest.raises(_socket.gaierror): + await resolver.resolve("nonexistent.invalid", 443) + + self._run(run()) + + def test_resolver_unparseable_ip_skipped(self): + resolver = _SSRFGuardResolver() + mock_infos = [ + (2, 1, 6, "", ("not-an-ip", 443)), + (2, 1, 6, "", ("104.18.7.8", 443)), + ] + + async def run(): + loop = asyncio.get_running_loop() + with patch.object(loop, "getaddrinfo", return_value=mock_infos): + result = await resolver.resolve("example.com", 443) + assert any(r["host"] == "104.18.7.8" for r in result) + + self._run(run()) + + @pytest.mark.asyncio + async def test_resolver_close_is_noop(self): + resolver = _SSRFGuardResolver() + await resolver.close() # Should not raise + + +class TestSSRFGuardTransport: + """_SSRFGuardTransport eliminates DNS-rebinding TOCTOU on the sync httpx path. + + It resolves the hostname, validates every returned IP, then rewrites the + URL to an IP literal so httpcore connects directly — no second DNS lookup. + """ + + def test_transport_blocks_hostname_resolving_to_private_ip(self): + """Raises when DNS returns a private IP for a hostname-based URL.""" + transport = _SSRFGuardTransport() + mock_infos = [(2, 1, 6, "", ("10.0.0.1", 443))] + request = httpx.Request("POST", "https://evil.internal/v1/chat") + + with patch("socket.getaddrinfo", return_value=mock_infos): + with pytest.raises(ValueError, match="private/reserved"): + transport.handle_request(request) + + def test_transport_blocks_aws_metadata_endpoint(self): + transport = _SSRFGuardTransport() + mock_infos = [(2, 1, 6, "", ("169.254.169.254", 80))] + request = httpx.Request("GET", "http://metadata.example.com/latest/meta-data/") + + with patch("socket.getaddrinfo", return_value=mock_infos): + with pytest.raises(ValueError, match="private/reserved"): + transport.handle_request(request) + + def test_transport_blocks_all_dns_answers(self): + """Raises when ANY answer resolves to a private IP (A-record rotation).""" + transport = _SSRFGuardTransport() + mock_infos = [ + (2, 1, 6, "", ("104.18.7.8", 443)), # public — passes + (2, 1, 6, "", ("169.254.169.254", 443)), # private — must still block + ] + request = httpx.Request("POST", "https://rebinding.example.com/v1") + + with patch("socket.getaddrinfo", return_value=mock_infos): + with pytest.raises(ValueError, match="private/reserved"): + transport.handle_request(request) + + def test_transport_pins_ip_and_preserves_hostname(self): + """Forwards request with URL rewritten to pinned IP; original hostname + in Host header and sni_hostname extension for TLS.""" + transport = _SSRFGuardTransport() + mock_infos = [(2, 1, 6, "", ("104.18.7.8", 443))] + request = httpx.Request( + "POST", + "https://api.openai.com/v1/chat/completions", + content=b'{"model":"gpt-4"}', + ) + mock_response = Mock(spec=httpx.Response) + + with patch("socket.getaddrinfo", return_value=mock_infos): + with patch.object( + httpx.HTTPTransport, "handle_request", return_value=mock_response + ) as mock_super: + result = transport.handle_request(request) + + assert result is mock_response + forwarded: httpx.Request = mock_super.call_args[0][0] + # URL host is the pinned IP — httpcore skips DNS, TOCTOU window closed + assert forwarded.url.host == "104.18.7.8" + # Original hostname preserved for virtual hosting + assert forwarded.headers.get("host") == "api.openai.com" + # Original hostname preserved for TLS SNI so cert validation passes + assert forwarded.extensions.get("sni_hostname") == "api.openai.com" + + def test_transport_dns_failure_passes_through(self): + """DNS failure is not treated as a block — httpcore propagates the error.""" + import socket as _socket + + transport = _SSRFGuardTransport() + request = httpx.Request("POST", "https://nonexistent.invalid/v1") + mock_response = Mock(spec=httpx.Response) + + with patch("socket.getaddrinfo", side_effect=_socket.gaierror("DNS fail")): + with patch.object( + httpx.HTTPTransport, "handle_request", return_value=mock_response + ): + result = transport.handle_request(request) + + assert result is mock_response + + def test_transport_public_ip_literal_allowed(self): + """Public IP-literal URLs are allowed and bypass DNS (no resolution needed).""" + transport = _SSRFGuardTransport() + request = httpx.Request("POST", "https://104.18.7.8/v1/chat") + mock_response = Mock(spec=httpx.Response) + + # getaddrinfo must NOT be called for an IP-literal URL + with patch("socket.getaddrinfo") as mock_dns: + with patch.object( + httpx.HTTPTransport, "handle_request", return_value=mock_response + ): + result = transport.handle_request(request) + + mock_dns.assert_not_called() + assert result is mock_response + + def test_transport_blocks_private_ip_literal(self): + """Private IP-literal in URL is blocked directly — no DNS lookup required.""" + transport = _SSRFGuardTransport() + request = httpx.Request("POST", "http://10.0.0.1/api") + + with patch("socket.getaddrinfo") as mock_dns: + with pytest.raises(ValueError, match="private/reserved"): + transport.handle_request(request) + + mock_dns.assert_not_called() + + def test_transport_blocks_aws_metadata_ip_literal(self): + """169.254.169.254 as an IP literal is blocked without DNS lookup.""" + transport = _SSRFGuardTransport() + request = httpx.Request("GET", "http://169.254.169.254/latest/meta-data/") + + with patch("socket.getaddrinfo") as mock_dns: + with pytest.raises(ValueError, match="private/reserved"): + transport.handle_request(request) + + mock_dns.assert_not_called() + + def test_transport_blocks_ipv6_loopback_literal(self): + """IPv6 loopback ::1 as an IP literal is blocked without DNS lookup.""" + transport = _SSRFGuardTransport() + request = httpx.Request("GET", "http://[::1]/admin") + + with patch("socket.getaddrinfo") as mock_dns: + with pytest.raises(ValueError, match="private/reserved"): + transport.handle_request(request) + + mock_dns.assert_not_called() + + def test_transport_ip_literal_bypassed_when_flag_set(self): + """allow_requests_to_internal_ips=True skips IP-literal validation too.""" + transport = _SSRFGuardTransport() + request = httpx.Request("POST", "http://10.0.0.1/internal") + mock_response = Mock(spec=httpx.Response) + + litellm.allow_requests_to_internal_ips = True + try: + with patch.object( + httpx.HTTPTransport, "handle_request", return_value=mock_response + ): + result = transport.handle_request(request) + assert result is mock_response + finally: + litellm.allow_requests_to_internal_ips = False + + def test_transport_bypassed_when_flag_set(self): + """allow_requests_to_internal_ips=True disables all transport-level checks.""" + transport = _SSRFGuardTransport() + mock_infos = [(2, 1, 6, "", ("10.0.0.1", 443))] + request = httpx.Request("POST", "https://internal.corp/v1/chat") + mock_response = Mock(spec=httpx.Response) + + litellm.allow_requests_to_internal_ips = True + try: + with patch("socket.getaddrinfo", return_value=mock_infos): + with patch.object( + httpx.HTTPTransport, "handle_request", return_value=mock_response + ): + result = transport.handle_request(request) + assert result is mock_response + finally: + litellm.allow_requests_to_internal_ips = False + + def test_get_ssrf_safe_sync_client_uses_guard_transport(self): + """_get_ssrf_safe_sync_client() produces an HTTPHandler backed by + _SSRFGuardTransport so the sync code path gets connect-time validation.""" + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + handler = _get_ssrf_safe_sync_client() + assert isinstance(handler, HTTPHandler) + # The underlying httpx.Client must use _SSRFGuardTransport + assert isinstance(handler.client._transport, _SSRFGuardTransport) + + def test_ssrf_safe_client_blocks_private_ip_at_connect_time(self): + """End-to-end: _get_ssrf_safe_sync_client raises on a private-IP hostname + even when called via HTTPHandler.post(), proving the sync code path works.""" + handler = _get_ssrf_safe_sync_client() + mock_infos = [(2, 1, 6, "", ("10.0.0.1", 443))] + + with patch("socket.getaddrinfo", return_value=mock_infos): + with pytest.raises(ValueError, match="private/reserved"): + handler.post("https://evil.internal/v1/chat", json={"test": True}) + + def test_ssrf_safe_client_blocks_private_ip_literal(self): + """End-to-end: _get_ssrf_safe_sync_client raises on a private IP-literal URL + even without DNS, proving the sync path catches direct IP addresses.""" + handler = _get_ssrf_safe_sync_client() + + # No DNS mock needed — the transport must block before any resolution + with patch("socket.getaddrinfo") as mock_dns: + with pytest.raises(ValueError, match="private/reserved"): + handler.post("http://10.0.0.1/api", json={"test": True}) + + mock_dns.assert_not_called() + + def test_transport_skips_unparseable_ip_in_dns_answers(self): + """Unparseable entries in getaddrinfo answers are skipped; valid entries + still checked — mirrors _SSRFGuardResolver behaviour on the sync path.""" + transport = _SSRFGuardTransport() + # One entry is unparseable; the other is a safe public IP. + mock_infos = [ + (2, 1, 6, "", ("not-an-ip", 443)), + (2, 1, 6, "", ("104.18.7.8", 443)), + ] + request = httpx.Request("POST", "https://api.example.com/v1/chat") + mock_response = Mock(spec=httpx.Response) + + with patch("socket.getaddrinfo", return_value=mock_infos): + with patch.object( + httpx.HTTPTransport, "handle_request", return_value=mock_response + ): + result = transport.handle_request(request) + + assert result is mock_response + + def test_transport_blocks_redirect_to_private_ip_literal(self): + """Redirect targets with private IP literals are blocked by the transport. + httpx calls handle_request for each redirect hop, so the guard applies.""" + transport = _SSRFGuardTransport() + # Simulate httpx presenting the redirect target URL to the transport. + redirect_request = httpx.Request( + "GET", "http://169.254.169.254/latest/meta-data/" + ) + + with patch("socket.getaddrinfo") as mock_dns: + with pytest.raises(ValueError, match="private/reserved"): + transport.handle_request(redirect_request) + + mock_dns.assert_not_called() + + def test_transport_blocks_redirect_to_private_hostname(self): + """Redirect targets resolving to private IPs are blocked by the transport.""" + transport = _SSRFGuardTransport() + mock_infos = [(2, 1, 6, "", ("10.0.0.1", 80))] + redirect_request = httpx.Request("GET", "http://internal.corp/sensitive") + + with patch("socket.getaddrinfo", return_value=mock_infos): + with pytest.raises(ValueError, match="private/reserved"): + transport.handle_request(redirect_request) + + +class TestSSRFSafeClientCachingAndSSL: + """_get_ssrf_safe_sync_client must reuse a cached client and forward SSL config.""" + + def setup_method(self): + # Clear the cache before each test so tests are independent + cache = getattr(litellm, "in_memory_llm_clients_cache", None) + if cache is not None: + try: + cache.delete_cache("ssrf_safe_httpx_client") + except Exception: + pass + + def test_client_is_cached_across_calls(self): + """Two consecutive calls return the same HTTPHandler instance.""" + client1 = _get_ssrf_safe_sync_client() + client2 = _get_ssrf_safe_sync_client() + assert client1 is client2 + + def test_ssl_config_is_forwarded_to_transport(self): + """get_ssl_configuration() return value is passed as verify= to + _SSRFGuardTransport so custom CA bundles and ssl_context objects are + honoured by the SSRF-guarded sync client.""" + from litellm.llms.custom_httpx.aiohttp_handler import ( + _SSRF_SAFE_CLIENT_CACHE_KEY, + ) + + cache = getattr(litellm, "in_memory_llm_clients_cache", None) + if cache: + try: + cache.delete_cache(_SSRF_SAFE_CLIENT_CACHE_KEY) + except Exception: + pass + + sentinel_ssl = object() # unique sentinel — proves it was forwarded + with patch( + "litellm.llms.custom_httpx.aiohttp_handler.get_ssl_configuration", + return_value=sentinel_ssl, + ): + with patch.object( + _SSRFGuardTransport, "__init__", wraps=_SSRFGuardTransport.__init__ + ) as spy: + try: + _get_ssrf_safe_sync_client() + except Exception: + pass # construction may fail with a non-SSL sentinel + if spy.call_args: + _, kwargs = spy.call_args + assert kwargs.get("verify") is sentinel_ssl + + def test_ssl_certificate_is_forwarded_to_transport(self): + """SSL_CERTIFICATE env-var (used for mTLS / custom CA) is forwarded to + _SSRFGuardTransport as cert= so it is not silently dropped.""" + from litellm.llms.custom_httpx.aiohttp_handler import ( + _SSRF_SAFE_CLIENT_CACHE_KEY, + ) + + cache = getattr(litellm, "in_memory_llm_clients_cache", None) + if cache: + try: + cache.delete_cache(_SSRF_SAFE_CLIENT_CACHE_KEY) + except Exception: + pass + + with patch( + "litellm.llms.custom_httpx.aiohttp_handler.os.getenv", + side_effect=lambda k, d=None: ( + "/path/to/client.pem" if k == "SSL_CERTIFICATE" else d + ), + ): + with patch.object( + _SSRFGuardTransport, "__init__", wraps=_SSRFGuardTransport.__init__ + ) as spy: + try: + _get_ssrf_safe_sync_client() + except Exception: + pass + if spy.call_args: + _, kwargs = spy.call_args + assert kwargs.get("cert") == "/path/to/client.pem" + + +class TestSSRFTraceConfig: + """Tests for the on_request_start trace hook that blocks IP-literal aiohttp + redirect targets — the gap _SSRFGuardResolver cannot cover because aiohttp + skips DNS resolution when the Location header is already an IP literal.""" + + def _run(self, coro): + return asyncio.run(coro) + + def _mock_params(self, url: str): + params = unittest.mock.Mock() + params.url = url + return params + + def test_trace_callback_blocks_aws_metadata_ip_literal(self): + from litellm.llms.custom_httpx.aiohttp_handler import _on_ssrf_request_start + + async def run(): + with pytest.raises(ValueError, match="private/reserved"): + await _on_ssrf_request_start( + unittest.mock.Mock(), + unittest.mock.Mock(), + self._mock_params("http://169.254.169.254/latest/meta-data/"), + ) + + self._run(run()) + + def test_trace_callback_blocks_private_rfc1918_ip_literal(self): + from litellm.llms.custom_httpx.aiohttp_handler import _on_ssrf_request_start + + async def run(): + with pytest.raises(ValueError, match="private/reserved"): + await _on_ssrf_request_start( + unittest.mock.Mock(), + unittest.mock.Mock(), + self._mock_params("http://10.0.0.1/internal"), + ) + + self._run(run()) + + def test_trace_callback_allows_public_ip_literal(self): + from litellm.llms.custom_httpx.aiohttp_handler import _on_ssrf_request_start + + async def run(): + # Public IP — must not raise + await _on_ssrf_request_start( + unittest.mock.Mock(), + unittest.mock.Mock(), + self._mock_params("https://104.18.7.8/v1/chat/completions"), + ) + + self._run(run()) + + def test_trace_callback_passes_hostname_urls_through(self): + """Hostname URLs are not IP literals — resolver handles them at DNS time.""" + from litellm.llms.custom_httpx.aiohttp_handler import _on_ssrf_request_start + + async def run(): + await _on_ssrf_request_start( + unittest.mock.Mock(), + unittest.mock.Mock(), + self._mock_params("https://api.openai.com/v1/chat/completions"), + ) + + self._run(run()) + + def test_trace_callback_bypassed_when_flag_set(self): + from litellm.llms.custom_httpx.aiohttp_handler import _on_ssrf_request_start + + litellm.allow_requests_to_internal_ips = True + try: + + async def run(): + # Private IP literal must NOT raise when flag is set + await _on_ssrf_request_start( + unittest.mock.Mock(), + unittest.mock.Mock(), + self._mock_params("http://169.254.169.254/latest/meta-data/"), + ) + + self._run(run()) + finally: + litellm.allow_requests_to_internal_ips = False