From 98dbe80ada9d2530e9764a7b8267d376a5c17290 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Thu, 12 Nov 2026 14:05:05 -0800 Subject: [PATCH] fix(proxy): refuse blocked and over-limit keys before reading the request body (#51428) * fix(proxy): refuse blocked and over-limit keys before reading the request body A client retrying 429s with multi-megabyte payloads made every worker buffer each payload in full before the refusal went out, so the rejections themselves drove memory. Auth now checks the cached key for a block or an exhausted key-level rate limit from the headers alone and answers before the body is read; admitted requests and cache misses take the unchanged path. * fix(proxy): pick the effective key header for the early refusal and keep its failure log The headers-only check now reads x-litellm-api-key when present, skips a missing Authorization header instead of raising, stamps the route on the cached identity so the failure is logged, and honors an exception a callback returns from the failure hook. * fix(proxy): honor litellm_key_header_name in the early refusal The headers-only check now reads the key from the header general_settings.litellm_key_header_name names when the request carries it, the same precedence the full auth path applies. * fix(proxy): keep the early-rejection path inside the type gate ceilings Read the configured key header through general_settings_view() instead of importing the untyped general_settings global, and drop the isinstance on an exception post_call_failure_hook already types as HTTPException. --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 83 ++++++++++++++++++- .../hooks/parallel_request_limiter_v3.py | 20 +++++ .../authorization/test_reject_before_body.py | 75 +++++++++++++++++ .../test_user_api_key_auth_request_flow.py | 79 ++++++++++++++++++ .../hooks/test_parallel_request_limiter_v3.py | 29 +++++++ 5 files changed, 285 insertions(+), 1 deletion(-) create mode 100644 tests/integration/authorization/test_reject_before_body.py diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 02d0e8736e3..fd361e4a539 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2232,6 +2232,84 @@ async def _read_request_body_deferring_parse_failure( return populate_request_with_path_params(request_data=parsed_body, request=request), None +async def _reject_from_headers_before_reading_body( + request: Request, api_key: str | None, custom_litellm_key_header: str | None, route: str +) -> None: + """Refuse a request its headers already condemn, a blocked key or an exhausted key limit, before reading its body. + + Only a key already in the cache is considered, so a miss changes nothing and the full auth path still runs. + A retry flood of large payloads against a limited key used to buffer every payload just to answer 429. + """ + if not RouteChecks.is_llm_api_route(route=route): + return + from litellm.proxy import proxy_server + from litellm.proxy.common_request_processing import proxy_exception_from_http_exception + from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + configured_header: Final = proxy_server.general_settings_view().get("litellm_key_header_name") + bearer_token: Final = ( + get_api_key_from_custom_header(request=request, custom_litellm_key_header_name=configured_header) + if isinstance(configured_header, str) and configured_header in request.headers + else _get_bearer_token_or_received_api_key(custom_litellm_key_header) + if isinstance(custom_litellm_key_header, str) + else get_bearer_token(api_key=api_key) + if isinstance(api_key, str) + else "" + ) + if not bearer_token: + return + hashed_token: Final = hash_token(bearer_token) + try: + cached_key: Final = IdentityStore.key_from_principal( + await IdentityStore(prisma_client, user_api_key_cache, check_cache_only=True).resolve( + hashed_token=hashed_token + ) + ) + except Exception: # noqa: BLE001 # any cache miss defers to the full auth path below + return + parent_otel_span: Final[Span | None] = getattr(request.state, "parent_otel_span", None) + if cached_key.blocked is True: + await UserAPIKeyAuthExceptionHandler.handle_authentication_error( + e=ProxyException( + message="Authentication Error, Key is blocked. Update via `/key/unblock` if you're an admin.", + type=ProxyErrorTypes.auth_error, + param=None, + code=status.HTTP_401_UNAUTHORIZED, + ), + request=request, + request_data={}, + route=route, + parent_otel_span=parent_otel_span, + api_key=bearer_token, + resolved_identity=cached_key, + ) + return + limiter: Final = proxy_logging_obj.get_proxy_hook("parallel_request_limiter") + if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3): + return + rate_limited: Final = cached_key.model_copy( + update={ + "api_key": cached_key.api_key or hashed_token, + "parent_otel_span": parent_otel_span, + "request_route": normalize_request_route(route), + } + ) + try: + await limiter.raise_if_over_limit_before_body(rate_limited) + except HTTPException as rejection: + transformed: Final = await proxy_logging_obj.post_call_failure_hook( # pyright: ignore[reportUnknownMemberType] # bare dict in sig + request_data={}, + original_exception=rejection, + user_api_key_dict=rate_limited, + route=route, + ) + failure: Final = transformed if transformed is not None else rejection + if isinstance(failure, ProxyException): + raise failure + raise proxy_exception_from_http_exception(failure, dict(failure.headers or {})) + + async def _record_unparsable_body_failure( user_api_key_dict: UserAPIKeyAuth, body_parse_exception: ProxyException, @@ -4592,9 +4670,12 @@ async def user_api_key_auth( # close, and the trace never reaches the backend. _ensure_parent_otel_span_on_request_state(request) + route: Final[str] = get_request_route(request=request) + await _reject_from_headers_before_reading_body( + request=request, api_key=api_key, custom_litellm_key_header=custom_litellm_key_header, route=route + ) request_data, body_parse_exception = await _read_request_body_deferring_parse_failure(request=request) phase_event("litellm.request.body_parsed") - route: Final[str] = get_request_route(request=request) ## CHECK IF ROUTE IS ALLOWED # Run the whole auth phase inside a live ``auth`` span so the DB lookups it diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 39409002058..602f955d8a3 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -3950,6 +3950,26 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): *await self._create_tag_rate_limit_descriptors(data), ] + async def raise_if_over_limit_before_body(self, user_api_key_dict: UserAPIKeyAuth) -> None: + """Refuse a request from its headers alone once a limit that needs no body is exhausted. + + Counters are only read, never incremented, so an admitted request is still charged by + ``async_pre_call_hook`` after its body is parsed. Limits that need the model are not checked here. + """ + descriptors: Final = await self._build_request_rate_limit_descriptors( + user_api_key_dict=user_api_key_dict, data={}, call_type=None + ) + if not descriptors: + return + response: Final = await self.should_rate_limit( + descriptors=descriptors, + parent_otel_span=user_api_key_dict.parent_otel_span, + read_only=True, + skip_tpm_check=self.tpm_reservation_enabled, + ) + if response["overall_code"] == "OVER_LIMIT": + self._handle_rate_limit_error(response=response, descriptors=descriptors) + async def _release_request_capacity_when_admitted( self, admission: asyncio.Task[RateLimitResponse], diff --git a/tests/integration/authorization/test_reject_before_body.py b/tests/integration/authorization/test_reject_before_body.py new file mode 100644 index 00000000000..e3e63095f96 --- /dev/null +++ b/tests/integration/authorization/test_reject_before_body.py @@ -0,0 +1,75 @@ +"""A request the proxy can already refuse from its headers gets its refusal before the body arrives. + +A client retrying 429s with multi-megabyte image payloads used to make every worker buffer each payload +in full before answering, so the refusals themselves consumed the memory. The rejection must not depend +on the body, so these tests send only the headers and the first kilobyte of a declared 8 MB body, then +read the answer through the same socket. +""" + +import socket +import uuid +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway + +pytestmark: Final = pytest.mark.shares_worker_state + +_DECLARED_BODY_BYTES: Final = 8 * 1024 * 1024 +_SENT_BODY_BYTES: Final = 1024 + + +def _small_chat(gateway: Gateway, model: str, key: str) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"reject before body {uuid.uuid4().hex}"}]}, + key=key, + ) + + +def _status_before_body_completes(gateway: Gateway, model: str, key: str) -> int: + """POST a chat completion, send 1 KB of its declared 8 MB body, and return the status the proxy answers with. + + The socket stays open with the body unfinished, so an answer here can only have been made from the + headers. A proxy that waits for the whole body never answers and the read times out. + """ + url: Final = gateway.client.base_url + head: Final = ( + f"POST /v1/chat/completions HTTP/1.1\r\nHost: {url.host}\r\nAuthorization: Bearer {key}\r\n" + f"Content-Type: application/json\r\nContent-Length: {_DECLARED_BODY_BYTES}\r\n\r\n" + ).encode() + partial_body: Final = (b'{"model": "' + model.encode() + b'", "messages": [{"role": "user", "content": "').ljust( + _SENT_BODY_BYTES, b"A" + ) + with socket.create_connection((url.host, url.port or 80), timeout=gateway.client.timeout.read) as connection: + connection.sendall(head + partial_body) + status_line: Final = connection.recv(4096).split(b"\r\n", 1)[0] + return int(status_line.split()[1]) + + +def test_rate_limited_key_is_refused_before_its_body_arrives(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model() + key: Final = scenario.key(models=[model], rpm_limit=1, max_parallel_requests=1) + admitted: Final = _small_chat(gateway, model, key) + assert admitted.status_code == 200, admitted.text + refused: Final = _small_chat(gateway, model, key) + assert refused.status_code == 429, refused.text + + assert _status_before_body_completes(gateway, model, key) == 429 + + +def test_blocked_key_is_refused_before_its_body_arrives(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model() + key: Final = scenario.key(models=[model]) + admitted: Final = _small_chat(gateway, model, key) + assert admitted.status_code == 200, admitted.text + blocked: Final = gateway.request("POST", "/key/block", {"key": key}) + assert blocked.status_code == 200, blocked.text + refused: Final = _small_chat(gateway, model, key) + assert refused.status_code == 401, refused.text + + assert _status_before_body_completes(gateway, model, key) == 401 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 af7ee104351..d0acb13230b 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 @@ -70,6 +70,8 @@ from litellm.proxy._types import ( JWTRoutingOverride, ) from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler +from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 +from litellm.proxy.utils import InternalUsageCache from litellm.proxy.auth.auth_checks import ( OrganizationNotFoundError, TeamNotFoundError, @@ -14270,3 +14272,80 @@ async def test_should_skip_budget_checks_spares_only_zero_spend_mcp_discovery( ) is expected ) + + +def _request_whose_body_must_not_be_read(api_key: str) -> Request: + async def receive() -> dict[str, object]: + raise AssertionError("the body was read before the request was refused") + + return Request( + { + "type": "http", + "method": "POST", + "path": "/v1/chat/completions", + "headers": [(b"authorization", f"Bearer {api_key}".encode()), (b"content-length", b"8388608")], + }, + receive=receive, + ) + + +async def _refusal_before_body( + cached_key: UserAPIKeyAuth, api_key: str, limiter: PROXY_MaxParallelRequestsHandler_v3 | None +) -> ProxyException: + from litellm.proxy.utils import hash_token + + user_api_key_cache: Final = DualCache() + await cache_key_object( + hashed_token=hash_token(api_key), + user_api_key_obj=cached_key, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=None, + ) + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.get_proxy_hook.return_value = limiter + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + with ( + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.open_telemetry_logger", None), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), + patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj), + pytest.raises(ProxyException) as exc, + ): + await user_api_key_auth(request=_request_whose_body_must_not_be_read(api_key), api_key=f"Bearer {api_key}") + proxy_logging_obj.post_call_failure_hook.assert_awaited_once() + return exc.value + + +@pytest.mark.asyncio +async def test_rate_limited_cached_key_is_refused_before_the_body_is_read() -> None: + from litellm.proxy.utils import hash_token + + api_key: Final = "sk-flooding-key" + cached_key: Final = UserAPIKeyAuth(token=hash_token(api_key), rpm_limit=1, models=["gpt-5.6"]) + limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(DualCache())) + await limiter.async_pre_call_hook( + UserAPIKeyAuth(api_key=hash_token(api_key), rpm_limit=1), + DualCache(), + {"model": "gpt-5.6", "messages": []}, + "acompletion", + ) + + refusal: Final = await _refusal_before_body(cached_key, api_key, limiter) + + assert refusal.code == "429", refusal.message + assert refusal.headers["rate_limited_by"] == "api_key" + + +@pytest.mark.asyncio +async def test_blocked_cached_key_is_refused_before_the_body_is_read() -> None: + from litellm.proxy.utils import hash_token + + api_key: Final = "sk-blocked-key" + cached_key: Final = UserAPIKeyAuth(token=hash_token(api_key), blocked=True, models=["gpt-5.6"]) + + refusal: Final = await _refusal_before_body(cached_key, api_key, None) + + assert refusal.code == "401", refusal.message + assert "Key is blocked" in refusal.message diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 40144736ebd..e6662736da6 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -10178,3 +10178,32 @@ async def test_team_model_limits_stop_being_hard_caps_when_borrowing_is_on( assert exc.value.status_code == 429 assert exc.value.headers is not None assert exc.value.headers["rate_limit_type"] == rate_limit_type + + +@pytest.mark.asyncio +async def test_raise_if_over_limit_before_body_refuses_an_exhausted_key_without_charging_it() -> None: + cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-exhausted"), rpm_limit=1) + requests: Final = _counter("api_key", auth.api_key, "requests") + + assert await handler.raise_if_over_limit_before_body(auth) is None + await handler.async_pre_call_hook(auth, cache, {"model": "gpt-5.6", "messages": []}, "acompletion") + charged: Final = _local(handler, requests) + + with pytest.raises(ProxyRateLimitError) as exc: + await handler.raise_if_over_limit_before_body(auth) + assert exc.value.status_code == 429 + assert exc.value.headers["rate_limited_by"] == "api_key" + assert _local(handler, requests) == charged, "the headers-only refusal must not count as a request" + + +@pytest.mark.asyncio +async def test_raise_if_over_limit_before_body_skips_limits_that_need_the_model() -> None: + cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-per-model"), metadata={"model_rpm_limit": {"gpt-5.6": 1}}) + + await handler.async_pre_call_hook(auth, cache, {"model": "gpt-5.6", "messages": []}, "acompletion") + + assert await handler.raise_if_over_limit_before_body(auth) is None