diff --git a/litellm/constants.py b/litellm/constants.py index cd7c0176428..76a5484a7b1 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1851,6 +1851,11 @@ LOGIN_THROTTLE_UNKNOWN_SOURCE: Final = "unknown" LOGIN_THROTTLE_MAX_TRACKED_COUNTERS: Final = 20_000 LOGIN_THROTTLE_MAX_TRACKED_BLOCKS: Final = 10_000 LOGIN_THROTTLE_NOT_BLOCKED: Final = (0, 0) +KEY_REVOKE_MAX_KEYS_PER_REQUEST: Final = 100 +KEY_REVOKE_MAX_REQUESTS_PER_SOURCE: Final = 60 +KEY_REVOKE_WINDOW_SECONDS: Final = 3600 +KEY_REVOKE_THROTTLE_CACHE_KEY_PREFIX: Final = "key_revoke" +KEY_REVOKE_THROTTLE_MAX_TRACKED_SOURCES: Final = 20_000 LITELLM_PROXY_ADMIN_NAME: Final = "default_user_id" LITELLM_PROXY_BUDGET_NAME: Final = "litellm-proxy-budget" GLOBAL_PROXY_SPEND_CACHE_KEY: Final = f"{LITELLM_PROXY_ADMIN_NAME}:spend" diff --git a/litellm/models/verification_token.py b/litellm/models/verification_token.py index 89048c56a9f..98dabab752d 100644 --- a/litellm/models/verification_token.py +++ b/litellm/models/verification_token.py @@ -5,7 +5,9 @@ Canonical definition for ``litellm_verificationtoken``. Re-exported from ``litellm.proxy._types`` for backwards compatibility. """ +from collections.abc import Iterable from datetime import datetime +from typing import Protocol from pydantic import ConfigDict @@ -77,3 +79,9 @@ class LiteLLM_DeletedVerificationToken(LiteLLM_VerificationToken): litellm_changed_by: str | None = None model_config = ConfigDict(protected_namespaces=()) + + +class KeyRevocationStore(Protocol): + async def find_by_tokens(self, tokens: Iterable[str]) -> tuple[LiteLLM_VerificationToken, ...]: ... + + async def block_tokens(self, tokens: Iterable[str]) -> None: ... diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c9960d794c6..259b137d5e7 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -753,6 +753,7 @@ class LiteLLMRoutes(enum.Enum): "/public/skill_hub", "/public/litellm_model_cost_map", "/moyai/connect/exchange", + "/key/revoke", ) ) diff --git a/litellm/proxy/auth/key_revoke_throttle.py b/litellm/proxy/auth/key_revoke_throttle.py new file mode 100644 index 00000000000..3c5e319f9c8 --- /dev/null +++ b/litellm/proxy/auth/key_revoke_throttle.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass +from functools import cache +from typing import Final, Protocol, TypeAlias + +from fastapi import Request +from pydantic import TypeAdapter, ValidationError +from redis.exceptions import RedisError + +from litellm._internal_context import with_service_target +from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.constants import ( + KEY_REVOKE_MAX_REQUESTS_PER_SOURCE, + KEY_REVOKE_THROTTLE_CACHE_KEY_PREFIX, + KEY_REVOKE_THROTTLE_MAX_TRACKED_SOURCES, + KEY_REVOKE_WINDOW_SECONDS, +) +from litellm.proxy.auth.login_throttle import TRUSTED_PROXY_RANGES_KEY, declared_proxy_ranges, source_group +from litellm.proxy.auth.network import TrustedProxyConfig, resolve_client_ip + +_UNKNOWN_SOURCE: Final = "unknown" +_REDIS_FAILURES: Final = (RedisError, RedisCircuitBreakerOpenError, OSError, asyncio.TimeoutError, ValidationError) + +_COUNT_LUA: Final = ( + "local count = redis.call('INCR', KEYS[1]) " + "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[1]) end " + "return {count, redis.call('TTL', KEYS[1])}" +) + +_Hit: TypeAlias = tuple[int, int] +_HIT: Final = TypeAdapter[_Hit](_Hit) + +_COUNTERS: Final = InMemoryCache( + max_size_in_memory=KEY_REVOKE_THROTTLE_MAX_TRACKED_SOURCES, default_ttl=KEY_REVOKE_WINDOW_SECONDS +) + + +class CounterStore(Protocol): + def increment_cache(self, key: str, value: float, *, ttl: int) -> float: ... + + +class ScriptRunner(Protocol): + def async_register_script(self, script: str) -> Callable[..., Awaitable[object]]: ... + + +@cache +def warn_key_revoke_allowance_is_per_peer() -> None: + verbose_proxy_logger.warning( + "%s is not set or not a valid list of ranges, so POST /key/revoke counts its hourly request allowance per " + "connecting address, and behind a load balancer every client shares one allowance. Set it to the address " + "ranges of the proxies in front of LiteLLM, or to an empty list when clients connect directly, to count " + "each client separately.", + TRUSTED_PROXY_RANGES_KEY, + ) + + +@dataclass(frozen=True, slots=True) +class KeyRevokeRefused: + retry_after: int + + +@dataclass(frozen=True, slots=True) +class KeyRevokeThrottle: + source: str + counters: CounterStore + redis_cache: ScriptRunner | None + + @classmethod + def from_request( + cls, + request: Request, + general_settings: Mapping[str, object] | None, + redis_cache: ScriptRunner | None, + counters: CounterStore = _COUNTERS, + ) -> KeyRevokeThrottle: + proxies: Final = None if general_settings is None else declared_proxy_ranges(general_settings) + if proxies is None: + warn_key_revoke_allowance_is_per_peer() + client_ip, _ = resolve_client_ip( + request, TrustedProxyConfig(use_forwarded_for=bool(proxies), trusted_proxy_cidrs=proxies or ()) + ) + return cls(source=source_group(client_ip or _UNKNOWN_SOURCE), counters=counters, redis_cache=redis_cache) + + async def check(self) -> KeyRevokeRefused | None: + count, retry_after = await self._count() + if count <= KEY_REVOKE_MAX_REQUESTS_PER_SOURCE: + return None + return KeyRevokeRefused(retry_after=retry_after) + + @with_service_target("key_revoke_throttle") + async def _count(self) -> _Hit: + key: Final = f"{KEY_REVOKE_THROTTLE_CACHE_KEY_PREFIX}:{self.source}" + if self.redis_cache is not None: + try: + hit: Final = _HIT.validate_python( + await self.redis_cache.async_register_script(_COUNT_LUA)([key], [KEY_REVOKE_WINDOW_SECONDS]) + ) + except _REDIS_FAILURES as err: + verbose_proxy_logger.warning( + "Redis failed while counting POST /key/revoke requests; using this worker's own count until it " + "recovers: %s", + err, + ) + else: + return hit + return int(self.counters.increment_cache(key, 1, ttl=KEY_REVOKE_WINDOW_SECONDS)), KEY_REVOKE_WINDOW_SECONDS diff --git a/litellm/proxy/deps/repositories.py b/litellm/proxy/deps/repositories.py index e37ecfd582a..0e312726872 100644 --- a/litellm/proxy/deps/repositories.py +++ b/litellm/proxy/deps/repositories.py @@ -5,10 +5,12 @@ from fastapi import Depends from litellm.models.deployment_warning import DeploymentWarningIgnoreStore from litellm.models.gateway_activity import GatewayActivityStore from litellm.models.user_banner import UserBannerStore +from litellm.models.verification_token import KeyRevocationStore from litellm.proxy.deps.database import DatabaseDep from litellm.repositories.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository from litellm.repositories.gateway_activity_repository import GatewayActivityRepository from litellm.repositories.user_banner_repository import UserBannerRepository +from litellm.repositories.verification_token_repository import VerificationTokenRepository async def get_user_banner_repository(database: DatabaseDep) -> UserBannerStore | None: @@ -34,3 +36,12 @@ async def get_gateway_activity_repository(database: DatabaseDep) -> GatewayActiv GatewayActivityRepositoryDep: TypeAlias = Annotated[ GatewayActivityStore | None, Depends(get_gateway_activity_repository) ] + + +async def get_verification_token_repository(database: DatabaseDep) -> KeyRevocationStore | None: + return None if database is None else VerificationTokenRepository(database.prisma, use_writer=True) + + +VerificationTokenRepositoryDep: TypeAlias = Annotated[ + KeyRevocationStore | None, Depends(get_verification_token_repository) +] diff --git a/litellm/proxy/management_endpoints/key_revocation_endpoints.py b/litellm/proxy/management_endpoints/key_revocation_endpoints.py new file mode 100644 index 00000000000..237f68adba0 --- /dev/null +++ b/litellm/proxy/management_endpoints/key_revocation_endpoints.py @@ -0,0 +1,153 @@ +import re +from collections.abc import Awaitable, Callable, Sequence +from datetime import datetime, timezone +from typing import Annotated, Final, TypeAlias + +from fastapi import APIRouter, Depends, Request, status +from pydantic import ConfigDict, Field + +from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid4 +from litellm.constants import KEY_REVOKE_MAX_KEYS_PER_REQUEST +from litellm.models.verification_token import KeyRevocationStore, LiteLLM_VerificationToken +from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames, ProxyErrorTypes, ProxyException +from litellm.proxy.auth.auth_checks import delete_cache_key_objects +from litellm.proxy.auth.key_revoke_throttle import KeyRevokeThrottle +from litellm.proxy.deps.repositories import VerificationTokenRepositoryDep +from litellm.proxy.management_helpers.audit_logs import create_audit_logs_for_updates, is_audit_logging_enabled +from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.proxy.auth.user_api_key_auth import hash_token + +router: Final = APIRouter() + +PUBLIC_KEY_REVOCATION_ACTOR: Final = "litellm_public_key_revocation" +_RAW_KEY: Final = re.compile(r"sk-[A-Za-z0-9_-]+") +_MAX_RAW_KEY_LENGTH: Final = 100 + +KeyRevocationAudit: TypeAlias = Callable[[Sequence[LiteLLM_AuditLogs]], Awaitable[None]] + + +class KeyRevokeRequest(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + keys: str | Annotated[tuple[str, ...], Field(min_length=1, max_length=KEY_REVOKE_MAX_KEYS_PER_REQUEST)] = Field( + description=f"One virtual key (sk-...) or a list of up to {KEY_REVOKE_MAX_KEYS_PER_REQUEST} keys to block." + ) + + +class KeyRevokeResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + +async def get_key_revoke_throttle(request: Request) -> KeyRevokeThrottle: + from litellm.proxy.proxy_server import general_settings_view, redis_usage_cache + + return KeyRevokeThrottle.from_request(request, general_settings_view(), redis_usage_cache) + + +async def get_key_revocation_audit() -> KeyRevocationAudit: + return create_audit_logs_for_updates + + +KeyRevokeThrottleDep: TypeAlias = Annotated[KeyRevokeThrottle, Depends(get_key_revoke_throttle)] +KeyRevocationAuditDep: TypeAlias = Annotated[KeyRevocationAudit, Depends(get_key_revocation_audit)] + + +def _is_raw_key(entry: str) -> bool: + return len(entry) <= _MAX_RAW_KEY_LENGTH and _RAW_KEY.fullmatch(entry) is not None + + +def revocable_hashes(keys: str | Sequence[str]) -> tuple[str, ...]: + entries: Final = (keys,) if isinstance(keys, str) else keys + return tuple(hash_token(entry) for entry in entries if _is_raw_key(entry)) + + +def _audit_row(row: LiteLLM_VerificationToken, hashed_token: str) -> LiteLLM_AuditLogs: + return LiteLLM_AuditLogs( + id=str(uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=PUBLIC_KEY_REVOCATION_ACTOR, + changed_by_api_key=None, + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id=hashed_token, + action="blocked", + updated_values="{}", + before_value=row.model_dump_json(), + ) + + +async def _block_live_keys( + tokens: KeyRevocationStore, hashed_tokens: Sequence[str], audit: KeyRevocationAudit, source: str +) -> None: + from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache + + newly_blocked: Final = tuple( + (row.token, row) for row in await tokens.find_by_tokens(hashed_tokens) if row.token and row.blocked is not True + ) + if not newly_blocked: + return + blocked_tokens: Final = tuple(token for token, _ in newly_blocked) + await tokens.block_tokens(blocked_tokens) + await delete_cache_key_objects( + hashed_tokens=blocked_tokens, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj + ) + if is_audit_logging_enabled(): + await audit(tuple(_audit_row(row, token) for token, row in newly_blocked)) + verbose_proxy_logger.warning( + "Blocked %d virtual key(s) reported through POST /key/revoke from source %s; hashed tokens: %s", + len(blocked_tokens), + source, + ", ".join(blocked_tokens), + ) + + +def _too_many_requests(retry_after: int) -> ProxyException: + return ProxyException( + message="Too many revocation requests from this address. Retry later.", + type="rate_limit_error", + param=None, + code=status.HTTP_429_TOO_MANY_REQUESTS, + headers={"Retry-After": str(retry_after)}, + ) + + +def _revocation_failed() -> ProxyException: + return ProxyException( + message="Could not process the revocation request. Retry it.", + type=ProxyErrorTypes.internal_server_error, + param=None, + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + + +@router.post( + "/key/revoke", + tags=["key management"], + status_code=status.HTTP_202_ACCEPTED, + response_model=KeyRevokeResponse, +) +async def revoke_keys( + data: KeyRevokeRequest, + throttle: KeyRevokeThrottleDep, + tokens: VerificationTokenRepositoryDep, + audit: KeyRevocationAuditDep, +) -> KeyRevokeResponse: + """ + Block virtual keys that were published by accident, without an admin. + + Anyone holding a key can report it, so this route takes no authentication. Every request that passes + validation gets 202 with an empty body, whether or not any submitted key exists. Each source address can + send 60 requests an hour. + """ + refused: Final = await throttle.check() + if refused is not None: + raise _too_many_requests(refused.retry_after) + hashed_tokens: Final = revocable_hashes(data.keys) + if tokens is None or not hashed_tokens: + return KeyRevokeResponse() + try: + await _block_live_keys(tokens, hashed_tokens, audit, throttle.source) + except Exception as err: # noqa: BLE001 # every storage or cache failure gets the same retryable answer + verbose_proxy_logger.exception("POST /key/revoke could not block the submitted keys: %s", err) + raise _revocation_failed() from err + return KeyRevokeResponse() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c5fcbfb2564..22a94083e5b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -746,6 +746,9 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( from litellm.proxy.management_endpoints.key_management_endpoints import ( router as key_management_router, ) +from litellm.proxy.management_endpoints.key_revocation_endpoints import ( + router as key_revocation_router, +) from litellm.proxy.management_endpoints.management_v1 import ( router as management_v1_router, ) @@ -21849,6 +21852,7 @@ app.include_router(pass_through_router) app.include_router(health_router) app.include_router(deployment_warning_router) app.include_router(key_management_router) +app.include_router(key_revocation_router) app.include_router(internal_user_router) app.include_router(daily_activity_router) app.include_router(password_management_router) diff --git a/litellm/repositories/verification_token_repository.py b/litellm/repositories/verification_token_repository.py index b20fb47306e..5025f9dad53 100644 --- a/litellm/repositories/verification_token_repository.py +++ b/litellm/repositories/verification_token_repository.py @@ -3,7 +3,7 @@ VerificationToken repository for database operations on LiteLLM_VerificationToke """ import json -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from datetime import datetime from types import TracebackType from typing import TYPE_CHECKING, Final, Protocol @@ -17,6 +17,7 @@ from litellm.repositories.base_repository import ( DbRecord, record_to_dict, ) +from litellm.repositories.chunked_in import update_many_in from litellm.repositories.prisma_protocols import TableActions if TYPE_CHECKING: @@ -57,6 +58,9 @@ class _PrismaClientView(Protocol): @property def db(self) -> _PrismaVerificationTokenDb: ... + @property + def writer_db(self) -> _PrismaVerificationTokenDb: ... + _JSON_ENCODED_TOKEN_FIELDS: Final = ( "aliases", @@ -74,10 +78,14 @@ _JSON_ENCODED_TOKEN_FIELDS: Final = ( class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): """Repository for verification token (API key) database operations.""" + def __init__(self, prisma_client: object, *, use_writer: bool = False) -> None: + super().__init__(prisma_client) + self._use_writer: Final = use_writer + @property def _db(self) -> _PrismaVerificationTokenDb: client: Final[_PrismaClientView] = self.prisma_client - return client.db + return client.writer_db if self._use_writer else client.db @property def table(self) -> TableActions["PrismaVerificationToken"]: @@ -437,6 +445,20 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): data["updated_by"] = updated_by return await self.update(token, data, id_field="token") + async def find_by_tokens(self, tokens: Iterable[str]) -> tuple[LiteLLM_VerificationToken, ...]: + return tuple(await self.find_many_in("token", tokens)) + + async def block_tokens(self, tokens: Iterable[str]) -> None: + from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at + + await update_many_in( + self.table, + "token", + tokens, + data=with_settings_updated_at({"blocked": True}), + atomicity="per_chunk_ok", + ) + async def unblock_token(self, token: str, updated_by: str | None = None) -> LiteLLM_VerificationToken | None: """Unblock a token.""" data: Final[dict[str, object]] = {"blocked": False} diff --git a/scripts/check_openapi_docs.py b/scripts/check_openapi_docs.py index 26bb034bfa6..66d0be64d30 100755 --- a/scripts/check_openapi_docs.py +++ b/scripts/check_openapi_docs.py @@ -7,7 +7,8 @@ count them per rule the same delta-vs-base way as the other lint gates. The rule OA001 no request body schema (the handler reads ``request.body()`` and never told FastAPI what is in it, see ``inline_request_body`` in custom_openapi_spec.py for the fix) OA002 no description (the handler has no docstring and the route no ``description=``) -OA003 the 200 response has no schema (no ``response_model`` and no ``responses={200: ...}``) +OA003 the success response (the lowest 2xx code) has no schema (no ``response_model`` and no + ``responses={200: ...}``) """ from __future__ import annotations @@ -35,13 +36,17 @@ def _json_schema(container: object) -> object: return _as_mapping(content.get("application/json")).get("schema") +def _success_response(responses: Mapping[str, object]) -> object: + return next((responses[code] for code in sorted(responses) if code.startswith("2")), None) + + def _operation_codes(operation: Mapping[str, object]) -> Iterator[str]: if not _json_schema(operation.get("requestBody")): yield "OA001" description: Final = operation.get("description") if not isinstance(description, str) or not description.strip(): yield "OA002" - if not _json_schema(_as_mapping(operation.get("responses")).get("200")): + if not _json_schema(_success_response(_as_mapping(operation.get("responses")))): yield "OA003" diff --git a/scripts/openapi_docs_gate.py b/scripts/openapi_docs_gate.py index de340bf84e7..845fb8df077 100755 --- a/scripts/openapi_docs_gate.py +++ b/scripts/openapi_docs_gate.py @@ -5,7 +5,7 @@ Sibling of scripts/test_quality_gate.py, pointed at the proxy's OpenAPI document instead of a source tree. Each rule is counted over every POST operation the proxy serves at HEAD and at the merge-base with the branch this change merges into, and the gate fails only when a rule grew past the merge-base count, so a new endpoint has to -ship with a request schema, a description and a 200 schema while the undocumented +ship with a request schema, a description and a success response schema while the undocumented routes that already exist are only ever allowed to shrink. The merge-base counts come from scripts/lint_base_counts.py: the disk cache, then the @@ -167,7 +167,8 @@ def cmd_check(base: str) -> None: print( "Give every new POST route a request body schema (a typed body parameter, or " "`openapi_extra={'requestBody': inline_request_body(Model, example)}` when the handler reads " - "`request.body()`), a docstring, and a 200 schema (`response_model=` or `responses={200: {'model': ...}}`), " + "`request.body()`), a docstring, and a schema on its lowest 2xx response (`response_model=` or " + "`responses={200: {'model': ...}}`), " "or document an equal number of existing routes; the ceiling is the merge-base count. Run " "`python scripts/check_openapi_docs.py` to see every finding." ) diff --git a/tests/integration/management/test_key_revoke.py b/tests/integration/management/test_key_revoke.py new file mode 100644 index 00000000000..f47289a2355 --- /dev/null +++ b/tests/integration/management/test_key_revoke.py @@ -0,0 +1,166 @@ +import json +import uuid +from collections.abc import Iterator, Mapping +from pathlib import Path +from typing import Final + +import httpx +import pytest +from pydantic import JsonValue + +from litellm.constants import KEY_REVOKE_MAX_REQUESTS_PER_SOURCE, KEY_REVOKE_WINDOW_SECONDS +from litellm.proxy.utils import hash_token +from tests.integration._support.client import eventually, gateway_from_environment, string_value +from tests.integration._support.process import OwnedProxy, OwnedState, WorkerHealthcheck, owned_proxy_process + +MODEL: Final = "integration-key-revoke" +_FRESH: Final = httpx.Limits(max_keepalive_connections=0) + + +def _config(directory: Path, upstream_url: str) -> Path: + config: Final = directory / "key_revoke.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": MODEL, + "litellm_params": { + "model": f"openai/{MODEL}", + "api_base": f"{upstream_url}/v1", + "api_key": "integration-provider-key", + }, + } + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "trusted_proxy_ranges": ["127.0.0.1/32", "::1/128"], + }, + } + ) + ) + return config + + +@pytest.fixture(scope="module") +def revoking(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + directory: Final = tmp_path_factory.mktemp("key-revoke").resolve() + with ( + gateway_from_environment() as shared, + owned_proxy_process( + shared, + directory, + {"STORE_MODEL_IN_DB": "False"}, + config=_config(directory, shared.upstream_url), + workers=2, + healthcheck=WorkerHealthcheck.LENIENT, + state=OwnedState.PRIVATE, + ) as owned, + ): + yield owned + + +def _source() -> str: + return f"198.51.{uuid.uuid4().int % 250}.{uuid.uuid4().int % 250 + 1}" + + +def _post(owned: OwnedProxy, path: str, body: Mapping[str, JsonValue], headers: Mapping[str, str]) -> httpx.Response: + with httpx.Client(base_url=owned.gateway.client.base_url, timeout=30, trust_env=False, limits=_FRESH) as client: + return client.post(path, json=body, headers=dict(headers)) + + +def _revoke(owned: OwnedProxy, keys: JsonValue, source: str) -> httpx.Response: + return _post(owned, "/key/revoke", {"keys": keys}, {"X-Forwarded-For": source}) + + +def _chat(owned: OwnedProxy, key: str) -> httpx.Response: + return _post( + owned, + "/v1/chat/completions", + {"model": MODEL, "messages": [{"role": "user", "content": "key revoke"}]}, + {"Authorization": f"Bearer {key}"}, + ) + + +def _workers(owned: OwnedProxy) -> tuple[int, int]: + first, second = sorted(eventually(owned.worker_pids, lambda pids: len(pids) == 2)) + return first, second + + +def _chat_on_each_worker(owned: OwnedProxy, key: str) -> tuple[httpx.Response, ...]: + def on(idle: int) -> httpx.Response: + with owned.suspended(idle): + return _chat(owned, key) + + return tuple(on(idle) for idle in reversed(_workers(owned))) + + +def _key(owned: OwnedProxy) -> str: + return string_value(owned.gateway.post("/key/generate", {"models": [MODEL]})["key"]) + + +def _blocked(owned: OwnedProxy, key: str) -> bool: + rows: Final = owned.gateway.rows( + 'SELECT blocked FROM "LiteLLM_VerificationToken" WHERE token = %s', (hash_token(key),) + ) + assert len(rows) == 1, rows + return rows[0]["blocked"] is True + + +def test_a_revoked_key_is_refused_by_every_worker_that_had_it_cached(revoking: OwnedProxy) -> None: + key: Final = _key(revoking) + assert [response.status_code for response in _chat_on_each_worker(revoking, key)] == [200, 200] + + response: Final = _revoke(revoking, key, _source()) + + assert (response.status_code, response.content) == (202, b"{}") + refused: Final = eventually( + lambda: _chat_on_each_worker(revoking, key), + lambda answers: [answer.status_code for answer in answers] == [401, 401], + return_last_on_timeout=True, + ) + assert [answer.status_code for answer in refused] == [401, 401], [answer.text for answer in refused] + assert all("Key is blocked" in answer.text for answer in refused), [answer.text for answer in refused] + assert _blocked(revoking, key) + + +def test_every_key_state_gets_the_same_answer_and_only_raw_keys_block(revoking: OwnedProxy) -> None: + live: Final = _key(revoking) + already_blocked: Final = _key(revoking) + revoking.gateway.post("/key/block", {"key": already_blocked}) + hashed_only: Final = _key(revoking) + mixed_live: Final = _key(revoking) + submissions: Final[tuple[JsonValue, ...]] = ( + live, + already_blocked, + f"sk-never-existed-{uuid.uuid4().hex}", + hash_token(hashed_only), + "not a key", + [hash_token(hashed_only), "junk", f"sk-never-{uuid.uuid4().hex}", mixed_live], + ) + + answers: Final = tuple(_revoke(revoking, keys, _source()) for keys in submissions) + + assert [(answer.status_code, answer.content) for answer in answers] == [(202, b"{}")] * len(submissions) + assert (_blocked(revoking, live), _blocked(revoking, mixed_live)) == (True, True) + assert not _blocked(revoking, hashed_only) + assert _chat(revoking, hashed_only).status_code == 200 + + +def test_the_hourly_allowance_is_counted_once_across_workers(revoking: OwnedProxy) -> None: + source: Final = _source() + never_existed: Final = f"sk-never-existed-{uuid.uuid4().hex}" + half: Final = KEY_REVOKE_MAX_REQUESTS_PER_SOURCE // 2 + + def on(idle: int) -> tuple[int, ...]: + with revoking.suspended(idle): + return tuple(_revoke(revoking, never_existed, source).status_code for _ in range(half)) + + admitted: Final = tuple(status for idle in reversed(_workers(revoking)) for status in on(idle)) + refused: Final = _revoke(revoking, never_existed, source) + + assert admitted == (202,) * (half * 2) + assert refused.status_code == 429, refused.text + assert 0 < int(refused.headers["retry-after"]) <= KEY_REVOKE_WINDOW_SECONDS + assert _revoke(revoking, never_existed, _source()).status_code == 202 diff --git a/tests/unit/proxy/auth/test_key_revoke_throttle.py b/tests/unit/proxy/auth/test_key_revoke_throttle.py new file mode 100644 index 00000000000..25bbec769ce --- /dev/null +++ b/tests/unit/proxy/auth/test_key_revoke_throttle.py @@ -0,0 +1,138 @@ +import asyncio +from collections.abc import Awaitable, Callable, Mapping, Sequence +from typing import Final + +import pytest +from redis.exceptions import ConnectionError as RedisConnectionError +from redis.exceptions import TimeoutError as RedisTimeoutError +from starlette.requests import Request + +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.constants import KEY_REVOKE_MAX_REQUESTS_PER_SOURCE, KEY_REVOKE_WINDOW_SECONDS +from litellm.proxy.auth.key_revoke_throttle import KeyRevokeRefused, KeyRevokeThrottle + +pytestmark = pytest.mark.asyncio + + +def _request(peer: str, forwarded_for: str | None = None) -> Request: + headers: Final = () if forwarded_for is None else ((b"x-forwarded-for", forwarded_for.encode()),) + return Request({"type": "http", "client": (peer, 40000), "headers": list(headers)}) + + +class _CountingRedis: + def __init__(self) -> None: + self.counts: dict[str, int] = {} + self.ttls: dict[str, int] = {} + + def async_register_script(self, script: str) -> Callable[[Sequence[str], Sequence[object]], Awaitable[object]]: + async def run(keys: Sequence[str], args: Sequence[object]) -> object: + key: Final = keys[0] + self.counts[key] = self.counts.get(key, 0) + 1 + self.ttls.setdefault(key, int(str(args[0]))) + return [self.counts[key], self.ttls[key]] + + return run + + +class _FailingRedis: + def __init__(self, outcome: object) -> None: + self.outcome: Final = outcome + + def async_register_script(self, script: str) -> Callable[[Sequence[str], Sequence[object]], Awaitable[object]]: + async def run(keys: Sequence[str], args: Sequence[object]) -> object: + if isinstance(self.outcome, BaseException): + raise self.outcome + return self.outcome + + return run + + +def _throttle( + request: Request, + settings: Mapping[str, object] | None = None, + redis: _CountingRedis | _FailingRedis | None = None, + counters: InMemoryCache | None = None, +) -> KeyRevokeThrottle: + return KeyRevokeThrottle.from_request( + request, + general_settings=settings, + redis_cache=redis, + counters=counters if counters is not None else InMemoryCache(), + ) + + +async def _spend_allowance(throttle: KeyRevokeThrottle) -> tuple[KeyRevokeRefused | None, ...]: + return tuple([await throttle.check() for _ in range(KEY_REVOKE_MAX_REQUESTS_PER_SOURCE)]) + + +async def test_the_request_past_the_hourly_allowance_is_refused_for_the_whole_window(): + throttle: Final = _throttle(_request("203.0.113.7")) + + assert await _spend_allowance(throttle) == (None,) * KEY_REVOKE_MAX_REQUESTS_PER_SOURCE + assert await throttle.check() == KeyRevokeRefused(retry_after=KEY_REVOKE_WINDOW_SECONDS) + + +async def test_redis_shares_one_count_across_workers_and_sets_the_retry_time(): + redis: Final = _CountingRedis() + first_worker: Final = _throttle(_request("203.0.113.7"), redis=redis) + second_worker: Final = _throttle(_request("203.0.113.7"), redis=redis) + half: Final = KEY_REVOKE_MAX_REQUESTS_PER_SOURCE // 2 + + admitted: Final = tuple([await worker.check() for worker in (first_worker, second_worker) for _ in range(half)]) + assert admitted == (None,) * (half * 2) + assert await second_worker.check() == KeyRevokeRefused(retry_after=KEY_REVOKE_WINDOW_SECONDS) + (key,) = redis.ttls + redis.ttls[key] = 1234 + assert await first_worker.check() == KeyRevokeRefused(retry_after=1234) + + +@pytest.mark.parametrize( + "outcome", + [ + RedisConnectionError("redis is down"), + RedisTimeoutError("redis timed out"), + RedisCircuitBreakerOpenError("circuit open"), + ConnectionRefusedError("refused"), + asyncio.TimeoutError(), + ["not", "a count"], + ], + ids=["connection", "timeout", "circuit-open", "socket", "asyncio-timeout", "malformed-reply"], +) +async def test_a_redis_failure_falls_back_to_this_workers_own_count(outcome): + throttle: Final = _throttle(_request("203.0.113.7"), redis=_FailingRedis(outcome)) + + assert await _spend_allowance(throttle) == (None,) * KEY_REVOKE_MAX_REQUESTS_PER_SOURCE + assert await throttle.check() == KeyRevokeRefused(retry_after=KEY_REVOKE_WINDOW_SECONDS) + + +async def test_without_trusted_proxy_ranges_clients_behind_one_peer_share_a_bucket(): + counters: Final = InMemoryCache() + first_client: Final = _throttle(_request("10.0.0.1", "198.51.100.1"), counters=counters) + second_client: Final = _throttle(_request("10.0.0.1", "198.51.100.2"), counters=counters) + + assert await _spend_allowance(first_client) == (None,) * KEY_REVOKE_MAX_REQUESTS_PER_SOURCE + assert await second_client.check() == KeyRevokeRefused(retry_after=KEY_REVOKE_WINDOW_SECONDS) + + +async def test_with_trusted_proxy_ranges_each_forwarded_client_has_its_own_bucket(): + counters: Final = InMemoryCache() + settings: Final = {"trusted_proxy_ranges": ["10.0.0.0/8"]} + first_client: Final = _throttle(_request("10.0.0.1", "198.51.100.1"), settings, counters=counters) + second_client: Final = _throttle(_request("10.0.0.1", "198.51.100.2"), settings, counters=counters) + + assert await _spend_allowance(first_client) == (None,) * KEY_REVOKE_MAX_REQUESTS_PER_SOURCE + assert await second_client.check() is None + assert await first_client.check() == KeyRevokeRefused(retry_after=KEY_REVOKE_WINDOW_SECONDS) + + +async def test_ipv6_addresses_in_one_64_share_a_bucket(): + counters: Final = InMemoryCache() + settings: Final = {"trusted_proxy_ranges": []} + first_address: Final = _throttle(_request("2001:db8:1:2::1"), settings, counters=counters) + same_prefix: Final = _throttle(_request("2001:db8:1:2:ffff:ffff:ffff:ffff"), settings, counters=counters) + other_prefix: Final = _throttle(_request("2001:db8:1:3::1"), settings, counters=counters) + + assert await _spend_allowance(first_address) == (None,) * KEY_REVOKE_MAX_REQUESTS_PER_SOURCE + assert await same_prefix.check() == KeyRevokeRefused(retry_after=KEY_REVOKE_WINDOW_SECONDS) + assert await other_prefix.check() is None diff --git a/tests/unit/proxy/deps/test_repositories.py b/tests/unit/proxy/deps/test_repositories.py index e802a4fda61..19a4e13edbd 100644 --- a/tests/unit/proxy/deps/test_repositories.py +++ b/tests/unit/proxy/deps/test_repositories.py @@ -15,6 +15,7 @@ from litellm.proxy.deps.repositories import ( DeploymentWarningIgnoreRepositoryDep, GatewayActivityRepositoryDep, UserBannerRepositoryDep, + VerificationTokenRepositoryDep, ) from litellm.repositories import database as database_runtime from tests._support.native_database import native_database @@ -41,6 +42,10 @@ def _app(database: database_runtime.Database | None) -> FastAPI: async def activity(activity: GatewayActivityRepositoryDep) -> tuple[GatewayActivity, ...] | None: return None if activity is None else await activity.get_activity("2026-01-01", "2026-01-31") + @app.get("/revocable") + async def revocable(tokens: VerificationTokenRepositoryDep) -> tuple[str | None, ...] | None: + return None if tokens is None else tuple(row.token for row in await tokens.find_by_tokens(["hashed"])) + return app @@ -105,6 +110,16 @@ def test_gateway_activity_reads_the_lifespan_reader_not_the_writer() -> None: assert response.json() == [row] +def test_key_revocation_looks_keys_up_on_the_writer_so_a_lagging_replica_cannot_hide_one() -> None: + prisma: Final = MagicMock() + prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + prisma.writer_db.litellm_verificationtoken.find_many = AsyncMock(return_value=[{"token": "hashed"}]) + with TestClient(_app(database_runtime.Database(prisma=prisma, native=NATIVE))) as client: + response: Final = client.get("/revocable") + assert response.status_code == 200 + assert response.json() == ["hashed"] + + def test_gateway_activity_without_database_is_unavailable() -> None: with TestClient(_app(None)) as client: response: Final = client.get("/activity") diff --git a/tests/unit/proxy/management_endpoints/test_key_revocation_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_revocation_endpoints.py new file mode 100644 index 00000000000..accd09d72d8 --- /dev/null +++ b/tests/unit/proxy/management_endpoints/test_key_revocation_endpoints.py @@ -0,0 +1,370 @@ +import json +import logging +from collections.abc import Mapping, Sequence +from datetime import datetime +from types import SimpleNamespace +from typing import Final + +import pytest +from fastapi.testclient import TestClient + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import ( + KEY_REVOKE_MAX_KEYS_PER_REQUEST, + KEY_REVOKE_MAX_REQUESTS_PER_SOURCE, + KEY_REVOKE_WINDOW_SECONDS, +) +from litellm.proxy import proxy_server +from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames, UserAPIKeyAuth +from litellm.proxy.auth.key_revoke_throttle import KeyRevokeThrottle +from litellm.proxy.deps.database import get_database +from litellm.proxy.management_endpoints.key_revocation_endpoints import ( + get_key_revocation_audit, + get_key_revoke_throttle, +) +from litellm.proxy.proxy_server import app +from litellm.proxy.utils import hash_token +from litellm.repositories.database import Database +from tests._master_key import MASTER_KEY +from tests._support.native_database import native_database + +client: Final = TestClient(app) +NATIVE: Final = native_database() + +LIVE: Final = "sk-live-0123456789abcdefABCDEF" +OTHER: Final = "sk-other_0123456789abcdef" +BLOCKED: Final = "sk-blocked-0123456789abcdef" +NEVER_EXISTED: Final = "sk-never-existed-0123456789" +LIVE_HASH: Final = hash_token(LIVE) +OTHER_HASH: Final = hash_token(OTHER) +BLOCKED_HASH: Final = hash_token(BLOCKED) +STORAGE_ERROR: Final = "connection to db-internal-7.corp:5432 reset while reading LiteLLM_VerificationToken" +ACCEPTED: Final = (202, b"{}") + + +class _TokenTable: + def __init__(self, rows: Mapping[str, Mapping[str, object]], fail_on: str | None) -> None: + self.rows: Final = {token: dict(row) for token, row in rows.items()} + self.updates: Final[list[tuple[str, ...]]] = [] + self.fail_on: Final = fail_on + + async def find_many(self, where: Mapping[str, Mapping[str, Sequence[str]]]) -> list[dict[str, object]]: + if self.fail_on == "find": + raise RuntimeError(STORAGE_ERROR) + return [dict(self.rows[token]) for token in where["token"]["in"] if token in self.rows] + + async def update_many(self, data: Mapping[str, object], where: Mapping[str, Mapping[str, Sequence[str]]]) -> int: + if self.fail_on == "update": + raise RuntimeError(STORAGE_ERROR) + tokens: Final = tuple(where["token"]["in"]) + self.updates.append(tokens) + for token in tokens: + self.rows[token].update(data) + return len(tokens) + + def blocked(self, hashed_token: str) -> bool: + return self.rows[hashed_token]["blocked"] is True + + +def _row(hashed_token: str, blocked: bool | None = None) -> dict[str, object]: + return {"token": hashed_token, "blocked": blocked, "team_id": f"team-{hashed_token[:8]}"} + + +def _install( + monkeypatch: pytest.MonkeyPatch, rows: Mapping[str, Mapping[str, object]], fail_on: str | None = None +) -> _TokenTable: + table: Final = _TokenTable(rows, fail_on) + prisma: Final = SimpleNamespace(writer_db=SimpleNamespace(litellm_verificationtoken=table)) + monkeypatch.setitem(app.dependency_overrides, get_database, lambda: Database(prisma=prisma, native=NATIVE)) + return table + + +def _revoke(keys: object, headers: Mapping[str, str] | None = None) -> tuple[int, bytes]: + response: Final = client.post("/key/revoke", json={"keys": keys}, headers=headers) + return response.status_code, response.content + + +@pytest.fixture(autouse=True) +def throttle(monkeypatch: pytest.MonkeyPatch) -> KeyRevokeThrottle: + fresh: Final = KeyRevokeThrottle(source="203.0.113.7", counters=InMemoryCache(), redis_cache=None) + monkeypatch.setitem(app.dependency_overrides, get_key_revoke_throttle, lambda: fresh) + return fresh + + +@pytest.fixture(autouse=True) +def audit_rows(monkeypatch: pytest.MonkeyPatch) -> list[LiteLLM_AuditLogs]: + recorded: Final[list[LiteLLM_AuditLogs]] = [] + + async def record(rows: Sequence[LiteLLM_AuditLogs]) -> None: + recorded.extend(rows) + + monkeypatch.setitem(app.dependency_overrides, get_key_revocation_audit, lambda: record) + monkeypatch.setattr(litellm, "store_audit_logs", True) + return recorded + + +def test_a_single_live_key_is_blocked_and_answered_with_an_empty_202(monkeypatch): + table: Final = _install(monkeypatch, {LIVE_HASH: _row(LIVE_HASH), OTHER_HASH: _row(OTHER_HASH)}) + + assert _revoke(LIVE) == ACCEPTED + assert table.blocked(LIVE_HASH) + assert isinstance(table.rows[LIVE_HASH]["settings_updated_at"], datetime) + assert not table.blocked(OTHER_HASH) + + +def test_the_log_names_the_hashed_token_and_source_but_never_the_raw_key(monkeypatch, caplog): + _install(monkeypatch, {LIVE_HASH: _row(LIVE_HASH)}) + verbose_proxy_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): + assert _revoke([LIVE, NEVER_EXISTED]) == ACCEPTED + finally: + verbose_proxy_logger.removeHandler(caplog.handler) + + messages: Final = tuple(record.getMessage() for record in caplog.records) + assert any(LIVE_HASH in message and "203.0.113.7" in message for message in messages), messages + assert not any(LIVE in message or NEVER_EXISTED in message for message in messages), messages + + +def test_every_key_in_a_full_batch_is_blocked(monkeypatch): + keys: Final = tuple(f"sk-batch-{i:03d}" for i in range(KEY_REVOKE_MAX_KEYS_PER_REQUEST)) + table: Final = _install(monkeypatch, {hash_token(key): _row(hash_token(key)) for key in keys}) + + assert _revoke(list(keys)) == ACCEPTED + assert [table.blocked(hash_token(key)) for key in keys] == [True] * KEY_REVOKE_MAX_KEYS_PER_REQUEST + + +def test_revoking_evicts_the_cached_key_object(monkeypatch): + _install(monkeypatch, {LIVE_HASH: _row(LIVE_HASH)}) + cache: Final = proxy_server.user_api_key_cache + cache.set_cache(key=LIVE_HASH, value=UserAPIKeyAuth(token=LIVE_HASH).model_dump()) + assert cache.get_cache(key=LIVE_HASH) is not None + + assert _revoke(LIVE) == ACCEPTED + assert cache.get_cache(key=LIVE_HASH) is None + + +def test_duplicate_entries_are_blocked_with_one_update_and_one_audit_row(monkeypatch, audit_rows): + table: Final = _install(monkeypatch, {LIVE_HASH: _row(LIVE_HASH)}) + + assert _revoke([LIVE, LIVE, LIVE]) == ACCEPTED + assert table.updates == [(LIVE_HASH,)] + assert [row.object_id for row in audit_rows] == [LIVE_HASH] + + +@pytest.mark.parametrize( + "body", + [ + {"keys": []}, + {"keys": [f"sk-batch-{i:03d}" for i in range(KEY_REVOKE_MAX_KEYS_PER_REQUEST + 1)]}, + {"keys": [LIVE, 5]}, + {"keys": 5}, + {"keys": None}, + {}, + ], + ids=["empty-list", "over-the-cap", "non-string-entry", "not-a-string-or-list", "null", "missing"], +) +def test_schema_violations_are_rejected_without_spending_the_allowance(monkeypatch, body): + table: Final = _install(monkeypatch, {LIVE_HASH: _row(LIVE_HASH)}) + + statuses: Final = tuple( + client.post("/key/revoke", json=body).status_code for _ in range(KEY_REVOKE_MAX_REQUESTS_PER_SOURCE) + ) + assert statuses == (422,) * KEY_REVOKE_MAX_REQUESTS_PER_SOURCE + assert not table.blocked(LIVE_HASH) + assert _revoke(LIVE) == ACCEPTED + assert table.blocked(LIVE_HASH) + + +def test_a_hashed_token_never_blocks_its_key(monkeypatch): + table: Final = _install(monkeypatch, {LIVE_HASH: _row(LIVE_HASH), OTHER_HASH: _row(OTHER_HASH)}) + + assert _revoke([LIVE_HASH, OTHER]) == ACCEPTED + assert not table.blocked(LIVE_HASH) + assert table.blocked(OTHER_HASH) + + +@pytest.mark.parametrize( + "junk", + [ + "sk-" + "a" * 98, + "sk-has space", + "sk-trailing-newline\n", + "SK-UPPERCASE-PREFIX", + "sk-", + "sk-semi;colon", + "not-a-key", + "", + ], + ids=[ + "101-chars", + "space", + "trailing-newline", + "uppercase-prefix", + "bare-prefix", + "semicolon", + "no-prefix", + "empty", + ], +) +def test_a_junk_entry_blocks_nothing_even_when_a_row_has_its_hash(monkeypatch, junk): + table: Final = _install(monkeypatch, {hash_token(junk): _row(hash_token(junk))}) + + assert _revoke(junk) == ACCEPTED + assert not table.blocked(hash_token(junk)) + + +@pytest.mark.parametrize("key", ["sk-" + "a" * 97, "sk-under_score-and-dash"], ids=["100-chars", "underscore-dash"]) +def test_every_raw_key_shape_up_to_100_chars_is_revocable(monkeypatch, key): + table: Final = _install(monkeypatch, {hash_token(key): _row(hash_token(key))}) + + assert _revoke(key) == ACCEPTED + assert table.blocked(hash_token(key)) + + +def test_an_already_blocked_key_gets_no_second_update_or_audit_row(monkeypatch, audit_rows): + table: Final = _install(monkeypatch, {BLOCKED_HASH: _row(BLOCKED_HASH, blocked=True)}) + + assert _revoke(BLOCKED) == ACCEPTED + assert table.blocked(BLOCKED_HASH) + assert table.updates == [] + assert audit_rows == [] + + +@pytest.mark.parametrize( + "authorization", + [f"Bearer {MASTER_KEY}", "Bearer sk-not-a-real-key", "Basic Zm9vOmJhcg==", "garbage"], + ids=["master-key", "unknown-key", "basic", "garbage"], +) +def test_an_authorization_header_changes_nothing(monkeypatch, authorization): + monkeypatch.setattr(proxy_server, "master_key", MASTER_KEY) + table: Final = _install(monkeypatch, {LIVE_HASH: _row(LIVE_HASH)}) + + assert _revoke(LIVE, headers={"Authorization": authorization}) == ACCEPTED + assert table.blocked(LIVE_HASH) + + +def test_without_a_database_the_request_is_accepted_and_does_nothing(monkeypatch, audit_rows): + monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None) + + assert _revoke(LIVE) == ACCEPTED + assert audit_rows == [] + + +def test_without_a_database_requests_still_spend_the_allowance(monkeypatch): + monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None) + + admitted: Final = tuple(_revoke(LIVE) for _ in range(KEY_REVOKE_MAX_REQUESTS_PER_SOURCE)) + assert admitted == (ACCEPTED,) * KEY_REVOKE_MAX_REQUESTS_PER_SOURCE + assert _revoke(LIVE)[0] == 429 + + +@pytest.mark.parametrize("fail_on", ["find", "update"]) +def test_a_storage_failure_answers_a_generic_retryable_500(monkeypatch, audit_rows, fail_on): + _install(monkeypatch, {LIVE_HASH: _row(LIVE_HASH)}, fail_on=fail_on) + + response: Final = client.post("/key/revoke", json={"keys": LIVE}) + + assert response.status_code == 500 + assert response.json() == { + "error": { + "message": "Could not process the revocation request. Retry it.", + "type": "internal_server_error", + "param": None, + "code": "500", + } + } + assert STORAGE_ERROR not in response.text + assert audit_rows == [] + + +def test_each_newly_blocked_key_gets_an_audit_row_naming_the_public_revocation(monkeypatch, audit_rows): + _install( + monkeypatch, + { + LIVE_HASH: _row(LIVE_HASH), + OTHER_HASH: _row(OTHER_HASH, blocked=False), + BLOCKED_HASH: _row(BLOCKED_HASH, blocked=True), + }, + ) + + assert _revoke([LIVE, OTHER, BLOCKED]) == ACCEPTED + assert sorted( + (row.object_id, row.action, row.changed_by, row.changed_by_api_key, row.table_name) for row in audit_rows + ) == sorted( + (hashed, "blocked", "litellm_public_key_revocation", None, LitellmTableNames.KEY_TABLE_NAME) + for hashed in (LIVE_HASH, OTHER_HASH) + ) + assert len({row.id for row in audit_rows}) == 2 + before_values: Final = tuple(json.loads(str(row.before_value)) for row in audit_rows) + assert sorted((before["team_id"], before["blocked"]) for before in before_values) == sorted( + [(f"team-{LIVE_HASH[:8]}", None), (f"team-{OTHER_HASH[:8]}", False)] + ) + + +def test_with_audit_logging_off_keys_are_blocked_without_audit_rows(monkeypatch, audit_rows): + monkeypatch.setattr(litellm, "store_audit_logs", False) + table: Final = _install(monkeypatch, {LIVE_HASH: _row(LIVE_HASH)}) + + assert _revoke(LIVE) == ACCEPTED + assert table.blocked(LIVE_HASH) + assert audit_rows == [] + + +def test_the_response_is_byte_identical_whatever_the_submitted_keys_are(monkeypatch): + table: Final = _install( + monkeypatch, + { + LIVE_HASH: _row(LIVE_HASH), + OTHER_HASH: _row(OTHER_HASH), + BLOCKED_HASH: _row(BLOCKED_HASH, blocked=True), + hash_token("sk-mixed-live-key"): _row(hash_token("sk-mixed-live-key")), + }, + ) + submissions: Final[tuple[object, ...]] = ( + LIVE, + BLOCKED, + NEVER_EXISTED, + OTHER_HASH, + "not a key at all", + [LIVE, BLOCKED, NEVER_EXISTED, OTHER_HASH, "junk", "sk-mixed-live-key"], + ) + + def answer(keys: object) -> tuple[int, bytes, str | None, str | None]: + response: Final = client.post("/key/revoke", json={"keys": keys}) + return ( + response.status_code, + response.content, + response.headers.get("content-type"), + response.headers.get("content-length"), + ) + + answers: Final = tuple(answer(keys) for keys in submissions) + + assert answers == ((202, b"{}", "application/json", "2"),) * len(submissions) + assert table.blocked(LIVE_HASH) + assert table.blocked(hash_token("sk-mixed-live-key")) + assert not table.blocked(OTHER_HASH) + + +@pytest.mark.parametrize("filler", [NEVER_EXISTED, "not a key", OTHER_HASH], ids=["unknown-key", "junk", "hash"]) +def test_past_the_hourly_allowance_a_request_is_refused_and_blocks_nothing(monkeypatch, filler): + table: Final = _install(monkeypatch, {LIVE_HASH: _row(LIVE_HASH)}) + + admitted: Final = tuple(_revoke(filler) for _ in range(KEY_REVOKE_MAX_REQUESTS_PER_SOURCE)) + refused: Final = client.post("/key/revoke", json={"keys": LIVE}) + + assert admitted == (ACCEPTED,) * KEY_REVOKE_MAX_REQUESTS_PER_SOURCE + assert refused.status_code == 429 + assert refused.headers["retry-after"] == str(KEY_REVOKE_WINDOW_SECONDS) + assert refused.json() == { + "error": { + "message": "Too many revocation requests from this address. Retry later.", + "type": "rate_limit_error", + "param": None, + "code": "429", + } + } + assert not table.blocked(LIVE_HASH) diff --git a/tests/unit/test_check_openapi_docs.py b/tests/unit/test_check_openapi_docs.py index d7122541416..2fee5927a5e 100644 --- a/tests/unit/test_check_openapi_docs.py +++ b/tests/unit/test_check_openapi_docs.py @@ -37,6 +37,26 @@ def test_only_the_missing_piece_is_reported() -> None: ) +def test_the_schema_is_judged_on_the_lowest_2xx_response() -> None: + accepted: Final = {**_DOCUMENTED, "responses": {"202": _RESPONSES["200"], "422": {"description": "Invalid"}}} + untyped_202: Final = {**_DOCUMENTED, "responses": {"202": {"description": "Accepted"}, "422": _RESPONSES["200"]}} + untyped_200: Final = {**_DOCUMENTED, "responses": {"200": {"description": "OK"}, "201": _RESPONSES["200"]}} + no_2xx: Final = {**_DOCUMENTED, "responses": {"422": _RESPONSES["200"]}} + spec: Final = { + "paths": { + "/a": {"post": accepted}, + "/b": {"post": untyped_202}, + "/c": {"post": untyped_200}, + "/d": {"post": no_2xx}, + } + } + assert checker.violations(spec) == ( + checker.Violation("POST", "/b", "OA003"), + checker.Violation("POST", "/c", "OA003"), + checker.Violation("POST", "/d", "OA003"), + ) + + def test_non_post_operations_are_not_gated() -> None: spec: Final = {"paths": {"/v1/thing": {"get": {}, "delete": {}, "put": {}, "post": _DOCUMENTED}}} assert checker.violations(spec) == () diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 96d4c0e1bc4..0183af0b5a4 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8836,6 +8836,30 @@ export interface paths { patch?: never; trace?: never; }; + "/key/revoke": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Revoke Keys + * @description Block virtual keys that were published by accident, without an admin. + * + * Anyone holding a key can report it, so this route takes no authentication. Every request that passes + * validation gets 202 with an empty body, whether or not any submitted key exists. Each source address can + * send 60 requests an hour. + */ + post: operations["revoke_keys_key_revoke_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/key/service-account/generate": { parameters: { query?: never; @@ -34233,6 +34257,16 @@ export interface components { /** Keys */ keys?: string[] | null; }; + /** KeyRevokeRequest */ + KeyRevokeRequest: { + /** + * Keys + * @description One virtual key (sk-...) or a list of up to 100 keys to block. + */ + keys: string | string[]; + }; + /** KeyRevokeResponse */ + KeyRevokeResponse: Record; /** KeySpendActivityRow */ KeySpendActivityRow: { /** Api Key */ @@ -64181,6 +64215,39 @@ export interface operations { }; }; }; + revoke_keys_key_revoke_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["KeyRevokeRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 202: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["KeyRevokeResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; generate_service_account_key_fn_key_service_account_generate_post: { parameters: { query?: never;