diff --git a/litellm/constants.py b/litellm/constants.py index 67bc6bc43fd..274854e1c40 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1654,6 +1654,7 @@ DEFAULT_SOFT_BUDGET: Final = get_env_float( RATE_LIMIT_ERROR_MESSAGE_FOR_VIRTUAL_KEY: Final = "LiteLLM Virtual Key user_api_key_hash" # Prefix of the 401 raised when a submitted virtual key is not shaped like one. INVALID_VIRTUAL_KEY_ERROR_MESSAGE: Final = "LiteLLM Virtual Key expected" +INVALID_API_KEY_ERROR_MESSAGE: Final = "Authentication Error, Invalid API key provided." # Attribute stamped on that 401 at its raise site so log routing recognises it by # provenance. Message text is caller-influenceable on other 401s, so it must not # be used to classify. diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 93dbfa9a9d2..85c3f741efb 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -36,6 +36,7 @@ from litellm.constants import ( DEFAULT_MAX_RECURSE_DEPTH, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, END_USER_RESTRICTED_REGISTRY_MAX_SIZE, + INVALID_API_KEY_ERROR_MESSAGE, LITELLM_PROXY_MASTER_KEY_ALIAS, MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE, REGISTRY_ERROR_NEGATIVE_CACHE_TTL, @@ -4225,7 +4226,7 @@ async def get_key_object( if _valid_token is None: raise ProxyException( - message=f"Authentication Error, Invalid proxy server token passed. key={hashed_token}, not found in db. Create key via `/key/generate` call.", + message=INVALID_API_KEY_ERROR_MESSAGE, type=ProxyErrorTypes.token_not_found_in_db, param="key", code=status.HTTP_401_UNAUTHORIZED, diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 0266131fc33..ae473a62271 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -28,6 +28,7 @@ from litellm.proxy.auth.auth_utils import ( # noqa: F401 # legacy module expor normalize_request_route, ) from litellm.proxy.auth.model_access_denied import ModelAccessDeniedHTTPException +from litellm.proxy.auth.resolvers.exceptions import KeyNotFoundError from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.types.services import ServiceTypes @@ -124,6 +125,15 @@ def _identity_log_suffix(resolved_identity: UserAPIKeyAuth | None) -> str: return f"\nKey Identity: {known}" if known else "" +def _auth_failure_log_message(e: Exception) -> str: + if not isinstance(e, KeyNotFoundError): + return str(e) + return ( + "Authentication Error, Invalid proxy server token passed. " + f"key={_escape_control_chars(e.hashed_token)}, not found in db. Create key via `/key/generate` call." + ) + + class UserAPIKeyAuthExceptionHandler: @staticmethod async def handle_authentication_error( @@ -197,7 +207,7 @@ class UserAPIKeyAuthExceptionHandler: logger.log( logging.WARNING if is_quiet_log else logging.ERROR, "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - %s\nRequester IP Address:%s%s", - e, + _auth_failure_log_message(e), requester_ip, _identity_log_suffix(resolved_identity), exc_info=True if litellm.log_client_error_tracebacks or not is_expected_client_error(e) else None, diff --git a/litellm/proxy/auth/resolvers/exceptions.py b/litellm/proxy/auth/resolvers/exceptions.py index fd57ba94aea..b7c071bd6b4 100644 --- a/litellm/proxy/auth/resolvers/exceptions.py +++ b/litellm/proxy/auth/resolvers/exceptions.py @@ -1,7 +1,10 @@ from __future__ import annotations +from typing import Final + from fastapi import status +from litellm.constants import INVALID_API_KEY_ERROR_MESSAGE from litellm.proxy._types import ProxyErrorTypes, ProxyException @@ -27,9 +30,10 @@ class KeyNotFoundError(IdentityResolutionError, ProxyException): """ def __init__(self, hashed_token: str) -> None: + self.hashed_token: Final = hashed_token ProxyException.__init__( self, - message=f"Authentication Error, Invalid proxy server token passed. key={hashed_token}, not found in db. Create key via `/key/generate` call.", + message=INVALID_API_KEY_ERROR_MESSAGE, type=ProxyErrorTypes.token_not_found_in_db, param="key", code=status.HTTP_401_UNAUTHORIZED, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 46eb3662283..a469c378db6 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2758,24 +2758,18 @@ async def user_api_key_auth_builder( type(api_key) if api_key is not None else "None" ) ) - abbreviated_api_key: Final = abbreviate_api_key(api_key=api_key) if api_key.startswith("sk-"): api_key = hash_token(token=api_key) - try: - with tracer.trace("litellm.proxy.auth.get_key_object_from_db"): - valid_token = IdentityStore.key_from_principal( - await IdentityStore( - prisma_client, - user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ).resolve(hashed_token=api_key) - ) - except ProxyException as e: - if e.code == 401 or e.code == "401": - e.message = f"Authentication Error, Invalid proxy server token passed. Received API Key = {abbreviated_api_key}, Key Hash (Token) ={api_key}. Unable to find token in cache or `LiteLLM_VerificationTokenTable`" - raise e + with tracer.trace("litellm.proxy.auth.get_key_object_from_db"): + valid_token = IdentityStore.key_from_principal( + await IdentityStore( + prisma_client, + user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ).resolve(hashed_token=api_key) + ) # update end-user params on valid token # These can change per request - it's important to update them here valid_token.end_user_id = end_user_params.get("end_user_id") diff --git a/tests/integration/authorization/test_jwt_key_mapping_lifecycle.py b/tests/integration/authorization/test_jwt_key_mapping_lifecycle.py index 03ca94dd561..8c678cc1a92 100644 --- a/tests/integration/authorization/test_jwt_key_mapping_lifecycle.py +++ b/tests/integration/authorization/test_jwt_key_mapping_lifecycle.py @@ -22,7 +22,7 @@ from tests.integration._support.client import ( eventually, gateway_from_environment, ) -from tests.integration._support.database import read_rows +from tests.integration._support.database import read_rows, write_rows from tests.integration._support.process import owned_proxy from tests.integration._support.wire import Reply, Request, Wire, wire_server @@ -118,6 +118,11 @@ class _ErrorResponse(BaseModel): error: _ProxyError +_UNKNOWN_KEY_ERROR: Final = _ProxyError( + message="Authentication Error, Invalid API key provided.", type="token_not_found_in_db", param="key", code="401" +) + + @dataclass(frozen=True, slots=True) class _JWTLifecycleRuntime: first: Gateway @@ -729,6 +734,49 @@ def test_deleting_a_jwt_mapped_key_stops_its_mapping_on_both_workers( _assert_upstream_requests(upstream, tuple(_expected_chat_request(backend, prompt) for prompt in prompts[:2])) +def test_a_cached_mapping_to_a_key_row_that_is_gone_answers_the_generic_401_on_both_workers( + jwt_lifecycle: _JWTLifecycleRuntime, +) -> None: + first: Final = jwt_lifecycle.first + second: Final = jwt_lifecycle.second + upstream: Final = jwt_lifecycle.upstream + model: Final = jwt_lifecycle.model + client_id: Final = "integration-client-" + uuid.uuid4().hex + prompts: Final = ( + "stale mapping warm primary " + uuid.uuid4().hex, + "stale mapping warm peer " + uuid.uuid4().hex, + "stale mapping primary " + uuid.uuid4().hex, + "stale mapping peer " + uuid.uuid4().hex, + ) + token: Final = _signed_jwt(jwt_lifecycle.private_key_a, client_id, jwt_lifecycle.issuer_a) + upstream.drain() + with first.scenario() as scenario: + key: Final = _create_key(first, model) + scenario.cleanups.callback(delete_key_if_present, first, key) + key_hash: Final = sha256(key.encode()).hexdigest() + mapping: Final = _create_mapping(first, key, client_id) + scenario.cleanups.callback(_delete_mapping_if_present, first, mapping.id) + _assert_chat(first, model, token, prompts[0], key_hash) + _assert_chat(second, model, token, prompts[1], key_hash) + + update: Final = first.request("POST", "/key/update", {"key": key, "metadata": {"rotated": "out-of-band"}}) + assert update.status_code == 200, update.text + write_rows('DELETE FROM "LiteLLM_VerificationToken" WHERE token=%s', (key_hash,)) + assert read_rows('SELECT id FROM "LiteLLM_JWTKeyMapping" WHERE id=%s', (mapping.id,)) == [] + + second_refusal: Final = eventually( + lambda: _chat(second, model, token, prompts[3]), + lambda response: response.status_code == 401, + seconds=3, + ) + first_refusal: Final = _chat(first, model, token, prompts[2]) + for refusal in (first_refusal, second_refusal): + assert refusal.status_code == 401, refusal.text + assert _ErrorResponse.model_validate_json(refusal.content).error == _UNKNOWN_KEY_ERROR, refusal.text + assert key_hash not in refusal.text + _assert_upstream_requests(upstream, tuple(_expected_chat_request(model, prompt) for prompt in prompts[:2])) + + def test_issuer_scoped_mappings_with_the_same_claim_resolve_per_issuer( jwt_lifecycle: _JWTLifecycleRuntime, ) -> None: diff --git a/tests/integration/authorization/test_unknown_key_generic_401.py b/tests/integration/authorization/test_unknown_key_generic_401.py new file mode 100644 index 00000000000..8e6ff3fa3ab --- /dev/null +++ b/tests/integration/authorization/test_unknown_key_generic_401.py @@ -0,0 +1,403 @@ +import asyncio +import json +import uuid +from collections.abc import Mapping +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import Final, Literal + +import anthropic +import httpx +import openai +import pytest +import websockets +from pydantic import BaseModel, ConfigDict, JsonValue +from websockets.exceptions import InvalidStatus + +from tests.integration._support.client import Gateway, delete_key_if_present, eventually, object_value, string_value +from tests.integration._support.database import read_rows, write_rows +from tests.integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from tests.integration._support.provider import SharedProvider + +MODEL: Final = "anthropic/claude-haiku-4-5" +GENERIC_MESSAGE: Final = "Authentication Error, Invalid API key provided." +CHAT: Final[Mapping[str, JsonValue]] = {"model": MODEL, "messages": [{"role": "user", "content": "unknown key"}]} +MESSAGES: Final[Mapping[str, JsonValue]] = { + "model": MODEL, + "max_tokens": 16, + "messages": [{"role": "user", "content": "unknown key"}], +} +RESPONSES: Final[Mapping[str, JsonValue]] = {"model": MODEL, "input": "unknown key"} +GEMINI: Final[Mapping[str, JsonValue]] = {"contents": [{"parts": [{"text": "unknown key"}]}]} + +Credential = Literal["authorization", "x-api-key", "query"] +Side = Literal["gateway", "peer"] + + +class _ProxyError(BaseModel): + model_config = ConfigDict(frozen=True) + + message: str + type: str + param: str + code: str + + +class _ErrorResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + error: _ProxyError + + +GENERIC_ERROR: Final = _ProxyError(message=GENERIC_MESSAGE, type="token_not_found_in_db", param="key", code="401") + + +@dataclass(frozen=True, slots=True) +class _Route: + method: Literal["GET", "POST"] + path: str + body: Mapping[str, JsonValue] | None + credential: Credential + + +ROUTES: Final[Mapping[str, _Route]] = { + "chat": _Route("POST", "/v1/chat/completions", CHAT, "authorization"), + "chat_stream": _Route("POST", "/v1/chat/completions", {**CHAT, "stream": True}, "authorization"), + "messages": _Route("POST", "/v1/messages", MESSAGES, "x-api-key"), + "messages_stream": _Route("POST", "/v1/messages", {**MESSAGES, "stream": True}, "x-api-key"), + "messages_bearer": _Route("POST", "/v1/messages", MESSAGES, "authorization"), + "responses": _Route("POST", "/v1/responses", RESPONSES, "authorization"), + "responses_stream": _Route("POST", "/v1/responses", {**RESPONSES, "stream": True}, "authorization"), + "embeddings": _Route("POST", "/v1/embeddings", {"model": MODEL, "input": "unknown key"}, "authorization"), + "completions": _Route("POST", "/v1/completions", {"model": MODEL, "prompt": "unknown key"}, "authorization"), + "models": _Route("GET", "/v1/models", None, "authorization"), + "anthropic_passthrough": _Route("POST", "/anthropic/v1/messages", MESSAGES, "x-api-key"), + "gemini_passthrough": _Route("POST", "/gemini/v1beta/models/gemini-2.5-flash:generateContent", GEMINI, "query"), + "key_info": _Route("GET", "/key/info", None, "authorization"), + "user_info": _Route("GET", "/user/info", None, "authorization"), +} +BURST_ROUTES: Final = ("chat", "chat_stream", "messages", "messages_stream", "responses", "responses_stream") + + +def _unknown_key() -> str: + return f"sk-{uuid.uuid4().hex}" + + +def _digest(key: str) -> str: + return sha256(key.encode()).hexdigest() + + +def _headers(credential: Credential, key: str) -> dict[str, str]: + match credential: + case "authorization": + return {"Authorization": f"Bearer {key}"} + case "x-api-key": + return {"x-api-key": key, "anthropic-version": "2023-06-01"} + case "query": + return {} + + +def _params(credential: Credential, key: str) -> dict[str, str]: + return {"key": key} if credential == "query" else {} + + +def _send(gateway: Gateway, route: _Route, key: str) -> httpx.Response: + return gateway.client.request( + route.method, + route.path, + json=None if route.body is None else dict(route.body), + headers=_headers(route.credential, key), + params=_params(route.credential, key), + ) + + +def _assert_generic_refusal(response: httpx.Response, *secrets: str) -> None: + assert response.status_code == 401, response.text + assert _ErrorResponse.model_validate_json(response.content).error == GENERIC_ERROR, response.text + for secret in secrets: + assert secret not in response.text, response.text + + +def _side(gateway: Gateway, peer: Gateway, side: Side) -> Gateway: + return gateway if side == "gateway" else peer + + +@pytest.mark.parametrize("side", ("gateway", "peer")) +@pytest.mark.parametrize("route", tuple(ROUTES)) +def test_an_unknown_key_gets_the_generic_401_without_reaching_the_provider( + gateway: Gateway, peer: Gateway, provider: SharedProvider, route: str, side: Side +) -> None: + key: Final = _unknown_key() + _assert_generic_refusal(_send(_side(gateway, peer, side), ROUTES[route], key), key, key[-4:], _digest(key)) + assert provider.received() == () + + +def test_the_openai_sdk_raises_authentication_error_carrying_the_generic_message( + gateway: Gateway, provider: SharedProvider +) -> None: + key: Final = _unknown_key() + client: Final = openai.OpenAI(base_url=f"{gateway.client.base_url}/v1", api_key=key, max_retries=0) + with pytest.raises(openai.AuthenticationError) as raised: + client.chat.completions.create(model=MODEL, messages=[{"role": "user", "content": "unknown key"}]) + assert raised.value.status_code == 401 + assert GENERIC_MESSAGE in raised.value.message + assert _digest(key) not in raised.value.message + assert provider.received() == () + + +def test_the_anthropic_sdk_raises_authentication_error_carrying_the_generic_message( + gateway: Gateway, provider: SharedProvider +) -> None: + key: Final = _unknown_key() + client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=key, max_retries=0) + with pytest.raises(anthropic.AuthenticationError) as raised: + client.messages.create(model=MODEL, max_tokens=16, messages=[{"role": "user", "content": "unknown key"}]) + assert raised.value.status_code == 401 + assert GENERIC_MESSAGE in raised.value.message + assert _digest(key) not in raised.value.message + assert provider.received() == () + + +HOSTILE_KEYS: Final[Mapping[str, str]] = {"five_kb": f"sk-{'a' * 5120}", "bare_prefix": "sk-"} + + +@pytest.mark.parametrize("hostile", tuple(HOSTILE_KEYS)) +def test_a_hostile_unknown_key_gets_the_generic_401_and_valid_keys_keep_working( + gateway: Gateway, peer: Gateway, provider: SharedProvider, hostile: str +) -> None: + key: Final = HOSTILE_KEYS[hostile] + for candidate in (gateway, peer): + _assert_generic_refusal(_send(candidate, ROUTES["chat"], key), _digest(key)) + with gateway.scenario() as scenario: + valid: Final = scenario.key(models=[MODEL]) + listed: Final = tuple(candidate.request("GET", "/v1/models", key=valid) for candidate in (gateway, peer)) + assert [response.status_code for response in listed] == [200, 200], [response.text for response in listed] + assert gateway.client.get("/health/liveliness").status_code == 200 + assert provider.received() == () + + +def test_two_different_unknown_keys_in_both_headers_get_the_generic_401_naming_neither( + gateway: Gateway, provider: SharedProvider +) -> None: + bearer: Final = _unknown_key() + api_key: Final = _unknown_key() + headers: Final = {**_headers("authorization", bearer), **_headers("x-api-key", api_key)} + for path, body in (("/v1/chat/completions", CHAT), ("/v1/messages", MESSAGES)): + _assert_generic_refusal( + gateway.client.post(path, json=dict(body), headers=headers), + bearer, + api_key, + _digest(bearer), + _digest(api_key), + ) + assert provider.received() == () + + +def test_the_same_unknown_key_sent_twice_across_both_proxies_gets_the_same_generic_401( + gateway: Gateway, peer: Gateway, provider: SharedProvider +) -> None: + key: Final = _unknown_key() + responses: Final = tuple(_send(candidate, ROUTES["chat"], key) for candidate in (gateway, peer, gateway, peer)) + for response in responses: + _assert_generic_refusal(response, key, _digest(key)) + assert len({response.content for response in responses}) == 1 + assert provider.received() == () + + +def _refused_once_evicted(candidate: Gateway, key: str) -> httpx.Response: + return eventually( + lambda: httpx.get( + f"{candidate.client.base_url}/v1/models", headers={"Authorization": f"Bearer {key}"}, trust_env=False + ), + lambda response: response.status_code == 401, + seconds=75, + ) + + +def test_a_deleted_key_gets_the_generic_401_on_both_proxies( + gateway: Gateway, peer: Gateway, provider: SharedProvider +) -> None: + key: Final = string_value(gateway.post("/key/generate", {"models": [MODEL]})["key"]) + try: + warmed: Final = tuple(candidate.request("GET", "/v1/models", key=key) for candidate in (gateway, peer)) + assert [response.status_code for response in warmed] == [200, 200], [response.text for response in warmed] + gateway.post("/key/delete", {"keys": [key]}) + for refusal in (_refused_once_evicted(gateway, key), _refused_once_evicted(peer, key)): + _assert_generic_refusal(refusal, key, _digest(key)) + finally: + delete_key_if_present(gateway, key) + assert provider.received() == () + + +def test_an_unknown_key_failure_row_keeps_its_hash_and_records_the_generic_message( + gateway: Gateway, provider: SharedProvider +) -> None: + key: Final = _unknown_key() + _assert_generic_refusal(_send(gateway, ROUTES["chat"], key), key) + rows: Final = eventually( + lambda: read_rows('SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (_digest(key),)), + lambda found: len(found) == 1, + seconds=15, + ) + assert rows[0]["status"] == "failure" + error: Final = object_value(object_value(rows[0]["metadata"])["error_information"]) + assert (error["error_class"], error["error_code"], error["error_message"]) == ( + "KeyNotFoundError", + "401", + GENERIC_MESSAGE, + ) + assert key not in json.dumps(rows[0]["metadata"]) + assert provider.received() == () + + +def _control_error(response: httpx.Response) -> _ProxyError: + assert response.status_code == 401, response.text + return _ErrorResponse.model_validate_json(response.content).error + + +@pytest.mark.parametrize("side", ("gateway", "peer")) +def test_requests_without_an_sk_key_keep_their_own_401(gateway: Gateway, peer: Gateway, side: Side) -> None: + client: Final = _side(gateway, peer, side).client + assert _control_error(client.get("/v1/models")) == _ProxyError( + message="Authentication Error, No api key passed in.", type="auth_error", param="None", code="401" + ) + assert _control_error(client.get("/v1/models", headers={"Authorization": "Bearer not-a-real-key"})) == _ProxyError( + message="LiteLLM Virtual Key expected. Received=not-****-key, expected to start with 'sk-'.", + type="auth_error", + param="None", + code="401", + ) + + +def test_a_blocked_key_keeps_its_own_401_on_both_proxies(gateway: Gateway, peer: Gateway) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key(models=[MODEL]) + gateway.post("/key/block", {"key": key}) + expected: Final = _ProxyError( + message="Authentication Error, Key is blocked. Update via `/key/unblock` if you're an admin.", + type="auth_error", + param="None", + code="401", + ) + for candidate in (gateway, peer): + assert _control_error(_refused_once_evicted(candidate, key)) == expected + + +def test_an_expired_key_keeps_its_own_401_on_both_proxies(gateway: Gateway, peer: Gateway) -> None: + with gateway.scenario() as scenario: + key: Final = scenario.key(models=[MODEL]) + write_rows( + "UPDATE \"LiteLLM_VerificationToken\" SET expires = now() - interval '1 hour' WHERE token=%s", + (_digest(key),), + ) + errors: Final = tuple( + _control_error(candidate.request("GET", "/v1/models", key=key)) for candidate in (gateway, peer) + ) + abbreviations: Final = {f"sk-...{key[-4:]}", f"sk-...{_digest(key)[-4:]}"} + for error in errors: + assert error.message.startswith("Authentication Error - Expired Key. Key Expiry time "), error + assert (error.type, error.code) == ("expired_key", "401"), error + assert error.param in abbreviations, error + + +def test_the_mcp_endpoint_keeps_its_own_401_for_an_unknown_key(gateway: Gateway, peer: Gateway) -> None: + key: Final = _unknown_key() + initialize: Final = { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "audit", "version": "1"}, + }, + } + headers: Final = {"Authorization": f"Bearer {key}", "Accept": "application/json, text/event-stream"} + responses: Final = tuple( + candidate.client.post("/mcp/", json=initialize, headers=headers) for candidate in (gateway, peer) + ) + for response in responses: + assert response.status_code == 401, response.text + assert response.json() == { + "detail": { + "error": "authentication_required", + "message": "Authenticate with the gateway to use the MCP endpoint.", + } + } + + +async def _realtime_refusal(candidate: Gateway, key: str) -> InvalidStatus: + url: Final = str(candidate.client.base_url).rstrip("/").replace("http://", "ws://") + with pytest.raises(InvalidStatus) as refusal: + async with websockets.connect( + f"{url}/v1/realtime?model={MODEL}", additional_headers={"Authorization": f"Bearer {key}"} + ): + pass + return refusal.value + + +async def test_the_realtime_upgrade_refuses_an_unknown_key_on_both_proxies_without_echoing_it( + gateway: Gateway, peer: Gateway, provider: SharedProvider +) -> None: + key: Final = _unknown_key() + refusals: Final = await asyncio.gather(*(_realtime_refusal(candidate, key) for candidate in (gateway, peer))) + assert [refusal.response.status_code for refusal in refusals] == [403, 403] + seen: Final = " ".join( + f"{refusal.response.reason_phrase} {refusal.response.headers} {refusal.response.body!r}" for refusal in refusals + ) + assert [secret for secret in (key, f"sk-...{key[-4:]}", _digest(key)) if secret in seen] == [] + assert provider.received() == () + + +async def _send_async(client: httpx.AsyncClient, route: _Route, key: str) -> httpx.Response: + return await client.request( + route.method, + route.path, + json=None if route.body is None else dict(route.body), + headers=_headers(route.credential, key), + ) + + +async def _burst(owned: OwnedProxy, keys: tuple[str, ...]) -> tuple[httpx.Response, ...]: + async with httpx.AsyncClient(base_url=str(owned.gateway.client.base_url), timeout=60, trust_env=False) as client: + return tuple( + await asyncio.gather( + *( + _send_async(client, ROUTES[BURST_ROUTES[index % len(BURST_ROUTES)]], key) + for index, key in enumerate(keys) + ) + ) + ) + + +def _log_after(owned: OwnedProxy, offset: int) -> str: + return owned.log.read_bytes()[offset:].decode(errors="replace") + + +def _console_line(key: str) -> str: + return f"key={_digest(key)}, not found in db" + + +@pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120)) +async def test_a_burst_across_a_worker_kill_answers_every_unknown_key_with_the_generic_401_and_logs_its_hash_once( + gateway: Gateway, tmp_path: Path +) -> None: + keys: Final = tuple(_unknown_key() for _ in range(40)) + with owned_proxy_process(gateway, tmp_path, {}, workers=2) as owned: + victim, _ = sorted(eventually(owned.worker_pids, lambda pids: len(pids) == 2)) + offset: Final = owned.log.stat().st_size + with owned.replacing_worker(victim) as kill: + kill.freeze() + during, _ = await asyncio.gather(_burst(owned, keys[:20]), asyncio.to_thread(kill)) + after: Final = await _burst(owned, keys[20:]) + window: Final = await asyncio.to_thread( + eventually, + lambda: _log_after(owned, offset), + lambda text: all(_console_line(key) in text for key in keys), + 30, + ) + for response, key in zip((*during, *after), keys, strict=True): + _assert_generic_refusal(response, key, _digest(key)) + assert [window.count(_console_line(key)) for key in keys] == [1] * len(keys) + assert [key for key in keys if key in window] == [] diff --git a/tests/integration/management/test_organization_lifecycle.py b/tests/integration/management/test_organization_lifecycle.py index 0193bf517b2..cdbf3c15b18 100644 --- a/tests/integration/management/test_organization_lifecycle.py +++ b/tests/integration/management/test_organization_lifecycle.py @@ -117,15 +117,10 @@ def _observation(text: str) -> dict[str, JsonValue]: } -def _token_not_found_error(key: str) -> dict[str, JsonValue]: - hashed: Final = sha256(key.encode()).hexdigest() +def _token_not_found_error() -> dict[str, JsonValue]: return { "error": { - "message": ( - "Authentication Error, Invalid proxy server token passed. " - f"Received API Key = sk-...{key[-4:]}, Key Hash (Token) ={hashed}. " - "Unable to find token in cache or `LiteLLM_VerificationTokenTable`" - ), + "message": "Authentication Error, Invalid API key provided.", "type": "token_not_found_in_db", "param": "key", "code": "401", @@ -191,14 +186,14 @@ def test_organization_delete_revokes_org_and_team_keys_on_gateway_and_peer(gatew for key in (organization_key, team_key): refused: Final = _chat(gateway, model, key, f"refused {uuid.uuid4().hex}") assert refused.status_code == 401, refused.text - assert refused.json() == _token_not_found_error(key), refused.text + assert refused.json() == _token_not_found_error(), refused.text for key in (organization_key, team_key): refused_peer: Final = eventually( lambda k=key: _chat(peer, model, k, f"refused {uuid.uuid4().hex}"), lambda response: response.status_code == 401, seconds=10, ) - assert refused_peer.json() == _token_not_found_error(key), refused_peer.text + assert refused_peer.json() == _token_not_found_error(), refused_peer.text refused_observations: Final = _observed(upstream) assert refused_observations == [], refused_observations for proxy in (gateway, peer): diff --git a/tests/integration/management/test_scim_user_lifecycle.py b/tests/integration/management/test_scim_user_lifecycle.py index 527728e0a60..18cb66b3d8b 100644 --- a/tests/integration/management/test_scim_user_lifecycle.py +++ b/tests/integration/management/test_scim_user_lifecycle.py @@ -451,11 +451,7 @@ def test_profile_put_without_groups_keeps_memberships_and_a_new_groups_list_move ) expected_refusal: Final = { "error": { - "message": ( - "Authentication Error, Invalid proxy server token passed. " - f"Received API Key = sk-...{key_a[-4:]}, Key Hash (Token) ={sha256(key_a.encode()).hexdigest()}. " - "Unable to find token in cache or `LiteLLM_VerificationTokenTable`" - ), + "message": "Authentication Error, Invalid API key provided.", "type": "token_not_found_in_db", "param": "key", "code": "401", diff --git a/tests/integration/management/test_user_delete_revokes_keys.py b/tests/integration/management/test_user_delete_revokes_keys.py index 0cf2db9c3e6..ba1d3ca42e3 100644 --- a/tests/integration/management/test_user_delete_revokes_keys.py +++ b/tests/integration/management/test_user_delete_revokes_keys.py @@ -40,15 +40,10 @@ def _observation(text: str) -> dict[str, JsonValue]: } -def _token_not_found_error(key: str) -> dict[str, JsonValue]: - hashed: Final = sha256(key.encode()).hexdigest() +def _token_not_found_error() -> dict[str, JsonValue]: return { "error": { - "message": ( - "Authentication Error, Invalid proxy server token passed. " - f"Received API Key = sk-...{key[-4:]}, Key Hash (Token) ={hashed}. " - "Unable to find token in cache or `LiteLLM_VerificationTokenTable`" - ), + "message": "Authentication Error, Invalid API key provided.", "type": "token_not_found_in_db", "param": "key", "code": "401", @@ -128,14 +123,14 @@ def test_user_delete_revokes_keys_and_memberships(gateway: Gateway, peer: Gatewa for key in keys: denied: Final = _chat(gateway, model, key, f"denied {uuid.uuid4().hex}") assert denied.status_code == 401, denied.text - assert denied.json() == _token_not_found_error(key), denied.text + assert denied.json() == _token_not_found_error(), denied.text for key in keys: denied_peer: Final = eventually( lambda k=key: _chat(peer, model, k, f"denied {uuid.uuid4().hex}"), lambda response: response.status_code == 401, seconds=10, ) - assert denied_peer.json() == _token_not_found_error(key), denied_peer.text + assert denied_peer.json() == _token_not_found_error(), denied_peer.text denied_observations: Final = _observed(upstream) assert denied_observations == [], denied_observations diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 1f8c85ce337..72179e661a3 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -647,6 +647,35 @@ async def test_get_key_object_should_raise_if_reconnect_fails_on_db_connection_e assert mock_prisma_client.get_data.await_count == 1 +@pytest.mark.asyncio +async def test_get_key_object_miss_keeps_the_looked_up_value_out_of_the_401_and_the_log(caplog): + """This 401 reaches unauthenticated callers through the auth pipeline, so its + message must not name the hash it looked up, the table it looked it up in, or + the endpoint that mints keys. Failure logging also calls this helper with the raw + key a caller sent, a provider key included, so a miss must not log that value.""" + from litellm.constants import INVALID_API_KEY_ERROR_MESSAGE + + mock_prisma_client = MagicMock() + mock_prisma_client.get_data = AsyncMock(return_value=None) + + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + with caplog.at_level("INFO", logger="LiteLLM Proxy"), pytest.raises(ProxyException) as exc_info: + await get_key_object( + hashed_token="a-secret-looking-hash", + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + + assert int(exc_info.value.code) == 401 + assert exc_info.value.message == INVALID_API_KEY_ERROR_MESSAGE + for leaked in ("a-secret-looking-hash", "VerificationToken", "/key/generate"): + assert leaked not in exc_info.value.message + assert "a-secret-looking-hash" not in caplog.text + + class _InFlightCountingPrisma: def __init__(self) -> None: self.in_flight = 0 diff --git a/tests/unit/proxy/auth/test_auth_exception_handler.py b/tests/unit/proxy/auth/test_auth_exception_handler.py index dd6e1f57fbc..d83a8aa2fcd 100644 --- a/tests/unit/proxy/auth/test_auth_exception_handler.py +++ b/tests/unit/proxy/auth/test_auth_exception_handler.py @@ -688,6 +688,69 @@ async def test_expired_key_error_log_names_the_key_owner( assert absent_fragment not in logged, logged +@pytest.mark.asyncio +@pytest.mark.parametrize( + "looked_up, logged_fragment, absent_fragment", + [ + pytest.param( + "d748053a" * 8, + "Exception occured - Authentication Error, Invalid proxy server token passed. key=" + + "d748053a" * 8 + + ", not found in db. Create key via `/key/generate` call.", + None, + id="operator_wording_and_hash_kept_in_log", + ), + pytest.param( + "abc\nforged log line", + "key=abc\\nforged log line, not found in db", + "\nforged log line", + id="control_chars_in_hash_cannot_forge_log_lines", + ), + ], +) +async def test_unknown_key_error_log_keeps_the_operator_wording_the_401_hides( + looked_up, logged_fragment, absent_fragment, caplog +): + """A wrong key gets a generic 401, while the one auth failure log line keeps the + operator wording and the hash that was looked up, next to the requester IP.""" + from litellm.constants import INVALID_API_KEY_ERROR_MESSAGE + from litellm.proxy.auth.resolvers.exceptions import KeyNotFoundError + + with ( + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=None, + ), + patch("litellm.proxy.auth.auth_exception_handler.seed_request_identity"), + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ), + ): + verbose_proxy_logger.propagate = True + try: + with caplog.at_level("ERROR", logger="LiteLLM Proxy"), pytest.raises(ProxyException) as exc_info: + await UserAPIKeyAuthExceptionHandler().handle_authentication_error( + KeyNotFoundError(looked_up), + MagicMock(), + {"model": "gpt-4o"}, + "/v1/chat/completions", + None, + "sk-raw-key", + ) + finally: + verbose_proxy_logger.propagate = False + + assert exc_info.value.message == INVALID_API_KEY_ERROR_MESSAGE + records = [r for r in caplog.records if "user_api_key_auth(): Exception occured" in r.getMessage()] + assert len(records) == 1, [r.getMessage() for r in caplog.records] + logged = records[0].getMessage() + assert "Requester IP Address:" in logged and logged_fragment in logged, logged + if absent_fragment is not None: + assert absent_fragment not in logged, logged + + @pytest.mark.asyncio async def test_auth_failure_without_resolved_identity_still_logs(): """When auth fails before any identity is resolved (e.g. an unknown key), diff --git a/tests/unit/proxy/auth/test_resolvers_exceptions.py b/tests/unit/proxy/auth/test_resolvers_exceptions.py index f88b7e2be14..31642230e4a 100644 --- a/tests/unit/proxy/auth/test_resolvers_exceptions.py +++ b/tests/unit/proxy/auth/test_resolvers_exceptions.py @@ -1,5 +1,6 @@ from __future__ import annotations +from litellm.constants import INVALID_API_KEY_ERROR_MESSAGE from litellm.proxy._types import ProxyErrorTypes, ProxyException from litellm.proxy.auth.resolvers.exceptions import ( IdentityResolutionError, @@ -31,5 +32,17 @@ def test_key_not_found_preserves_the_public_401_contract(): assert error.param == "key" +def test_key_not_found_message_hides_the_hash_but_the_raise_site_can_still_log_it(): + """The JWT key-mapping paths surface this exception's own message to the caller + without the rewrite the main key path applies, so the message itself has to be + generic. The hash stays reachable on the instance for the auth failure log line.""" + error = KeyNotFoundError("hashed-token") + + assert error.message == INVALID_API_KEY_ERROR_MESSAGE + assert "hashed-token" not in str(error) + assert "/key/generate" not in str(error) + assert error.hashed_token == "hashed-token" + + def test_key_not_in_cache_names_the_token(): assert "hashed-token" in str(KeyNotInCacheError("hashed-token")) diff --git a/tests/unit/proxy/auth/test_resolvers_store.py b/tests/unit/proxy/auth/test_resolvers_store.py index 5e644a084a5..607162000ab 100644 --- a/tests/unit/proxy/auth/test_resolvers_store.py +++ b/tests/unit/proxy/auth/test_resolvers_store.py @@ -4,8 +4,10 @@ from typing import Dict, Optional import pytest +from litellm.constants import INVALID_API_KEY_ERROR_MESSAGE from litellm.proxy._types import UserAPIKeyAuth, hash_token from litellm.proxy.auth.resolvers.exceptions import ( + KeyNotFoundError, NoDatabaseConnectionError, PrincipalMissingSourceKeyError, ) @@ -29,6 +31,14 @@ class _FakeCache: return None +class _FakeEmptyKeyTable: + """Stands in for the PrismaClient the store reads the key row through. Returning + no row is what a submitted key that was never issued looks like.""" + + async def get_data(self, *args, **kwargs): + return None + + async def test_resolve_returns_a_principal_projected_from_the_looked_up_key(): raw = "sk-live-abc" key = UserAPIKeyAuth(token=hash_token(raw), user_id="u-1", team_id="t-1") @@ -68,3 +78,17 @@ async def test_resolve_raises_without_a_db_connection(): store = IdentityStore(None, _FakeCache()) with pytest.raises(NoDatabaseConnectionError): await store.resolve(hashed_token="missing") + + +async def test_resolve_raises_a_generic_key_not_found_carrying_the_hash_when_no_key_row_exists(): + """The 401 message stays generic for the caller, while the hash rides on the + exception so the auth failure log line can still name it for the operator.""" + missing = hash_token("sk-live-not-a-key") + store = IdentityStore(_FakeEmptyKeyTable(), _FakeCache()) + + with pytest.raises(KeyNotFoundError) as exc_info: + await store.resolve(hashed_token=missing) + + assert exc_info.value.message == INVALID_API_KEY_ERROR_MESSAGE + assert missing not in str(exc_info.value) + assert exc_info.value.hashed_token == missing 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 a0441981447..bf5d57c2312 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 @@ -8132,6 +8132,31 @@ async def test_builder_returns_401_when_db_lookup_reports_missing_key(): assert int(exc_info.value.code) == status.HTTP_401_UNAUTHORIZED +@pytest.mark.asyncio +async def test_builder_401_for_unknown_key_discloses_nothing_about_the_lookup(): + """A wrong key reaches an unauthenticated caller, so its 401 body must carry no key + hash, no verification table, and no key-minting endpoint. Regression for the 401 that + shipped the key hash plus "Unable to find token in cache or + LiteLLM_VerificationTokenTable" to anyone who probed a public deployment. The hash the + operator still needs is on the auth failure log line, covered in test_auth_exception_handler.""" + from litellm.constants import INVALID_API_KEY_ERROR_MESSAGE + from litellm.proxy.auth.resolvers.exceptions import KeyNotFoundError + from litellm.proxy.proxy_server import hash_token + + submitted_key = "sk-db-lookup-test" + key_hash = hash_token(token=submitted_key) + get_key_object = AsyncMock(side_effect=KeyNotFoundError(key_hash)) + + with pytest.raises(ProxyException) as exc_info: + await _run_builder_with_key_lookup(get_key_object) + + body = json.dumps(exc_info.value.to_dict()) + assert int(exc_info.value.code) == status.HTTP_401_UNAUTHORIZED + assert exc_info.value.message == INVALID_API_KEY_ERROR_MESSAGE + for leaked in (key_hash, "VerificationToken", "/key/generate", submitted_key[-4:]): + assert leaked not in body, f"401 body leaks {leaked!r}: {body}" + + @pytest.mark.asyncio async def test_builder_succeeds_when_db_lookup_returns_valid_token(): """Regression guard: a valid key still authenticates. Proves the 503 diff --git a/tests/unit/proxy/management_endpoints/test_key_generate_prisma.py b/tests/unit/proxy/management_endpoints/test_key_generate_prisma.py index 3d9667ace1b..26f1603ccca 100644 --- a/tests/unit/proxy/management_endpoints/test_key_generate_prisma.py +++ b/tests/unit/proxy/management_endpoints/test_key_generate_prisma.py @@ -316,12 +316,12 @@ def test_call_with_invalid_key(prisma_client): print("got result", result) pytest.fail(f"This should have failed!. IT's an invalid key") - with pytest.raises(Exception, match="Authentication Error, Invalid proxy server token passed") as exc_info: + with pytest.raises(Exception, match="Authentication Error, Invalid API key provided") as exc_info: asyncio.run(test()) e = exc_info.value print("Got Exception", e) print(e.message) - assert "Authentication Error, Invalid proxy server token passed" in e.message + assert "Authentication Error, Invalid API key provided" in e.message @pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") @@ -3706,14 +3706,14 @@ async def test_auth_vertex_ai_route(prisma_client): request = Request(scope={"type": "http"}) request._url = URL(url=route) request._headers = {"Authorization": "Bearer sk-98765"} - with pytest.raises(Exception, match="Invalid proxy server token passed") as exc_info: + with pytest.raises(Exception, match="Invalid API key provided") as exc_info: await user_api_key_auth(request=request, api_key="Bearer " + "sk-98765") e = exc_info.value print(vars(e)) print("error str=", str(e.message)) error_str = str(e.message) assert e.code == "401" - assert "Invalid proxy server token passed" in error_str + assert "Invalid API key provided" in error_str @pytest.mark.asyncio