From 80cea73ccd99823b84df95d0f41185a82bc12e18 Mon Sep 17 00:00:00 2001
From: Drishna Trivedi <drishna4505@gmail.com>
Date: Tue, 19 May 2026 17:36:21 +0530
Subject: [PATCH] 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>
---
 .../llms/test_aiohttp_ssrf_protection.py      | 90 +++++++++++++++++++
 1 file changed, 90 insertions(+)

diff --git a/tests/test_litellm/llms/test_aiohttp_ssrf_protection.py b/tests/test_litellm/llms/test_aiohttp_ssrf_protection.py
index 0c08685eaff..2d114f9775a 100644
--- a/tests/test_litellm/llms/test_aiohttp_ssrf_protection.py
+++ b/tests/test_litellm/llms/test_aiohttp_ssrf_protection.py
@@ -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