fix(proxy): validate JSON object bodies before request handling (#48438)

This commit is contained in:
tin-berri 2026-10-26 19:11:38 -04:00 • committed by GitHub
parent db9eda19b4
commit bc0f598db5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 226 additions and 27 deletions

View file

@ -12,7 +12,7 @@ import litellm
from litellm._logging import verbose_proxy_logger, verbose_proxy_stdout_logger
from litellm.constants import EMPTY_MAPPING
from litellm.integrations.otel.runtime import seed_request_identity
from litellm.litellm_core_utils.core_helpers import is_expected_client_error
from litellm.litellm_core_utils.core_helpers import as_str_mapping, is_expected_client_error
from litellm.proxy._types import (
LitellmUserRoles,
ModelAccessDeniedProxyException,
@ -139,7 +139,7 @@ class UserAPIKeyAuthExceptionHandler:
async def handle_authentication_error(
e: Exception,
request: Request,
request_data: dict[str, object],
request_data: object,
route: str,
parent_otel_span: Span | None,
api_key: str,
@ -190,6 +190,8 @@ class UserAPIKeyAuthExceptionHandler:
request_route=route,
)
else:
request_mapping: Final = as_str_mapping(request_data)
request_context: Final[dict[str, object]] = dict(request_mapping) if request_mapping is not None else {}
# raise the exception to the caller
requester_ip: Final = get_request_ip_address(
request=request,
@ -229,7 +231,7 @@ class UserAPIKeyAuthExceptionHandler:
# without this the failed trace would carry no team/key attributes.
seed_request_identity(
user_api_key_dict,
model=request_data.get("model"),
model=request_context.get("model"),
)
# Budget checks live in tenant-scoped helpers (key / team / org / tag)
@ -242,14 +244,14 @@ class UserAPIKeyAuthExceptionHandler:
resolve_llm_provider_for_rate_limit,
)
budget_model: Final = request_data.get("model")
budget_model: Final = request_context.get("model")
_, e.llm_provider = resolve_llm_provider_for_rate_limit(
budget_model if isinstance(budget_model, str) else None
)
# Allow callbacks to transform the error response
transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook(
request_data=_with_client_context(request_data, requester_ip, user_agent),
request_data=_with_client_context(request_context, requester_ip, user_agent),
original_exception=e,
user_api_key_dict=user_api_key_dict,
error_type=ProxyErrorTypes.auth_error,

View file

@ -6,7 +6,7 @@ from typing import Annotated, Any, Final, Literal, Union, get_args, get_origin
import orjson
from fastapi import Request, UploadFile, status
from pydantic import TypeAdapter
from pydantic import TypeAdapter, ValidationError
from starlette._utils import get_route_path
from starlette.requests import HTTPConnection
from typing_extensions import NotRequired, ReadOnly, Required, assert_never
@ -162,13 +162,11 @@ def coerce_numeric_form_fields(
}
def _parse_binary_body(body: bytes) -> dict:
def _parse_binary_body(body: bytes) -> dict[str, object]:
"""JSON sent under a binary content type still parses; real binary (protobuf) carries no params -> {}."""
try:
parsed: Final = _DECODED_JSON.validate_python(orjson.loads(body))
if isinstance(parsed, dict):
return parsed
except orjson.JSONDecodeError:
return dict(_JSON_REQUEST_BODY.validate_python(orjson.loads(body)))
except (orjson.JSONDecodeError, ValidationError):
pass
return {}
@ -299,8 +297,9 @@ async def read_request_body(request: Request | None) -> dict:
body_str = re.sub(r"(?<![\uD800-\uDBFF])[\uDC00-\uDFFF]", "", body_str)
try:
parsed_body = json.loads(body_str)
json.dumps(parsed_body, ensure_ascii=False).encode("utf-8")
decoded_body: Final = _DECODED_JSON.validate_python(json.loads(body_str))
json.dumps(decoded_body, ensure_ascii=False).encode("utf-8")
parsed_body = dict(_JSON_REQUEST_BODY.validate_python(decoded_body))
except (json.JSONDecodeError, UnicodeEncodeError):
# json.loads accepts lone surrogate escapes that no provider can encode
verbose_proxy_logger.error("Invalid JSON payload received: %s", e)
@ -315,6 +314,34 @@ async def read_request_body(request: Request | None) -> dict:
safe_set_request_parsed_body(request=request, parsed_body=parsed_body)
return parsed_body
except ValidationError as e:
from fastapi.routing import APIRoute
from litellm.proxy.auth.auth_utils import (
request_dispatched_to_pass_through_endpoint,
request_dispatched_to_provider_pass_through,
)
from litellm.proxy.lens.buffering import LensRoute
route: Final[object] = request.scope.get("route") if request is not None else None
if request is not None and (
request_dispatched_to_pass_through_endpoint(request)
or request_dispatched_to_provider_pass_through(request)
or isinstance(route, LensRoute)
or (
isinstance(route, APIRoute)
and route.body_field is not None
and not route.body_field.field_info.is_required()
and (await request.body()).strip() == b"null"
)
):
return {}
raise ProxyException(
message="Invalid JSON payload: expected an object",
type="invalid_request_error",
param="request_body",
code=status.HTTP_400_BAD_REQUEST,
) from e
except (json.JSONDecodeError, orjson.JSONDecodeError, ProxyException) as e:
# Re-raise ProxyException as-is
verbose_proxy_logger.error("Invalid JSON payload received: %s", e)

View file

@ -1,5 +1,7 @@
import asyncio
import json
from collections.abc import Mapping
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
@ -29,6 +31,7 @@ import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import INVALID_VIRTUAL_KEY_ERROR_MARKER
from litellm.exceptions import BudgetExceededError
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import (
ModelAccessDeniedProxyException,
ProxyErrorTypes,
@ -812,6 +815,65 @@ def _http_request(client_host: str | None = "10.1.2.3", headers: dict[str, str]
)
class _AuthFailureRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.data: dict[str, object] | None = None
self.error: Exception | None = None
async def async_post_call_failure_hook(
self,
request_data: dict[str, object],
original_exception: Exception,
user_api_key_dict: UserAPIKeyAuth,
traceback_str: str | None = None,
) -> HTTPException | None:
self.data = dict(request_data)
self.error = original_exception
return None
@pytest.mark.asyncio
@pytest.mark.parametrize("request_data", [[], None, "metadata", 1, True])
@pytest.mark.parametrize("error_kind", ["auth", "budget", "database"])
async def test_auth_error_preserves_classification_with_non_object_context(
monkeypatch: pytest.MonkeyPatch, request_data: object, error_kind: str
) -> None:
import litellm.proxy.proxy_server as proxy_server
monkeypatch.setattr(litellm, "budget_exceeded_status_code", 429)
errors: Final = {
"auth": HTTPException(status_code=401, detail="Invalid API key"),
"budget": BudgetExceededError(current_cost=2, max_budget=1),
"database": httpx.ConnectError("Database unavailable"),
}
error: Final = errors[error_kind]
expected: Final = _as_proxy_exception(error)
recorder: Final = _AuthFailureRecorder()
monkeypatch.setattr(litellm, "callbacks", [recorder])
monkeypatch.setattr(proxy_server, "general_settings", {"allow_requests_on_db_unavailable": False})
with pytest.raises(ProxyException) as exc_info:
await UserAPIKeyAuthExceptionHandler.handle_authentication_error(
error,
_http_request(headers={"user-agent": "shape-regression"}),
request_data,
"/key/info",
None,
"sk-test",
)
assert exc_info.value.code == expected.code
assert exc_info.value.type == expected.type
assert exc_info.value.message == expected.message
assert recorder.error is error
assert recorder.data is not None
metadata: Final = recorder.data["metadata"]
assert isinstance(metadata, Mapping)
assert metadata["requester_ip_address"] == "10.1.2.3"
assert metadata["user_agent"] == "shape-regression"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"auth_error, general_settings, request_kwargs, expected_ip",

View file

@ -7724,7 +7724,8 @@ async def test_user_api_key_auth_does_not_overwrite_end_user_id_set_by_builder()
@pytest.mark.asyncio
async def test_user_api_key_auth_authenticates_before_raising_malformed_body_error():
@pytest.mark.parametrize("body", [b'{}{"model": "gpt-4o"}', b"[]", b"null", b"[1e400]"])
async def test_user_api_key_auth_authenticates_before_raising_malformed_body_error(body: bytes) -> None:
"""Regression (LIT-4780): a body that fails to parse must still be authenticated
first, so the rejected request's trace carries the caller's key / team / user
identity instead of an anonymous root span. The parse error is re-raised
@ -7742,10 +7743,11 @@ async def test_user_api_key_auth_authenticates_before_raising_malformed_body_err
"headers": [(b"content-type", b"application/json")],
"method": "POST",
"path": "/chat/completions",
"query_string": b"probe=1",
}
)
request._url = URL(url="/chat/completions")
request._body = b'{}{"model": "gpt-4o"}'
request._body = body
attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None)
originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}

View file

@ -64,6 +64,100 @@ def _starlette_request(
return Request(scope, receive)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"body",
[b"[]", b'[["model", "test"]]', b"null", b'"s"', b"1", b"true", b"false"],
)
@pytest.mark.parametrize("path", ["/v1/messages", "/vertex_ai/v1/messages", "/lens/datasets"])
async def test_read_request_body_rejects_non_object_json(body: bytes, path: str) -> None:
request: Final = _starlette_request(body, "application/json", path)
for _ in range(2):
with pytest.raises(ProxyException) as exc_info:
await read_request_body(request)
assert exc_info.value.code == "400"
assert exc_info.value.type == "invalid_request_error"
assert exc_info.value.param == "request_body"
assert exc_info.value.message == "Invalid JSON payload: expected an object"
assert _safe_get_request_parsed_body(request) is None
assert await request.body() == body
@pytest.mark.asyncio
@pytest.mark.parametrize("body", [b"[1e400]", b"1e400", b"[NaN]"])
async def test_read_request_body_rejects_non_objects_from_json_fallback(body: bytes) -> None:
with pytest.raises(ProxyException) as exc_info:
await read_request_body(_starlette_request(body, "application/json"))
assert exc_info.value.code == "400"
assert exc_info.value.message == "Invalid JSON payload: expected an object"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"body, invalid_unicode",
[
(b"[]", False),
(b"null", False),
(b'"value"', False),
(b"[1e400]", False),
(b'"\\ud800"', True),
(b'["\\ud800"]', True),
(b'{"value":"\\ud800"}', True),
],
)
@pytest.mark.parametrize("dispatch", ["custom", "provider", "lens"])
async def test_pass_through_body_validation_preserves_bytes(dispatch: str, body: bytes, invalid_unicode: bool) -> None:
from fastapi.routing import APIRoute
from litellm.proxy.lens.adapter import lens_request
from litellm.proxy.lens.buffering import LensRoute
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import vertex_proxy_route
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import create_pass_through_route
path: Final = {"custom": "/v1/messages", "provider": "/vertex_ai/predict", "lens": "/lens/datasets"}[dispatch]
route: Final = {
"custom": APIRoute(path, create_pass_through_route(path, "https://upstream.test"), methods=["POST"]),
"provider": APIRoute("/vertex_ai/{endpoint:path}", vertex_proxy_route, methods=["POST"]),
"lens": LensRoute("/lens/{path:path}", lens_request, methods=["POST"]),
}[dispatch]
incoming: Final = _starlette_request(body, "application/json", path)
_, dispatched_scope = route.matches(incoming.scope)
request: Final = Request({**incoming.scope, **dispatched_scope}, incoming.receive)
for _ in range(2):
if invalid_unicode:
with pytest.raises(ProxyException) as exc_info:
await read_request_body(request)
assert exc_info.value.code == "400"
else:
assert await read_request_body(request) == {}
assert await request.body() == body
@pytest.mark.asyncio
@pytest.mark.parametrize("optional", [False, True])
@pytest.mark.parametrize("body", [b"null", b" \nnull\t", b"[]", b'"value"'])
async def test_null_body_respects_the_dispatched_body_requirement(optional: bool, body: bytes) -> None:
from fastapi.routing import APIRoute
from litellm.proxy.management_endpoints.key_management_endpoints import generate_key_fn
from litellm.proxy.spend_tracking.spend_management_endpoints import global_spend_end_users
endpoint: Final = global_spend_end_users if optional else generate_key_fn
route: Final = APIRoute("/audit", endpoint, methods=["POST"])
incoming: Final = _starlette_request(body, "application/json", "/audit")
_, dispatched_scope = route.matches(incoming.scope)
request: Final = Request({**incoming.scope, **dispatched_scope}, incoming.receive)
if optional and body.strip() == b"null":
assert await read_request_body(request) == {}
assert await request.json() is None
else:
with pytest.raises(ProxyException) as exc_info:
await read_request_body(request)
assert exc_info.value.code == "400"
assert await request.body() == body
@pytest.mark.asyncio
async def test_read_request_body_marks_body_received_once_with_its_size(monkeypatch: pytest.MonkeyPatch):
events: list[tuple[str, dict[str, str | int]]] = [] # mutable-ok: recorder for the injected phase_event double
@ -131,14 +225,20 @@ async def test_read_raw_json_body_is_none_for_form_bodies():
@pytest.mark.asyncio
@pytest.mark.parametrize("content_type", ["application/x-protobuf", "application/protobuf; charset=binary"])
async def test_protobuf_body_is_not_parsed_as_json(content_type):
# OTLP trace exports (POST /v1/traces) are binary protobuf; arbitrary bytes like these
# used to hit the JSON surrogate-repair path and fail auth with a 400.
body = b"\n\xa2\x01\n\x1c\n\x0cservice.name\x12\x0c\n\nswarm\xed\xa0\x80\xff"
request = _starlette_request(body, content_type)
@pytest.mark.parametrize(
"body",
[
b"\n\xa2\x01\n\x1c\n\x0cservice.name\x12\x0c\n\nswarm\xed\xa0\x80\xff",
pytest.param(b"100000000", id="protobuf-also-json-number"),
],
)
async def test_protobuf_body_is_not_parsed_as_json(content_type: str, body: bytes) -> None:
request: Final = _starlette_request(body, content_type)
assert await read_request_body(request) == {}
assert await request.body() == body # body is still readable by the endpoint
assert await read_request_body(request) == {}
assert await read_raw_json_body(request) == body
assert await request.body() == body
@pytest.mark.asyncio

View file

@ -249,7 +249,10 @@ async def test_admission_releases_upload_on_disconnect_or_cancellation(cancel: b
@pytest.mark.asyncio
@pytest.mark.parametrize("reject_auth", [False, True])
async def test_admitted_bytes_survive_auth_parsing_and_remain_reserved_until_sent(reject_auth: bool) -> None:
@pytest.mark.parametrize("body", [b" {} ", b"[]", b"null", b'"s"'])
async def test_admitted_bytes_survive_auth_parsing_and_remain_reserved_until_sent(
reject_auth: bool, body: bytes
) -> None:
app: Final = FastAPI()
router: Final = APIRouter(route_class=LensRoute)
@ -264,9 +267,11 @@ async def test_admitted_bytes_survive_auth_parsing_and_remain_reserved_until_sen
if reject_auth:
raise HTTPException(401, "Invalid key")
async with httpx.AsyncClient(
transport=httpx.MockTransport(lambda _: httpx.Response(200, content=b"result"))
) as upstream:
def upstream_response(incoming: httpx.Request) -> httpx.Response:
assert incoming.content == body
return httpx.Response(200, content=b"result")
async with httpx.AsyncClient(transport=httpx.MockTransport(upstream_response)) as upstream:
@router.post("/lens/datasets")
async def endpoint(request: Request, auth: Annotated[None, Depends(auth_reader)]) -> Response:
@ -284,14 +289,14 @@ async def test_admitted_bytes_survive_auth_parsing_and_remain_reserved_until_sen
sending.set()
await finish.wait()
incoming: Final = request(b" {} ")
incoming: Final = request(body)
task: Final = asyncio.create_task(app(incoming.scope, incoming.receive, send))
try:
await asyncio.wait_for(sending.wait(), 2)
async with httpx.AsyncClient(
transport=httpx.ASGITransport(app=app), base_url="http://gateway.test"
) as client:
blocked: Final = await client.post("/lens/datasets", content=b" {} " * 2)
blocked: Final = await client.post("/lens/datasets", content=b" " * (11 - len(body)))
unrelated: Final = await client.post("/ordinary", content=b"ordinary body")
assert blocked.status_code == 503
assert unrelated.content == b"ordinary body"

View file

@ -4357,6 +4357,7 @@ class TestVertexRawPredictStreamingClassification:
"type": "http",
"method": "POST",
"path": f"/vertex_ai/{endpoint}",
"path_params": {"endpoint": endpoint},
"headers": [
(b"content-type", b"application/json"),
(b"x-litellm-api-key", b"test-key"),