mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-02 07:13:02 +00:00
feat(proxy): add POST /key/revoke so anyone who finds a published key can turn it off
Anyone who finds an accidentally published virtual key can now turn it off without an admin. The route takes one key or up to 100, blocks each live one and drops it from every worker's cache, and answers every valid request with the same empty 202. Each source address gets 60 requests an hour, counted in Redis when it is configured so all workers share one count The OpenAPI docs gate now checks the schema of a route's lowest 2xx response instead of only 200, so a route whose success code is 202 can pass it
This commit is contained in:
parent
0f4d5c6ad0
commit
e8354313aa
16 changed files with 1103 additions and 6 deletions
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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: ...
|
||||
|
|
|
|||
|
|
@ -753,6 +753,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/public/skill_hub",
|
||||
"/public/litellm_model_cost_map",
|
||||
"/moyai/connect/exchange",
|
||||
"/key/revoke",
|
||||
)
|
||||
)
|
||||
|
||||
|
|
|
|||
111
litellm/proxy/auth/key_revoke_throttle.py
Normal file
111
litellm/proxy/auth/key_revoke_throttle.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
]
|
||||
|
|
|
|||
153
litellm/proxy/management_endpoints/key_revocation_endpoints.py
Normal file
153
litellm/proxy/management_endpoints/key_revocation_endpoints.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
)
|
||||
|
|
|
|||
166
tests/integration/management/test_key_revoke.py
Normal file
166
tests/integration/management/test_key_revoke.py
Normal file
|
|
@ -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
|
||||
138
tests/unit/proxy/auth/test_key_revoke_throttle.py
Normal file
138
tests/unit/proxy/auth/test_key_revoke_throttle.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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) == ()
|
||||
|
|
|
|||
67
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
67
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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<string, never>;
|
||||
/** 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;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue