diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 52f30e31641..6cb281cc700 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -664,7 +664,9 @@ class AsyncHTTPHandler: logging_obj: LiteLLMLoggingObject | None = None, files: RequestFiles | None = None, content: _RequestContent | None = None, + follow_redirects: bool | None = None, ): + _follow_redirects: Final = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT start_time: Final = time.time() try: if timeout is None: @@ -684,7 +686,7 @@ class AsyncHTTPHandler: files=files, content=request_content, ) - response: Final = await self.client.send(req, stream=stream) + response: Final = await self.client.send(req, stream=stream, follow_redirects=_follow_redirects) response.raise_for_status() return response except (httpx.RemoteProtocolError, httpx.ConnectError): @@ -699,6 +701,7 @@ class AsyncHTTPHandler: params=params, headers=headers, stream=stream, + follow_redirects=follow_redirects, ) finally: await new_client.aclose() @@ -732,7 +735,9 @@ class AsyncHTTPHandler: timeout: float | httpx.Timeout | None = None, stream: bool = False, content: _RequestContent | None = None, + follow_redirects: bool | None = None, ): + _follow_redirects: Final = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT try: if timeout is None: timeout = self.timeout @@ -750,7 +755,7 @@ class AsyncHTTPHandler: timeout=timeout, content=request_content, ) - response: Final = await self.client.send(req) + response: Final = await self.client.send(req, follow_redirects=_follow_redirects) response.raise_for_status() return response except (httpx.RemoteProtocolError, httpx.ConnectError): @@ -765,6 +770,7 @@ class AsyncHTTPHandler: params=params, headers=headers, stream=stream, + follow_redirects=follow_redirects, ) finally: await new_client.aclose() @@ -796,7 +802,9 @@ class AsyncHTTPHandler: timeout: float | httpx.Timeout | None = None, stream: bool = False, content: _RequestContent | None = None, + follow_redirects: bool | None = None, ): + _follow_redirects: Final = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT try: if timeout is None: timeout = self.timeout @@ -814,7 +822,7 @@ class AsyncHTTPHandler: timeout=timeout, content=request_content, ) - response: Final = await self.client.send(req) + response: Final = await self.client.send(req, follow_redirects=_follow_redirects) response.raise_for_status() return response except (httpx.RemoteProtocolError, httpx.ConnectError): @@ -829,6 +837,7 @@ class AsyncHTTPHandler: params=params, headers=headers, stream=stream, + follow_redirects=follow_redirects, ) finally: await new_client.aclose() @@ -860,7 +869,9 @@ class AsyncHTTPHandler: timeout: float | httpx.Timeout | None = None, stream: bool = False, content: _RequestContent | None = None, + follow_redirects: bool | None = None, ): + _follow_redirects: Final = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT try: if timeout is None: timeout = self.timeout @@ -878,7 +889,7 @@ class AsyncHTTPHandler: timeout=timeout, content=request_content, ) - response: Final = await self.client.send(req, stream=stream) + response: Final = await self.client.send(req, stream=stream, follow_redirects=_follow_redirects) response.raise_for_status() return response except (httpx.RemoteProtocolError, httpx.ConnectError): @@ -893,6 +904,7 @@ class AsyncHTTPHandler: params=params, headers=headers, stream=stream, + follow_redirects=follow_redirects, ) finally: await new_client.aclose() @@ -911,12 +923,14 @@ class AsyncHTTPHandler: headers: dict | None = None, stream: bool = False, content: _RequestContent | None = None, + follow_redirects: bool | None = None, ): """ Making POST request for a single connection client. Used for retrying connection client errors. """ + _follow_redirects: Final = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT # Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix) request_data, request_content = _prepare_request_data_and_content(data, content) @@ -929,7 +943,7 @@ class AsyncHTTPHandler: headers=headers, content=request_content, ) - response: Final = await client.send(req, stream=stream) + response: Final = await client.send(req, stream=stream, follow_redirects=_follow_redirects) response.raise_for_status() return response diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py index 1553487ee67..89eb79e4ce7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py @@ -469,7 +469,7 @@ async def http_request( try: validated_url, host_header = validate_url(url) headers = {**(headers or {}), "Host": host_header} - except SSRFError as e: + except (SSRFError, ValueError) as e: verbose_proxy_logger.warning("Custom code http_request SSRF blocked: %s", e) return _http_error_response(f"Blocked: {e}") diff --git a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py index b6e7f3b7025..74c85dbc3e2 100644 --- a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py +++ b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py @@ -1,3 +1,4 @@ +import inspect import socket import httpx @@ -226,6 +227,12 @@ def test_missing_apply_guardrail_raises(): class _FakeAsyncClient: + """Mimics the REAL AsyncHTTPHandler method signatures. + + If the primitives pass a kwarg the production handler does not accept, + these fakes raise TypeError exactly like production would. + """ + def __init__(self, response=None): self.response = response self.calls = [] @@ -236,20 +243,123 @@ class _FakeAsyncClient: raise self.response return self.response - async def get(self, **kwargs): - return await self._record(**kwargs) + async def get( + self, + url, + params=None, + headers=None, + follow_redirects=None, + timeout=None, + ): + return await self._record( + url=url, + params=params, + headers=headers, + follow_redirects=follow_redirects, + timeout=timeout, + ) - async def post(self, **kwargs): - return await self._record(**kwargs) + async def post( + self, + url, + data=None, + json=None, + params=None, + headers=None, + timeout=None, + stream=False, + logging_obj=None, + files=None, + content=None, + follow_redirects=None, + ): + return await self._record( + url=url, + data=data, + json=json, + params=params, + headers=headers, + timeout=timeout, + follow_redirects=follow_redirects, + ) - async def put(self, **kwargs): - return await self._record(**kwargs) + async def put( + self, + url, + data=None, + json=None, + params=None, + headers=None, + timeout=None, + stream=False, + content=None, + follow_redirects=None, + ): + return await self._record( + url=url, + data=data, + json=json, + params=params, + headers=headers, + timeout=timeout, + follow_redirects=follow_redirects, + ) - async def delete(self, **kwargs): - return await self._record(**kwargs) + async def patch( + self, + url, + data=None, + json=None, + params=None, + headers=None, + timeout=None, + stream=False, + content=None, + follow_redirects=None, + ): + return await self._record( + url=url, + data=data, + json=json, + params=params, + headers=headers, + timeout=timeout, + follow_redirects=follow_redirects, + ) - async def patch(self, **kwargs): - return await self._record(**kwargs) + async def delete( + self, + url, + data=None, + json=None, + params=None, + headers=None, + timeout=None, + stream=False, + content=None, + follow_redirects=None, + ): + return await self._record( + url=url, + data=data, + json=json, + params=params, + headers=headers, + timeout=timeout, + follow_redirects=follow_redirects, + ) + + +@pytest.mark.parametrize("method", ["get", "post", "put", "patch", "delete"]) +def test_async_http_handler_accepts_follow_redirects(method): + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + params = inspect.signature(getattr(AsyncHTTPHandler, method)).parameters + assert "follow_redirects" in params, ( + f"AsyncHTTPHandler.{method} must accept follow_redirects or the " + "guardrail HTTP primitives cannot disable redirects" + ) + assert params["follow_redirects"].default is None def _ok_response(status_code=200): @@ -276,6 +386,11 @@ def _public_dns(monkeypatch): "http://192.168.1.1/", "http://169.254.169.254/latest/meta-data/iam/security-credentials/", "http://[::1]/", + # Malformed port: is_valid_url accepts it (scheme + netloc), but + # urlparse raises ValueError when validate_url reads parsed.port. + # Must surface as a structured error, not an escaped exception. + "http://example.com:99999/", + "http://example.com:notaport/", ], ) async def test_http_primitives_block_internal_targets(url, monkeypatch):