diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index f6637a1fdf3..2eb25cc6105 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -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, diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index a38769ac41b..9f7fed40ab3 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -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"(? 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) diff --git a/tests/unit/proxy/auth/test_auth_exception_handler.py b/tests/unit/proxy/auth/test_auth_exception_handler.py index d83a8aa2fcd..e21491fd996 100644 --- a/tests/unit/proxy/auth/test_auth_exception_handler.py +++ b/tests/unit/proxy/auth/test_auth_exception_handler.py @@ -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", diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 108f7ce2983..2d9c781b0d1 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -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} diff --git a/tests/unit/proxy/common_utils/test_http_parsing_utils.py b/tests/unit/proxy/common_utils/test_http_parsing_utils.py index cd5820fa2fd..354feb0bab1 100644 --- a/tests/unit/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/unit/proxy/common_utils/test_http_parsing_utils.py @@ -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 diff --git a/tests/unit/proxy/lens/test_buffering.py b/tests/unit/proxy/lens/test_buffering.py index b544b70d36b..2b868e72096 100644 --- a/tests/unit/proxy/lens/test_buffering.py +++ b/tests/unit/proxy/lens/test_buffering.py @@ -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" diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 0186a33dbc4..60043daa402 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -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"),