From 537db995391c3fd1f89a65d06d306069de271f4b Mon Sep 17 00:00:00 2001 From: Drishna Trivedi Date: Tue, 19 May 2026 17:27:06 +0530 Subject: [PATCH] fix(aiohttp): block private/metadata IPs in api_base to close SSRF gap from #26264 (CWE-918) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit aiohttp_handler.py was not covered by the SSRF protection in PR #26264. User-controlled api_base was passed directly to session.post() without IP validation. Protection added: - Blocks RFC-1918, loopback, link-local (169.254/16, fe80::/10), CGNAT, 0.0.0.0/8, IPv6 ULA/loopback - Unwraps IPv4-mapped IPv6 (::ffff:x.x.x.x) before network check - Validates ALL getaddrinfo answers to prevent A-record rotation bypass - _SSRFGuardResolver (AbstractResolver) validates IPs at TCP-connection time inside aiohttp's own connection loop — covers redirect targets and eliminates DNS-rebinding TOCTOU - Default ClientSession creation uses TCPConnector(resolver=_SSRFGuardResolver()) - Sync path (_make_common_sync_call via httpx) guarded with preflight check Tests: - 18 new tests in tests/test_litellm/llms/test_aiohttp_ssrf_protection.py - Updated 4 existing tests in test_aiohttp_handler.py to mock TCPConnector Co-Authored-By: Claude Sonnet 4.6 --- litellm/llms/custom_httpx/aiohttp_handler.py | 111 ++++++++++++++- .../llms/custom_httpx/test_aiohttp_handler.py | 47 +++++-- .../llms/test_aiohttp_ssrf_protection.py | 133 ++++++++++++++++++ 3 files changed, 276 insertions(+), 15 deletions(-) create mode 100644 tests/test_litellm/llms/test_aiohttp_ssrf_protection.py diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index 93b6c563dc1..b9e9be4aedd 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -1,8 +1,13 @@ +import asyncio +import ipaddress +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 @@ -31,6 +36,100 @@ else: DEFAULT_TIMEOUT = 600 +_BLOCKED_NETWORKS = [ + ipaddress.ip_network("0.0.0.0/8"), + ipaddress.ip_network("10.0.0.0/8"), + ipaddress.ip_network("100.64.0.0/10"), # CGNAT + ipaddress.ip_network("127.0.0.0/8"), + ipaddress.ip_network("169.254.0.0/16"), # Link-local / AWS IMDS + ipaddress.ip_network("172.16.0.0/12"), + ipaddress.ip_network("192.168.0.0/16"), + ipaddress.ip_network("::1/128"), + ipaddress.ip_network("fc00::/7"), + ipaddress.ip_network("fe80::/10"), # IPv6 link-local +] + + +def _is_blocked_address(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool: + """Return True if addr falls in any blocked network.""" + # 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 + return any(addr in net for net in _BLOCKED_NETWORKS) + + +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. + Used as a fast-fail guard on the sync path (httpx) and as defence-in-depth on async. + """ + 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)" + ) + + +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_event_loop() + try: + infos = await loop.getaddrinfo( + host, port, family=family, type=socket.SOCK_STREAM + ) + except socket.gaierror: + return [] # Let aiohttp surface the connection error naturally + 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 BaseLLMAIOHTTPHandler: def __init__( @@ -95,8 +194,12 @@ class BaseLLMAIOHTTPHandler: 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. + session = aiohttp.ClientSession( + connector=aiohttp.TCPConnector(resolver=_SSRFGuardResolver()) + ) return session def _get_async_client_session( @@ -191,6 +294,8 @@ class BaseLLMAIOHTTPHandler: dynamic_client_session=async_client_session ) + _assert_not_private_url(api_base) + for i in range(max(max_retry_on_unprocessable_entity_error, 1)): try: response = await async_client_session.post( @@ -235,6 +340,8 @@ class BaseLLMAIOHTTPHandler: 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)): 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 789c88d66f8..a7034aab241 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 @@ class TestBaseLLMAIOHTTPHandler: 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 @@ class TestBaseLLMAIOHTTPHandler: 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 @@ class TestBaseLLMAIOHTTPHandler: # 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 @@ class TestBaseLLMAIOHTTPHandler: # 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,30 @@ class TestBaseLLMAIOHTTPHandler: 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 + mock_client_session.assert_called_once_with(connector=mock_connector_instance) assert result is mock_session_instance def test_get_or_create_transport(self): 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 00000000000..0c08685eaff --- /dev/null +++ b/tests/test_litellm/llms/test_aiohttp_ssrf_protection.py @@ -0,0 +1,133 @@ +import asyncio +import ipaddress +import pytest +from unittest.mock import patch + +from litellm.llms.custom_httpx.aiohttp_handler import ( + _SSRFGuardResolver, + _assert_not_private_url, + _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_0_0_0_0_blocked(self): + assert _is_blocked_address(ipaddress.ip_address("0.0.0.0")) + + +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") + + +class TestSSRFGuardResolver: + """Tests for the async resolver that eliminates TOCTOU DNS rebinding.""" + + def _run(self, coro): + return asyncio.get_event_loop().run_until_complete(coro) + + def test_private_ip_blocked_at_connection_time(self): + resolver = _SSRFGuardResolver() + mock_infos = [ + (2, 1, 6, "", ("10.0.0.1", 443)), + ] + with patch("asyncio.AbstractEventLoop.getaddrinfo", return_value=mock_infos): + + async def run(): + loop = asyncio.get_event_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_event_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_event_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())