Fix/aiohttp ssrf cwe918

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

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,105 @@
 
 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.
+    Used as a fast-fail guard on the sync path (httpx) and as defence-in-depth on async.
+
+    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)"
+            )
+
+
+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 BaseLLMAIOHTTPHandler:
     def __init__(
@@ -95,8 +199,12 @@ 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.
+            session = aiohttp.ClientSession(
+                connector=aiohttp.TCPConnector(resolver=_SSRFGuardResolver())
+            )
             return session
 
     def _get_async_client_session(
@@ -147,26 +255,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 +290,12 @@ async def _make_common_async_call(
             dynamic_client_session=async_client_session
         )
 
+        # SSRF validation on the async path is handled at TCP-connect time by
+        # _SSRFGuardResolver (attached to the TCPConnector).  A synchronous
+        # socket.getaddrinfo() preflight would block the event loop, so it is
+        # intentionally omitted here.  The sync path retains _assert_not_private_url
+        # because httpx has no equivalent connection-time resolver hook.
+
         for i in range(max(max_retry_on_unprocessable_entity_error, 1)):
             try:
                 response = await async_client_session.post(
@@ -235,6 +340,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)):

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,30 @@ 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
+        mock_client_session.assert_called_once_with(connector=mock_connector_instance)
         assert result is mock_session_instance
 
     def test_get_or_create_transport(self):