mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
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:
parent
537db99539
commit
80cea73ccd
1 changed files with 90 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue