scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

Merge 8a54ce569c into c2c2a623c0

Browse source
This commit is contained in:
Nguyễn Anh Bình 2026-09-13 18:04:42 +00:00 • committed by GitHub
commit c5fc749951
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 274 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
litellm/llms/custom_httpx/http_handler.py
View file

@ -750,7 +750,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:
@ -770,7 +772,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):
@ -785,6 +787,7 @@ class AsyncHTTPHandler:
params=params,
headers=headers,
stream=stream,
follow_redirects=follow_redirects,
)
finally:
await new_client.aclose()
@ -892,7 +895,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
@ -910,7 +915,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):
@ -925,6 +930,7 @@ class AsyncHTTPHandler:
params=params,
headers=headers,
stream=stream,
follow_redirects=follow_redirects,
)
finally:
await new_client.aclose()
@ -956,7 +962,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
@ -974,7 +982,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):
@ -989,6 +997,7 @@ class AsyncHTTPHandler:
params=params,
headers=headers,
stream=stream,
follow_redirects=follow_redirects,
)
finally:
await new_client.aclose()
@ -1007,12 +1016,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)
@ -1025,7 +1036,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

70
litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py
View file

@ -15,7 +15,9 @@ import httpx
from pydantic import JsonValue
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
from litellm.types.llms.custom_http import httpxSpecialProvider
@ -463,6 +465,14 @@ async def http_request(
Uses LiteLLM's global cached AsyncHTTPHandler for connection pooling
and better performance.
SSRF protection: the URL is resolved and validated (via
``litellm_core_utils.url_utils.validate_url``) before connecting, and
redirects are never followed. Requests targeting loopback, private,
link-local, or cloud-metadata addresses are blocked. Operators can
allow specific internal hosts via ``user_url_allowed_hosts`` in
``general_settings``, or disable validation with
``litellm.user_url_validation = False``.
Args:
url: The URL to request
method: HTTP method (GET, POST, PUT, DELETE, PATCH). Defaults to GET.
@ -496,6 +506,20 @@ async def http_request(
if not is_valid_url(url):
return _http_error_response(f"Invalid URL: {url}")
# SSRF protection: block requests that resolve to internal/private/metadata
# targets before any connection is made. Legitimate internal services can
# be reached by adding them to `user_url_allowed_hosts` in general_settings.
# `litellm.user_url_validation = False` disables this check entirely.
validated_url = url
request_headers = headers
if getattr(litellm, "user_url_validation", True):
try:
validated_url, host_header = validate_url(url)
request_headers = {**(headers or {}), "Host": host_header}
except (SSRFError, ValueError) as e:
verbose_proxy_logger.warning("Custom code http_request SSRF blocked: %s", e)
return _http_error_response(f"Blocked: {e}")
# Validate and normalize method
method = method.upper()
allowed_methods: Final = {"GET", "POST", "PUT", "DELETE", "PATCH"}
@ -515,7 +539,7 @@ async def http_request(
)
try:
response: Final = await _execute_http_request(client, method, url, headers, body, timeout)
response: Final = await _execute_http_request(client, method, validated_url, request_headers, body, timeout)
return _http_success_response(response)
except httpx.TimeoutException as e:
@ -540,19 +564,51 @@ async def _execute_http_request(
body: JsonValue,
timeout: float,
) -> httpx.Response:
"""Execute the HTTP request using the appropriate client method."""
"""Execute the HTTP request using the appropriate client method.
Redirects are disabled on every method so a 3xx response cannot bypass
the SSRF validation performed by the caller.
"""
json_body, data_body = _prepare_http_body(body)
if method == "GET":
return await client.get(url=url, headers=headers)
return await client.get(url=url, headers=headers, follow_redirects=False)
elif method == "POST":
return await client.post(url=url, headers=headers, json=json_body, data=data_body, timeout=timeout)
return await client.post(
url=url,
headers=headers,
json=json_body,
data=data_body,
timeout=timeout,
follow_redirects=False,
)
elif method == "PUT":
return await client.put(url=url, headers=headers, json=json_body, data=data_body, timeout=timeout)
return await client.put(
url=url,
headers=headers,
json=json_body,
data=data_body,
timeout=timeout,
follow_redirects=False,
)
elif method == "DELETE":
return await client.delete(url=url, headers=headers, json=json_body, data=data_body, timeout=timeout)
return await client.delete(
url=url,
headers=headers,
json=json_body,
data=data_body,
timeout=timeout,
follow_redirects=False,
)
elif method == "PATCH":
return await client.patch(url=url, headers=headers, json=json_body, data=data_body, timeout=timeout)
return await client.patch(
url=url,
headers=headers,
json=json_body,
data=data_body,
timeout=timeout,
follow_redirects=False,
)
else:
raise ValueError(f"Unsupported HTTP method: {method}")

39
tests/test_litellm/llms/custom_httpx/test_http_handler.py
View file

@ -6,11 +6,13 @@ import pathlib
import ssl
import threading
import weakref
from typing import Final
from unittest.mock import MagicMock, patch
import certifi
import httpx
import pytest
import respx
from aiohttp import ClientSession, TCPConnector
import litellm
@ -1550,7 +1552,8 @@ def test_sync_force_ipv4_https_proxy_mount_uses_handler_ca_bundle(
@pytest.mark.asyncio
async def test_put_can_refuse_to_follow_a_redirect():
@pytest.mark.parametrize("method", ["post", "put", "patch", "delete"])
async def test_non_get_methods_can_refuse_to_follow_a_redirect(method: str) -> None:
"""The client follows redirects by default; a caller uploading to a URL it did not choose must be able to opt out."""
hops: list[str] = [] # mutable-ok: the fake transport records the paths it was asked for
@ -1564,19 +1567,49 @@ async def test_put_can_refuse_to_follow_a_redirect():
await handler.client.aclose()
handler.client = httpx.AsyncClient(transport=httpx.MockTransport(mock_handler), follow_redirects=True)
try:
followed = await handler.put("https://uploads.example/first", data=b"x")
followed = await getattr(handler, method)("https://uploads.example/first", data=b"x")
assert followed.status_code == 200
assert hops == ["/first", "/second"]
hops.clear()
with pytest.raises(MaskedHTTPStatusError) as refused:
await handler.put("https://uploads.example/first", data=b"x", follow_redirects=False)
await getattr(handler, method)("https://uploads.example/first", data=b"x", follow_redirects=False)
assert refused.value.status_code == 302
assert hops == ["/first"]
finally:
await handler.close()
@pytest.mark.asyncio
@pytest.mark.parametrize("error_type", [httpx.ConnectError, httpx.RemoteProtocolError])
@pytest.mark.parametrize("follow_redirects", [None, False])
async def test_post_retry_preserves_redirect_preference(
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
error_type: type[httpx.ConnectError] | type[httpx.RemoteProtocolError],
follow_redirects: bool | None,
) -> None:
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
first: Final = respx_mock.post("https://moderation.example/check").mock(
side_effect=[error_type("connection reset"), httpx.Response(302, headers={"Location": "/landing"})]
)
landing: Final = respx_mock.get("https://moderation.example/landing").respond(200)
handler: Final = AsyncHTTPHandler()
try:
if follow_redirects is False:
with pytest.raises(httpx.HTTPStatusError, match="302"):
await handler.post("https://moderation.example/check", json={"text": "hi"}, follow_redirects=False)
assert landing.call_count == 0
else:
response: Final = await handler.post("https://moderation.example/check", json={"text": "hi"})
assert response.status_code == 200
assert landing.call_count == 1
assert first.call_count == 2
assert first.calls[0].request.content == first.calls[1].request.content == b'{"text":"hi"}'
finally:
await handler.close()
@pytest.mark.asyncio
async def test_a_retried_put_stays_a_put_and_still_refuses_redirects():
"""

186
tests/test_litellm/proxy/guardrails/test_custom_code_security.py
View file

@ -1,7 +1,16 @@
import socket
from collections.abc import AsyncIterator
from typing import Final
import httpx
import pytest
import pytest_asyncio
from fastapi import HTTPException
import litellm
from litellm.exceptions import ModifyResponseException
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy.guardrails.guardrail_hooks.custom_code import primitives
from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import (
CustomCodeCompilationError,
CustomCodeGuardrail,
@ -77,18 +86,14 @@ def test_nfkc_homoglyph_rejected_at_compile():
[
# Literal dunder attribute access.
"def apply_guardrail(i, r, t):\n return str.__class__\n",
"def apply_guardrail(i, r, t):\n"
" return ().__class__.__bases__[0].__subclasses__()\n",
"def apply_guardrail(i, r, t):\n return ().__class__.__bases__[0].__subclasses__()\n",
# gi_code — on the transformer's restricted-names list.
"def apply_guardrail(i, r, t):\n"
" def g():\n yield 1\n"
" return g().gi_code\n",
"def apply_guardrail(i, r, t):\n def g():\n yield 1\n return g().gi_code\n",
# Import forms.
"import os\ndef apply_guardrail(i, r, t):\n return allow()\n",
"from subprocess import call\n"
"def apply_guardrail(i, r, t):\n return allow()\n",
"from subprocess import call\ndef apply_guardrail(i, r, t):\n return allow()\n",
# __import__ is rejected as an underscore-prefixed name.
"def apply_guardrail(i, r, t):\n" ' return __import__("os")\n',
'def apply_guardrail(i, r, t):\n return __import__("os")\n',
],
)
def test_compile_time_rejections(snippet: str):
@ -100,8 +105,7 @@ def test_compile_time_rejections(snippet: str):
"snippet",
[
# getattr is not in the sandbox builtins — NameError at call time.
"def apply_guardrail(i, r, t):\n"
' return getattr(str, "_"+"_class_"+"_")\n',
'def apply_guardrail(i, r, t):\n return getattr(str, "_"+"_class_"+"_")\n',
# setattr is guarded_setattr + full_write_guard — setting any attribute
# on a user-defined object raises TypeError, whether the name is a
# dunder or not.
@ -139,10 +143,7 @@ def test_documented_ssn_example_compiles_and_runs():
@pytest.mark.asyncio
async def test_async_guardrail_compiles_and_runs():
code = (
"async def apply_guardrail(inputs, request_data, input_type):\n"
" return allow()\n"
)
code = "async def apply_guardrail(inputs, request_data, input_type):\n return allow()\n"
guardrail = _compile(code)
from litellm.types.utils import GenericGuardrailAPIInputs
@ -156,10 +157,7 @@ async def test_async_guardrail_compiles_and_runs():
@pytest.mark.asyncio
async def test_custom_code_pre_call_block_uses_passthrough():
code = (
"def apply_guardrail(inputs, request_data, input_type):\n"
' return block("blocked by test")\n'
)
code = 'def apply_guardrail(inputs, request_data, input_type):\n return block("blocked by test")\n'
guardrail = _compile(code)
with pytest.raises(ModifyResponseException) as exc_info:
@ -176,10 +174,7 @@ async def test_custom_code_pre_call_block_uses_passthrough():
@pytest.mark.asyncio
async def test_custom_code_post_call_block_raises_http_400():
code = (
"def apply_guardrail(inputs, request_data, input_type):\n"
' return block("blocked by test")\n'
)
code = 'def apply_guardrail(inputs, request_data, input_type):\n return block("blocked by test")\n'
guardrail = _compile(code)
with pytest.raises(HTTPException) as exc_info:
@ -263,10 +258,7 @@ async def test_custom_code_allow_still_records_success_not_flagged():
def test_typical_sync_guardrail_still_works():
code = (
"def apply_guardrail(inputs, request_data, input_type):\n"
" return allow()\n"
)
code = "def apply_guardrail(inputs, request_data, input_type):\n return allow()\n"
guardrail = _compile(code)
assert guardrail._compiled_function is not None
@ -293,3 +285,145 @@ def test_augmented_assignment_works():
def test_missing_apply_guardrail_raises():
with pytest.raises(CustomCodeCompilationError, match="apply_guardrail"):
_compile("x = 1\n")
@pytest_asyncio.fixture
async def _http_requests(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[list[httpx.Request]]:
requests: Final[list[httpx.Request]] = []
def respond(request: httpx.Request) -> httpx.Response:
requests.append(request)
if request.url.path == "/redirect":
return httpx.Response(302, headers={"Location": "http://169.254.169.254/latest/meta-data/"})
return httpx.Response(200, json={"allowed": True})
handler: Final = AsyncHTTPHandler()
await handler.client.aclose()
async with httpx.AsyncClient(transport=httpx.MockTransport(respond), follow_redirects=True) as client:
handler.client = client
monkeypatch.setattr(primitives, "get_async_httpx_client", lambda **kwargs: handler)
monkeypatch.setattr(litellm, "user_url_validation", True)
monkeypatch.setattr(litellm, "user_url_allowed_hosts", [])
yield requests
def _resolve_to(monkeypatch: pytest.MonkeyPatch, address: str) -> None:
def getaddrinfo(host: str, port: int, *, proto: int) -> list[tuple[int, int, int, str, tuple[str, int]]]:
return [(socket.AF_INET6 if ":" in address else socket.AF_INET, socket.SOCK_STREAM, proto, "", (address, port))]
monkeypatch.setattr("litellm.litellm_core_utils.url_utils.socket.getaddrinfo", getaddrinfo)
@pytest.mark.asyncio
@pytest.mark.parametrize("method", ["GET", "POST", "PUT", "DELETE", "PATCH"])
@pytest.mark.parametrize(
("url", "resolved_ip"),
[
("http://127.0.0.1:8971/", "127.0.0.1"),
("http://localhost/admin", "127.0.0.1"),
("http://10.1.2.3/", "10.1.2.3"),
("http://192.168.1.1/", "192.168.1.1"),
("http://169.254.169.254/latest/meta-data/iam/security-credentials/", "169.254.169.254"),
("http://[::1]/", "::1"),
("http://example.com:99999/", "93.184.216.34"),
("http://example.com:notaport/", "93.184.216.34"),
],
)
async def test_http_primitives_block_internal_targets(
url: str,
resolved_ip: str,
method: str,
monkeypatch: pytest.MonkeyPatch,
_http_requests: list[httpx.Request],
) -> None:
_resolve_to(monkeypatch, resolved_ip)
result: Final = await primitives.http_request(url, method=method)
assert result["success"] is False
assert result["status_code"] == 0
assert result["error"] is not None and "Blocked" in result["error"]
assert _http_requests == []
@pytest.mark.asyncio
@pytest.mark.parametrize("method", ["GET", "POST", "PUT", "DELETE", "PATCH"])
async def test_http_primitives_allow_public_url(
method: str, monkeypatch: pytest.MonkeyPatch, _http_requests: list[httpx.Request]
) -> None:
_resolve_to(monkeypatch, "93.184.216.34")
headers: Final = {"Authorization": "Bearer t"}
result: Final = await primitives.http_request(
"http://moderation.example.com:8080/v1/check", method=method, headers=headers, body={"text": "hi"}
)
assert result["success"] is True
assert result["body"] == {"allowed": True}
assert len(_http_requests) == 1
request: Final = _http_requests[0]
assert str(request.url) == "http://93.184.216.34:8080/v1/check"
assert request.method == method
assert request.headers["Host"] == "moderation.example.com:8080"
assert request.headers["Authorization"] == "Bearer t"
assert headers == {"Authorization": "Bearer t"}
if method != "GET":
assert request.content == b'{"text":"hi"}'
@pytest.mark.asyncio
async def test_http_primitives_honor_allowlisted_internal_host(
monkeypatch: pytest.MonkeyPatch, _http_requests: list[httpx.Request]
) -> None:
_resolve_to(monkeypatch, "10.1.2.3")
blocked: Final = await primitives.http_get("http://internal-moderation.corp/check")
assert blocked["success"] is False
assert _http_requests == []
monkeypatch.setattr(litellm, "user_url_allowed_hosts", ["internal-moderation.corp"])
result: Final = await primitives.http_get("http://internal-moderation.corp/check")
assert result["success"] is True
assert len(_http_requests) == 1
assert str(_http_requests[0].url) == "http://10.1.2.3/check"
assert _http_requests[0].headers["Host"] == "internal-moderation.corp"
@pytest.mark.asyncio
async def test_http_primitives_validation_can_be_disabled(
monkeypatch: pytest.MonkeyPatch, _http_requests: list[httpx.Request]
) -> None:
monkeypatch.setattr(litellm, "user_url_validation", False)
def unexpected_resolution(host: str, port: int, *, proto: int) -> None:
pytest.fail("Disabled validation must not resolve DNS")
monkeypatch.setattr("litellm.litellm_core_utils.url_utils.socket.getaddrinfo", unexpected_resolution)
result: Final = await primitives.http_get("http://10.1.2.3/internal")
assert result["success"] is True
assert len(_http_requests) == 1
assert str(_http_requests[0].url) == "http://10.1.2.3/internal"
@pytest.mark.asyncio
@pytest.mark.parametrize("method", ["GET", "POST", "PUT", "DELETE", "PATCH"])
@pytest.mark.parametrize("validation_enabled", [True, False])
async def test_http_primitives_do_not_follow_redirects(
method: str,
validation_enabled: bool,
monkeypatch: pytest.MonkeyPatch,
_http_requests: list[httpx.Request],
) -> None:
_resolve_to(monkeypatch, "93.184.216.34")
monkeypatch.setattr(litellm, "user_url_validation", validation_enabled)
result: Final = await primitives.http_request(
"http://moderation.example.com/redirect", method=method, body={"text": "hi"}
)
assert len(_http_requests) == 1
assert _http_requests[0].url.path == "/redirect"
assert _http_requests[0].method == method
assert result["status_code"] == 302
assert result["success"] is False