scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 98

fix(aiohttp): block private/metadata IPs in api_base to close SSRF gap from #26264 (CWE-918)

Browse source 
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 <noreply@anthropic.com>
This commit is contained in:
Drishna Trivedi 2026-05-19 17:27:06 +05:30
parent cff3e0b75e
commit 537db99539
3 changed files with 276 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

111
litellm/llms/custom_httpx/aiohttp_handler.py
View file

@ -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)):

47
tests/test_litellm/llms/custom_httpx/test_aiohttp_handler.py
View file

@ -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):

133
tests/test_litellm/llms/test_aiohttp_ssrf_protection.py Normal file
View file

@ -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())