test(aiohttp): add coverage for all edge-case branches in SSRF guard

- Unparseable IP in getaddrinfo answer (lines 79-80, 111-112)
- DNS failure in _SSRFGuardResolver.resolve() (lines 105-106)
- _SSRFGuardResolver.close() noop (line 131)
- _make_common_async_call and _make_common_sync_call block private api_base (lines 297, 343)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Drishna Trivedi 2026-05-19 17:36:21 +05:30
parent 537db99539
commit 80cea73ccd

View file

@ -81,6 +81,61 @@ class TestAiohttpSSRFProtection:
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 TestSSRFGuardOnRequestMethods:
"""Verify _assert_not_private_url is actually called in the request paths."""
@pytest.mark.asyncio
async def test_make_common_async_call_blocks_private_ip(self):
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={},
)
def test_make_common_sync_call_blocks_private_ip(self):
from unittest.mock import 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_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={},
)
class TestSSRFGuardResolver:
"""Tests for the async resolver that eliminates TOCTOU DNS rebinding."""
@ -131,3 +186,38 @@ class TestSSRFGuardResolver:
await resolver.resolve("rebinding.example.com", 443)
self._run(run())
def test_resolver_dns_failure_returns_empty(self):
import socket as _socket
resolver = _SSRFGuardResolver()
async def run():
loop = asyncio.get_event_loop()
with patch.object(
loop, "getaddrinfo", side_effect=_socket.gaierror("DNS fail")
):
result = await resolver.resolve("nonexistent.invalid", 443)
assert result == []
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_event_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