From aaa09bb0867a27837bb8895d27b2ee0710abb990 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Wed, 29 Apr 2026 00:11:43 +0000 Subject: [PATCH] Add negative caching for invalid virtual keys - Introduced `InvalidVirtualKeyCache` to manage negative caching for unknown virtual keys, reducing unnecessary database lookups for invalid tokens. - Updated `user_api_key_auth.py` to utilize the new cache for token validation. - Added a new configuration constant for negative cache TTL. - Implemented unit tests for the new caching mechanism to ensure correct behavior in various scenarios. --- litellm/constants.py | 3 + litellm/proxy/auth/reject_invalid_tokens.py | 196 ++++++++++++++++++ litellm/proxy/auth/user_api_key_auth.py | 41 ++-- test-config/dev_config.yaml | 96 +++++++++ .../test_reject_invalid_tokens.py | 116 +++++++++++ 5 files changed, 428 insertions(+), 24 deletions(-) create mode 100644 litellm/proxy/auth/reject_invalid_tokens.py create mode 100644 test-config/dev_config.yaml create mode 100644 tests/proxy_unit_tests/test_reject_invalid_tokens.py diff --git a/litellm/constants.py b/litellm/constants.py index fc9f5730cdf..098bd4883bd 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1538,6 +1538,9 @@ DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int( os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60) ) DEFAULT_ACCESS_GROUP_CACHE_TTL = int(os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600)) +DEFAULT_INVALID_VIRTUAL_KEY_NEGATIVE_CACHE_TTL_SECONDS = int( + os.getenv("DEFAULT_INVALID_VIRTUAL_KEY_NEGATIVE_CACHE_TTL_SECONDS", "3600") +) # Sentry Scrubbing Configuration SENTRY_DENYLIST = [ diff --git a/litellm/proxy/auth/reject_invalid_tokens.py b/litellm/proxy/auth/reject_invalid_tokens.py new file mode 100644 index 00000000000..06f2cf8ee27 --- /dev/null +++ b/litellm/proxy/auth/reject_invalid_tokens.py @@ -0,0 +1,196 @@ +""" +Negative cache for unknown virtual keys (reduces repeat DB work on bad ``sk-`` hashes). + +:class:`InvalidVirtualKeyCache` holds: + +- :meth:`InvalidVirtualKeyCache.configured_ttl_seconds` — TTL from proxy settings, else + :data:`~litellm.constants.DEFAULT_INVALID_VIRTUAL_KEY_NEGATIVE_CACHE_TTL_SECONDS`. +- :meth:`InvalidVirtualKeyCache.check_invalid_token` — ``sk-`` format, then optional negative cache + + ``LiteLLM_VerificationToken`` probe. Returns ``True`` if the client should get **401** (except + malformed keys, which raise ``HTTPException``). Returns ``False`` if preflight passed—hash the raw + key and load from ``combined_view``. +- :meth:`InvalidVirtualKeyCache.allows_db_lookup` / :meth:`InvalidVirtualKeyCache.record_miss` — lower-level helpers. + +Cache keys: ``invalid_vk:{hashed_token}``. +""" + +from __future__ import annotations + +from typing import Any, Dict, Optional, Union + +from fastapi import HTTPException, status + +from litellm.constants import DEFAULT_INVALID_VIRTUAL_KEY_NEGATIVE_CACHE_TTL_SECONDS +from litellm._logging import verbose_logger, verbose_proxy_logger +from litellm.proxy._types import hash_token + +INVALID_VIRTUAL_KEY_CACHE_PREFIX = "invalid_vk:" +class InvalidVirtualKeyCache: + """Settings + negative cache for virtual keys that are not in the DB (or not yet).""" + + _prefix = INVALID_VIRTUAL_KEY_CACHE_PREFIX + + @staticmethod + def configured_ttl_seconds( + general_settings: Union[Dict[str, Any], Any], + ) -> Optional[float]: + """ + TTL for negative-caching unknown virtual keys. + + Reads ``invalid_virtual_key_cache_ttl`` from ``general_settings`` (top-level or + ``litellm_settings``). If unset, uses + :data:`litellm.constants.DEFAULT_INVALID_VIRTUAL_KEY_NEGATIVE_CACHE_TTL_SECONDS`. + A configured value of ``0`` or less disables negative caching (returns ``None``). + """ + raw: Any = None + if isinstance(general_settings, dict): + raw = general_settings.get("invalid_virtual_key_cache_ttl") + if raw is None: + litellm_settings = general_settings.get("litellm_settings") + if isinstance(litellm_settings, dict): + raw = litellm_settings.get("invalid_virtual_key_cache_ttl") + default_ttl = float(DEFAULT_INVALID_VIRTUAL_KEY_NEGATIVE_CACHE_TTL_SECONDS) + try: + if raw is None: + return default_ttl + v = float(raw) + except (TypeError, ValueError): + return default_ttl + if v <= 0: + return None + return v + + @classmethod + def _cache_key(cls, hashed_token: str) -> str: + return "{}{}".format(cls._prefix, hashed_token) + + @classmethod + async def allows_db_lookup( + cls, + *, + hashed_token: str, + user_api_key_cache: Any, + ttl_seconds: Optional[float], + ) -> bool: + """ + ``True`` if this hash may hit the database (entry not in the negative cache). + + When ``ttl_seconds`` is ``None``, negative caching is off — always ``True``. + """ + if ttl_seconds is None: + return True + + key = cls._cache_key(hashed_token) + try: + cached = await user_api_key_cache.async_get_cache(key=key) + except Exception as e: + verbose_proxy_logger.debug( + "InvalidVirtualKeyCache.allows_db_lookup: cache read failed, allowing query: %s", + e, + ) + return True + + return cached is None + + @classmethod + async def record_miss( + cls, + *, + hashed_token: str, + user_api_key_cache: Any, + ttl_seconds: float, + ) -> None: + """Remember this hash as a failed lookup until ``ttl_seconds`` elapses.""" + if ttl_seconds <= 0: + return + key = cls._cache_key(hashed_token) + try: + await user_api_key_cache.async_set_cache( + key=key, + value="", + ttl=ttl_seconds, + ) + except Exception as e: + verbose_proxy_logger.debug("InvalidVirtualKeyCache.record_miss: %s", e) + + @classmethod + async def check_invalid_token( + cls, + *, + api_key: Any, + prisma_client: Any, + user_api_key_cache: Any, + general_settings: Any, + ) -> bool: + """ + Virtual-key preflight: ``sk-`` shape → (if TTL on) negative cache → + ``litellm_verificationtoken`` row check. + + Returns ``True`` if the request should be rejected as an invalid virtual key (**401**). + Returns ``False`` if preflight passed; caller should ``hash_token(api_key)`` then call + ``get_key_object``. + + Malformed keys raise ``HTTPException`` (401) with masking details instead of returning bool. + """ + ttl_seconds = cls.configured_ttl_seconds(general_settings) + + if isinstance(api_key, str): + _masked_key = ( + "{}****{}".format(api_key[:4], api_key[-4:]) + if len(api_key) > 8 + else "****" + ) + if not api_key.startswith("sk-"): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=( + "LiteLLM Virtual Key expected. Received={}, expected to start with 'sk-'.".format( + _masked_key + ) + ), + ) + else: + verbose_logger.warning( + "litellm.proxy.proxy_server.user_api_key_auth(): Warning - Key is not a string. Got type={}".format( + type(api_key) if api_key is not None else "None" + ) + ) + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="LiteLLM Virtual Key expected.", + ) + + hashed_token = hash_token(token=api_key) + + if ttl_seconds is None: + return False + + if not await cls.allows_db_lookup( + hashed_token=hashed_token, + user_api_key_cache=user_api_key_cache, + ttl_seconds=ttl_seconds, + ): + return True + + token_probe_failed = False + try: + token_row = await prisma_client.db.litellm_verificationtoken.find_first( + where={"token": hashed_token}, + ) + except Exception as e: + verbose_proxy_logger.debug( + "InvalidVirtualKeyCache.check_invalid_token: verification token probe failed, continuing to combined_view: %s", + e, + ) + token_probe_failed = True + token_row = None + + if not token_probe_failed and token_row is None: + await cls.record_miss( + hashed_token=hashed_token, + user_api_key_cache=user_api_key_cache, + ttl_seconds=ttl_seconds, + ) + return True + + return False diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index b8db3cd2a7b..357230affe2 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -57,6 +57,7 @@ from litellm.proxy.auth.auth_utils import ( from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request +from litellm.proxy.auth.reject_invalid_tokens import InvalidVirtualKeyCache from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator from litellm.proxy.common_utils.http_parsing_utils import ( @@ -1154,7 +1155,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) ## Check DB - if ( prisma_client is None ): # if both master key + user key submitted, and user key != master key, and no db connected, raise an error @@ -1166,29 +1166,22 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) if valid_token is None: - if isinstance( - api_key, str - ): # if generated token, make sure it starts with sk-. - _masked_key = ( - "{}****{}".format(api_key[:4], api_key[-4:]) - if len(api_key) > 8 - else "****" + if await InvalidVirtualKeyCache.check_invalid_token( + api_key=api_key, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + general_settings=general_settings, + ): + raise ProxyException( + message="Authentication Error at InvalidVirtualKeyCache, Invalid proxy server token passed. Token (hash) = {}. Unable to find token in cache or `LiteLLM_VerificationTokenTable`".format( + hash_token(token=api_key), + ), + type=ProxyErrorTypes.token_not_found_in_db, + param="key", + code=status.HTTP_401_UNAUTHORIZED, ) - assert api_key.startswith( - "sk-" - ), "LiteLLM Virtual Key expected. Received={}, expected to start with 'sk-'.".format( - _masked_key - ) # prevent token hashes from being used - else: - verbose_logger.warning( - "litellm.proxy.proxy_server.user_api_key_auth(): Warning - Key is not a string. Got type={}".format( - type(api_key) if api_key is not None else "None" - ) - ) - abbreviated_api_key = abbreviate_api_key(api_key=api_key) - if api_key.startswith("sk-"): - api_key = hash_token(token=api_key) + api_key = hash_token(token=api_key) try: with tracer.trace("litellm.proxy.auth.get_key_object_from_db"): valid_token = await get_key_object( @@ -1200,8 +1193,8 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) except ProxyException as e: if e.code == 401 or e.code == "401": - e.message = "Authentication Error, Invalid proxy server token passed. Received API Key = {}, Key Hash (Token) ={}. Unable to find token in cache or `LiteLLM_VerificationTokenTable`".format( - abbreviated_api_key, api_key + e.message = "Authentication Error, Invalid proxy server token passed. Token (hash) = {}. Unable to find token in cache or `LiteLLM_VerificationTokenTable`".format( + api_key ) raise e # update end-user params on valid token diff --git a/test-config/dev_config.yaml b/test-config/dev_config.yaml new file mode 100644 index 00000000000..6e5c29e1275 --- /dev/null +++ b/test-config/dev_config.yaml @@ -0,0 +1,96 @@ +model_list: + - model_name: fake-openai-endpoint + litellm_params: + model: openai/fake-model + api_key: fake-key + api_base: https://exampleopenaiendpoint-production.up.railway.app/ + timeout: 40 + + - model_name: claude-opus-4-7 + litellm_params: + model: anthropic/claude-opus-4-7 + api_key: os.environ/ANTHROPIC_API_KEY + + # Bedrock: model id used in /bedrock/model/{...}/converse (must match the string clients send as `model`) + - model_name: global.anthropic.claude-sonnet-4-6 + litellm_params: + model: bedrock/global.anthropic.claude-sonnet-4-6 + custom_llm_provider: bedrock + aws_region_name: os.environ/AWS_REGION + # Credentials: AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY / AWS_REGION (or instance role) + + # Bedrock pass-through to a local mock: run `uv run python scripts/mock_bedrock_passthrough_target.py --port 9999` + # Test: curl -X POST "http://127.0.0.1:4000/bedrock/model/mock-bedrock-claude/converse" -H "Authorization: Bearer sk-1234" -H "Content-Type: application/json" -d '{"messages":[{"role":"user","content":[{"text":"hi"}]}]}' + # Do NOT use /bedrock/v1/messages (that is not a Bedrock path). For Anthropic /v1/messages use POST /v1/messages and set ANTHROPIC_BASE_URL to http://127.0.0.1:4000 (no /bedrock). + - model_name: mock-bedrock-claude + litellm_params: + model: bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0 + custom_llm_provider: bedrock + aws_region_name: us-west-2 + api_base: "http://127.0.0.1:9999" + + - model_name: fake-openai-gpt4 + litellm_params: + model: gpt-4 + api_key: fake-key + # Docker Compose service `mock-llm-provider` (host:8090). Requires network_mock: False for real HTTP. + api_base: http://mock-llm-provider/llm/ + + - model_name: fake-openai-gpt4 + litellm_params: + model: gpt-4 + api_key: fake-key + # Docker Compose service `mock-llm-provider` (host:8090). Requires network_mock: False for real HTTP. + api_base: http://mock-llm-provider/llm/ + +general_settings: + master_key: sk-1234 + health_check_details: False + database_url: os.environ/DATABASE_URL + # disable_spend_updates: True + # use_redis_transaction_buffer: True + +litellm_settings: + network_mock: False # must be false to call mock-llm-provider over the Docker network + json_logs: True + drop_params: True + telemetry: False + public_routes: ["LiteLLMRoutes.public_routes", "/health/liveliness"] + num_retries: 0 + set_verbose: False + request_timeout: 600 + enable_redis_auth_cache: True # Share virtual-key auth cache across workers via Redis. + # Eliminates per-pod DB round-trips on cache misses. Requires cache.type=redis below. + cache: true + cache_params: + type: redis + host: os.environ/RATELIMIT_REDIS_ENDPOINT + port: os.environ/RATELIMIT_REDIS_PORT + max_connections: 100 + callbacks: + - prometheus + # - "callbacks.overhead_metrics.overhead_logger" + # - "callbacks.otel_tracing.otel_tracing_logger" + # - "callbacks.usage_metrics.usage_metrics_logger" + service_callback: ["prometheus_system"] + default_key_generate_params: + max_budget: 500 + budget_duration: "30d" + +# Router Redis (same Compose Redis). Omit redis_password in `.env` if Redis has no auth. +router_settings: + routing_strategy: simple-shuffle + redis_host: os.environ/RATELIMIT_REDIS_ENDPOINT + redis_port: os.environ/RATELIMIT_REDIS_PORT + # redis_password: os.environ/REDIS_DEFAULT_PASSWORD + +# Uncomment when `opinionated_api` is installed (iFood production shape): +# guardrails: +# - guardrail_name: "llm-firewall" +# litellm_params: +# guardrail: opinionated_api.custom_guardrails.llm_firewall.LLMFirewall +# mode: "during_call" +# - guardrail_name: "llm-firewall-post" +# litellm_params: +# guardrail: opinionated_api.custom_guardrails.llm_firewall.LLMFirewall +# mode: "post_call" diff --git a/tests/proxy_unit_tests/test_reject_invalid_tokens.py b/tests/proxy_unit_tests/test_reject_invalid_tokens.py new file mode 100644 index 00000000000..28979c11ed7 --- /dev/null +++ b/tests/proxy_unit_tests/test_reject_invalid_tokens.py @@ -0,0 +1,116 @@ +# Unit tests for litellm.proxy.auth.reject_invalid_tokens.InvalidVirtualKeyCache + +import os +import sys +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm.proxy._types import hash_token +from litellm.proxy.auth.reject_invalid_tokens import InvalidVirtualKeyCache + + +def _sk_key() -> str: + return "sk-test-invalid-virtual-key" + +def _general_settings_positive_ttl() -> dict: + """Force negative-cache path on (avoid relying only on default constant).""" + return {"invalid_virtual_key_cache_ttl": 3600} + + +def _negative_cache_key_for(api_key: str) -> str: + return InvalidVirtualKeyCache._cache_key(hash_token(token=api_key)) + + +@pytest.mark.asyncio +async def test_check_invalid_token_empty_cache_db_miss_records_negative_entry(): + """ + Negative cache empty, Prisma finds no verification row → reject (True); + miss is recorded for repeat traffic. + """ + api_key = _sk_key() + hashed = hash_token(token=api_key) + neg_key = _negative_cache_key_for(api_key) + + user_api_key_cache = MagicMock() + user_api_key_cache.async_get_cache = AsyncMock(return_value=None) + user_api_key_cache.async_set_cache = AsyncMock() + + find_first = AsyncMock(return_value=None) + prisma_client = MagicMock() + prisma_client.db.litellm_verificationtoken.find_first = find_first + + result = await InvalidVirtualKeyCache.check_invalid_token( + api_key=api_key, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + general_settings=_general_settings_positive_ttl(), + ) + + assert result is True + find_first.assert_awaited_once_with(where={"token": hashed}) + user_api_key_cache.async_get_cache.assert_awaited_once_with(key=neg_key) + user_api_key_cache.async_set_cache.assert_awaited_once_with( + key=neg_key, + value="", + ttl=3600.0, + ) + + +@pytest.mark.asyncio +async def test_check_invalid_token_negative_cache_hit_short_circuits_even_if_db_has_row(): + """ + Hash already negative-cached → reject immediately without calling Prisma, + even if a row exists in DB (stale negative cache after key creation is possible). + """ + api_key = _sk_key() + neg_key = _negative_cache_key_for(api_key) + + user_api_key_cache = MagicMock() + user_api_key_cache.async_get_cache = AsyncMock(return_value="") + user_api_key_cache.async_set_cache = AsyncMock() + + find_first = AsyncMock(return_value=MagicMock(token=hash_token(token=api_key))) + prisma_client = MagicMock() + prisma_client.db.litellm_verificationtoken.find_first = find_first + + result = await InvalidVirtualKeyCache.check_invalid_token( + api_key=api_key, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + general_settings=_general_settings_positive_ttl(), + ) + + assert result is True + find_first.assert_not_called() + user_api_key_cache.async_set_cache.assert_not_called() + + +@pytest.mark.asyncio +async def test_check_invalid_token_cache_miss_db_hit_allows_auth_flow(): + """Negative cache empty, Prisma finds a verification row → preflight passes (False).""" + api_key = _sk_key() + hashed = hash_token(token=api_key) + neg_key = _negative_cache_key_for(api_key) + + user_api_key_cache = MagicMock() + user_api_key_cache.async_get_cache = AsyncMock(return_value=None) + user_api_key_cache.async_set_cache = AsyncMock() + + find_first = AsyncMock(return_value=MagicMock(token=hashed)) + prisma_client = MagicMock() + prisma_client.db.litellm_verificationtoken.find_first = find_first + + result = await InvalidVirtualKeyCache.check_invalid_token( + api_key=api_key, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + general_settings=_general_settings_positive_ttl(), + ) + + assert result is False + find_first.assert_awaited_once_with(where={"token": hashed}) + user_api_key_cache.async_get_cache.assert_awaited_once_with(key=neg_key) + user_api_key_cache.async_set_cache.assert_not_called()