diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index b9e9be4aedd..b8704a10b89 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -97,13 +97,13 @@ class _SSRFGuardResolver(AbstractResolver): async def resolve( self, host: str, port: int = 0, family: int = socket.AF_INET ) -> list: - loop = asyncio.get_event_loop() + loop = asyncio.get_running_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 + raise # Propagate so aiohttp wraps it in a ClientConnectorError for info in infos: raw_ip = info[4][0] try: diff --git a/tests/test_litellm/llms/test_aiohttp_ssrf_protection.py b/tests/test_litellm/llms/test_aiohttp_ssrf_protection.py index 2d114f9775a..69d552d9a92 100644 --- a/tests/test_litellm/llms/test_aiohttp_ssrf_protection.py +++ b/tests/test_litellm/llms/test_aiohttp_ssrf_protection.py @@ -141,22 +141,21 @@ 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) + return asyncio.run(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) + async def run(): + loop = asyncio.get_running_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()) + self._run(run()) def test_public_ip_passes_resolver(self): resolver = _SSRFGuardResolver() @@ -165,7 +164,7 @@ class TestSSRFGuardResolver: ] async def run(): - loop = asyncio.get_event_loop() + loop = asyncio.get_running_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" @@ -180,25 +179,25 @@ class TestSSRFGuardResolver: ] async def run(): - loop = asyncio.get_event_loop() + loop = asyncio.get_running_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()) - def test_resolver_dns_failure_returns_empty(self): + def test_resolver_dns_failure_propagates(self): import socket as _socket resolver = _SSRFGuardResolver() async def run(): - loop = asyncio.get_event_loop() + loop = asyncio.get_running_loop() with patch.object( loop, "getaddrinfo", side_effect=_socket.gaierror("DNS fail") ): - result = await resolver.resolve("nonexistent.invalid", 443) - assert result == [] + with pytest.raises(_socket.gaierror): + await resolver.resolve("nonexistent.invalid", 443) self._run(run()) @@ -210,7 +209,7 @@ class TestSSRFGuardResolver: ] async def run(): - loop = asyncio.get_event_loop() + loop = asyncio.get_running_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)