From 1fda6c85f30dad3ad2433068d9c1b01b5465afdb Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 14 Oct 2026 03:18:50 +0000 Subject: [PATCH] feat(control_plane): ocsf audit events, sink protocol, queued delivery and hash chain Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/control_plane/app.py | 32 ++- .../proxy/control_plane/audit/dependencies.py | 9 +- litellm/proxy/control_plane/audit/event.py | 33 +++ litellm/proxy/control_plane/audit/log.py | 127 ++++++---- litellm/proxy/control_plane/audit/ocsf.py | 226 +++++++++++++++++ litellm/proxy/control_plane/audit/pipeline.py | 228 ++++++++++++++++++ litellm/proxy/control_plane/audit/router.py | 29 ++- litellm/proxy/control_plane/audit/sinks.py | 123 ++++++++++ .../proxy/control_plane/auth/dependencies.py | 26 +- litellm/proxy/control_plane/auth/router.py | 6 + litellm/proxy/control_plane/common/tracing.py | 21 +- litellm/proxy/control_plane/container.py | 31 ++- litellm/proxy/control_plane/proxy/router.py | 78 +++++- litellm/proxy/control_plane/settings.py | 7 +- litellm/types/proxy/control_plane/audit.py | 9 + .../proxy/control_plane/audit/test_ocsf.py | 97 ++++++++ .../control_plane/audit/test_pipeline.py | 174 +++++++++++++ .../proxy/control_plane/audit/test_router.py | 57 +++++ .../proxy/control_plane/audit/test_sinks.py | 83 +++++++ .../proxy/control_plane/proxy/test_router.py | 104 +++++++- tests/unit/proxy/control_plane/test_app.py | 31 +++ 21 files changed, 1452 insertions(+), 79 deletions(-) create mode 100644 litellm/proxy/control_plane/audit/event.py create mode 100644 litellm/proxy/control_plane/audit/ocsf.py create mode 100644 litellm/proxy/control_plane/audit/pipeline.py create mode 100644 litellm/proxy/control_plane/audit/sinks.py create mode 100644 tests/unit/proxy/control_plane/audit/test_ocsf.py create mode 100644 tests/unit/proxy/control_plane/audit/test_pipeline.py create mode 100644 tests/unit/proxy/control_plane/audit/test_router.py create mode 100644 tests/unit/proxy/control_plane/audit/test_sinks.py diff --git a/litellm/proxy/control_plane/app.py b/litellm/proxy/control_plane/app.py index c92998748b2..3b0d688da99 100644 --- a/litellm/proxy/control_plane/app.py +++ b/litellm/proxy/control_plane/app.py @@ -1,7 +1,7 @@ """App factory for the LiteLLM global control plane. Enterprise only: refuses to build without a valid license.""" -from collections.abc import AsyncGenerator -from contextlib import asynccontextmanager +from dataclasses import dataclass +from types import TracebackType from typing import Final, Protocol from fastapi import FastAPI @@ -34,23 +34,35 @@ class LicenseRequired(RuntimeError): ) +@dataclass(frozen=True, slots=True) +class _Lifespan: + container: Container + + def __call__(self, app: FastAPI) -> "_Lifespan": + return self + + async def __aenter__(self) -> None: + self.container.audit_log.pipeline.start() + + async def __aexit__( + self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None + ) -> None: + try: + await self.container.audit_log.pipeline.close() + finally: + self.container.telemetry.shutdown() + + def create_app(container: Container, license_check: License) -> FastAPI: if not license_check.is_premium(): raise LicenseRequired() - @asynccontextmanager - async def lifespan(_app: FastAPI) -> AsyncGenerator[None, None]: - try: - yield - finally: - container.telemetry.shutdown() - app: Final = FastAPI( title=TITLE, version="1", + lifespan=_Lifespan(container), docs_url=f"{container.settings.api_prefix}/docs", openapi_url=f"{container.settings.api_prefix}/openapi.json", - lifespan=lifespan, ) container.telemetry.instrument(app) app.dependency_overrides[get_token_verifier] = lambda: container.token_verifier diff --git a/litellm/proxy/control_plane/audit/dependencies.py b/litellm/proxy/control_plane/audit/dependencies.py index c260c4a65b3..f47cafed8f1 100644 --- a/litellm/proxy/control_plane/audit/dependencies.py +++ b/litellm/proxy/control_plane/audit/dependencies.py @@ -1,5 +1,10 @@ -from litellm.proxy.control_plane.audit.log import AuditLog +from typing import Final + +from litellm.proxy.control_plane.audit.log import AuditLog, build_audit_log + +_UNBOUND: Final = build_audit_log(capacity=1) def get_audit_log() -> AuditLog: - raise RuntimeError("get_audit_log is not bound, build the app with create_app") + """Overridden in create_app. A throwaway ring by default so a bare router still works in tests.""" + return _UNBOUND diff --git a/litellm/proxy/control_plane/audit/event.py b/litellm/proxy/control_plane/audit/event.py new file mode 100644 index 00000000000..f2d07c9e845 --- /dev/null +++ b/litellm/proxy/control_plane/audit/event.py @@ -0,0 +1,33 @@ +"""The audit event value: what happened, by whom, on which dataplane, and which trace it belongs to.""" + +from dataclasses import dataclass +from datetime import datetime +from typing import Final, Literal, TypeAlias + +Decision: TypeAlias = Literal["allowed", "denied", "unreachable"] +EventKind: TypeAlias = Literal["api_activity", "authentication"] +KEY_PREFIX: Final = "sk-" +REDACTED_KEY: Final = "sk-[redacted]" +CONTROL_PLANE_ID: Final = "control-plane" + + +def redact_path(path: str) -> str: + """A virtual key in a path (``/key/sk-.../reset_spend``) must never reach the log or the audit API.""" + return "/".join(REDACTED_KEY if segment.startswith(KEY_PREFIX) else segment for segment in path.split("/")) + + +@dataclass(frozen=True, slots=True) +class AuditEvent: + at: datetime + actor: str + actor_email: str | None + dataplane_id: str + method: str + path: str + decision: Decision + status: int | None + credential: str | None + detail: str | None + kind: EventKind = "api_activity" + trace_id: str | None = None + span_id: str | None = None diff --git a/litellm/proxy/control_plane/audit/log.py b/litellm/proxy/control_plane/audit/log.py index a77671e17a0..b760adb9a89 100644 --- a/litellm/proxy/control_plane/audit/log.py +++ b/litellm/proxy/control_plane/audit/log.py @@ -1,47 +1,59 @@ -"""Who did what on which dataplane through the control plane. Bounded in memory plus a structured log line.""" +"""Who did what through the control plane. Every event becomes a chained OCSF record, kept in a bounded ring for the +audit API and handed to the configured sinks for delivery.""" -import json from collections import deque from collections.abc import Callable -from dataclasses import asdict, dataclass from datetime import UTC, datetime -from typing import Final, Literal, TypeAlias +from typing import Final, TypeAlias -from litellm._logging import verbose_proxy_logger +from opentelemetry.trace import format_span_id, format_trace_id, get_current_span + +from litellm.proxy.control_plane.audit.event import CONTROL_PLANE_ID, AuditEvent, Decision, EventKind, redact_path +from litellm.proxy.control_plane.audit.ocsf import AuditRecord, ChainState, Product, chain +from litellm.proxy.control_plane.audit.pipeline import NO_PIPELINE, AuditPipeline, Offer, Reservation -Decision: TypeAlias = Literal["allowed", "denied", "unreachable"] DEFAULT_CAPACITY: Final = 10_000 -KEY_PREFIX: Final = "sk-" -REDACTED_KEY: Final = "sk-[redacted]" - - -def redact_path(path: str) -> str: - """A virtual key in a path (``/key/sk-.../reset_spend``) must never reach the log or the audit API.""" - return "/".join(REDACTED_KEY if segment.startswith(KEY_PREFIX) else segment for segment in path.split("/")) - - -@dataclass(frozen=True, slots=True) -class AuditEvent: - at: datetime - actor: str - actor_email: str | None - dataplane_id: str - method: str - path: str - decision: Decision - status: int | None - credential: str | None - detail: str | None +UNKNOWN_PRODUCT: Final = Product(version="unknown") +Trace: TypeAlias = tuple[str | None, str | None] def utc_now() -> datetime: return datetime.now(UTC) -@dataclass(frozen=True, slots=True) +def current_trace() -> Trace: + context: Final = get_current_span().get_span_context() + if not context.is_valid: + return None, None + return format_trace_id(context.trace_id), format_span_id(context.span_id) + + class AuditLog: - events: deque[AuditEvent] # mutable-ok: a bounded ring buffer of events is the whole point - now: Callable[[], datetime] = utc_now + def __init__( + self, + records: deque[AuditRecord], # mutable-ok: a bounded ring buffer is the whole point + product: Product, + pipeline: AuditPipeline = NO_PIPELINE, + now: Callable[[], datetime] = utc_now, + trace: Callable[[], Trace] = current_trace, + ) -> None: + self._records = records + self._product = product + self._pipeline = pipeline + self._now = now + self._trace = trace + self._state = ChainState() + + @property + def pipeline(self) -> AuditPipeline: + return self._pipeline + + def accepting(self) -> bool: + return self._pipeline.accepting() + + def reserve(self) -> Reservation | None: + """Hold one queue slot per strict sink for a write that is about to go out, None when no slot is free.""" + return self._pipeline.reserve() def record( self, @@ -54,9 +66,12 @@ class AuditLog: status: int | None = None, credential: str | None = None, detail: str | None = None, + kind: EventKind = "api_activity", + reservation: Reservation | None = None, ) -> AuditEvent: + trace_id, span_id = self._trace() event: Final = AuditEvent( - at=self.now(), + at=self._now(), actor=actor, actor_email=actor_email, dataplane_id=dataplane_id, @@ -66,24 +81,48 @@ class AuditLog: status=status, credential=credential, detail=detail, + kind=kind, + trace_id=trace_id, + span_id=span_id, ) - self.events.append(event) - verbose_proxy_logger.info("control_plane.audit %s", json.dumps(asdict(event), default=str)) + self._state, record = chain(event, self._product, self._state) + self._records.append(record) + if reservation is not None: + reservation.release() + self._pipeline.publish(record) return event + def record_login(self, username: str, decision: Decision, detail: str | None = None) -> AuditEvent: + return self.record( + username, None, CONTROL_PLANE_ID, "POST", "/auth/login", decision, detail=detail, kind="authentication" + ) + + def recent_records( + self, limit: int, dataplane_ids: frozenset[str] | None = None, actor: str | None = None + ) -> tuple[AuditRecord, ...]: + """``dataplane_ids`` of ``None`` means every dataplane, an empty set means nothing is visible.""" + matching: Final = ( + record + for record in reversed(self._records) + if (dataplane_ids is None or record.event.dataplane_id in dataplane_ids) + and (actor is None or record.event.actor == actor) + ) + return tuple(record for _, record in zip(range(limit), matching, strict=False)) + def recent( self, limit: int, dataplane_ids: frozenset[str] | None = None, actor: str | None = None ) -> tuple[AuditEvent, ...]: - """``dataplane_ids`` of ``None`` means every dataplane, an empty set means nothing is visible.""" - newest_first: Final = reversed(self.events) - matching: Final = ( - event - for event in newest_first - if (dataplane_ids is None or event.dataplane_id in dataplane_ids) - and (actor is None or event.actor == actor) - ) - return tuple(event for _, event in zip(range(limit), matching, strict=False)) + return tuple(record.event for record in self.recent_records(limit, dataplane_ids, actor)) -def build_audit_log(capacity: int = DEFAULT_CAPACITY, now: Callable[[], datetime] = utc_now) -> AuditLog: - return AuditLog(events=deque(maxlen=capacity), now=now) +def build_audit_log( + capacity: int = DEFAULT_CAPACITY, + now: Callable[[], datetime] = utc_now, + pipeline: AuditPipeline = NO_PIPELINE, + product: Product = UNKNOWN_PRODUCT, + trace: Callable[[], Trace] = current_trace, +) -> AuditLog: + return AuditLog(deque(maxlen=capacity), product, pipeline, now, trace) + + +__all__ = ["AuditEvent", "AuditLog", "Decision", "Offer", "Reservation", "build_audit_log", "redact_path", "utc_now"] diff --git a/litellm/proxy/control_plane/audit/ocsf.py b/litellm/proxy/control_plane/audit/ocsf.py new file mode 100644 index 00000000000..a089bb675a5 --- /dev/null +++ b/litellm/proxy/control_plane/audit/ocsf.py @@ -0,0 +1,226 @@ +"""Audit events as OCSF 1.5 objects: API Activity (6003) for calls through the control plane, Authentication (3002) +for logins. Each record carries a per-process sequence and a sha256 chain over the previous record so a dropped or +edited line is detectable in whatever sink holds it.""" + +import hashlib +import json +from collections.abc import Mapping +from dataclasses import dataclass +from functools import reduce +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter, ValidationError + +from litellm.proxy.control_plane.audit.log import AuditEvent + +OCSF_VERSION: Final = "1.5.0" +API_ACTIVITY_CLASS: Final = 6003 +AUTHENTICATION_CLASS: Final = 3002 +APPLICATION_ACTIVITY_CATEGORY: Final = 6 +IAM_CATEGORY: Final = 3 +HASH_PREFIX: Final = "sha256:" +VENDOR: Final = "LiteLLM" +PRODUCT: Final = "LiteLLM Control Plane" +Json: TypeAlias = Mapping[str, object] +_JSON: Final = TypeAdapter(Mapping[str, object]) +_ACTIVITY_BY_METHOD: Final[Mapping[str, tuple[int, str]]] = { + "POST": (1, "Create"), + "GET": (2, "Read"), + "PUT": (3, "Update"), + "PATCH": (3, "Update"), + "DELETE": (4, "Delete"), +} +_LOGON: Final = (1, "Logon") +_UNKNOWN: Final = (0, "Unknown") + + +@dataclass(frozen=True, slots=True) +class ChainState: + sequence: int = 0 + prev_hash: str | None = None + + +@dataclass(frozen=True, slots=True) +class AuditRecord: + event: AuditEvent + sequence: int + hash: str + body: Json + + +@dataclass(frozen=True, slots=True) +class Product: + version: str + + +Outcome: TypeAlias = Literal["success", "failure"] + + +def _outcome(event: AuditEvent) -> Outcome: + if event.decision != "allowed": + return "failure" + return "failure" if event.status is not None and event.status >= 400 else "success" + + +def _status(event: AuditEvent) -> Json: + outcome: Final = _outcome(event) + return { + "status_id": 1 if outcome == "success" else 2, + "status": "Success" if outcome == "success" else "Failure", + **({} if event.status is None else {"status_code": str(event.status)}), + **({} if event.detail is None else {"status_detail": event.detail}), + } + + +def _severity(event: AuditEvent) -> Json: + if event.decision == "allowed": + return {"severity_id": 1, "severity": "Informational"} + return {"severity_id": 3, "severity": "Medium"} + + +def _actor(event: AuditEvent) -> Json: + user: Final = {"uid": event.actor, **({} if event.actor_email is None else {"email_addr": event.actor_email})} + return { + "user": user, + **({} if event.credential is None else {"session": {"credential_uid": event.credential}}), + "authorizations": [{"decision": event.decision}], + } + + +def _metadata(event: AuditEvent, product: Product, state: ChainState) -> Json: + return { + "version": OCSF_VERSION, + "product": {"name": PRODUCT, "vendor_name": VENDOR, "version": product.version}, + "sequence": state.sequence + 1, + **({} if event.trace_id is None else {"correlation_uid": event.trace_id}), + } + + +def _litellm(event: AuditEvent, state: ChainState) -> Json: + return { + "decision": event.decision, + "sequence": state.sequence + 1, + "prev_hash": state.prev_hash, + **({} if event.trace_id is None else {"trace_id": event.trace_id}), + **({} if event.span_id is None else {"span_id": event.span_id}), + } + + +def _api_activity(event: AuditEvent, product: Product, state: ChainState) -> Json: + activity_id, activity_name = _ACTIVITY_BY_METHOD.get(event.method, _UNKNOWN) + return { + "class_uid": API_ACTIVITY_CLASS, + "class_name": "API Activity", + "category_uid": APPLICATION_ACTIVITY_CATEGORY, + "category_name": "Application Activity", + "type_uid": API_ACTIVITY_CLASS * 100 + activity_id, + "activity_id": activity_id, + "activity_name": activity_name, + "time": int(event.at.timestamp() * 1000), + **_severity(event), + **_status(event), + "metadata": _metadata(event, product, state), + "actor": _actor(event), + "api": { + "operation": event.method, + "service": {"name": event.dataplane_id, "uid": "dataplane"}, + **({} if event.status is None else {"response": {"code": event.status}}), + }, + "http_request": {"http_method": event.method, "url": {"path": event.path}}, + **({} if event.status is None else {"http_response": {"code": event.status}}), + "unmapped": {"litellm": {"dataplane_id": event.dataplane_id, **_litellm(event, state)}}, + } + + +def _authentication(event: AuditEvent, product: Product, state: ChainState) -> Json: + activity_id, activity_name = _LOGON + return { + "class_uid": AUTHENTICATION_CLASS, + "class_name": "Authentication", + "category_uid": IAM_CATEGORY, + "category_name": "Identity & Access Management", + "type_uid": AUTHENTICATION_CLASS * 100 + activity_id, + "activity_id": activity_id, + "activity_name": activity_name, + "time": int(event.at.timestamp() * 1000), + **_severity(event), + **_status(event), + "metadata": _metadata(event, product, state), + "actor": _actor(event), + "user": {"uid": event.actor}, + "service": {"name": PRODUCT}, + "http_request": {"http_method": event.method, "url": {"path": event.path}}, + "unmapped": {"litellm": _litellm(event, state)}, + } + + +def to_ocsf(event: AuditEvent, product: Product, state: ChainState) -> Json: + match event.kind: + case "api_activity": + return _api_activity(event, product, state) + case "authentication": + return _authentication(event, product, state) + + +def canonical(body: Json) -> bytes: + return json.dumps(body, sort_keys=True, separators=(",", ":"), default=str).encode() + + +def chain(event: AuditEvent, product: Product, state: ChainState) -> tuple[ChainState, AuditRecord]: + """The next record and the state after it. The hash covers the whole body including the previous hash.""" + body: Final = to_ocsf(event, product, state) + digest: Final = HASH_PREFIX + hashlib.sha256(canonical(body)).hexdigest() + litellm_block: Final = dict(_unmapped_litellm(body)) + sealed: Final = {**body, "unmapped": {"litellm": {**litellm_block, "hash": digest}}} + record: Final = AuditRecord(event=event, sequence=state.sequence + 1, hash=digest, body=sealed) + return ChainState(sequence=state.sequence + 1, prev_hash=digest), record + + +def _child(body: Json, key: str) -> Json: + try: + return _JSON.validate_python(body.get(key)) + except ValidationError: + return {} + + +def section(body: Json, *keys: str) -> Json: + """The nested object at ``keys`` or an empty mapping when any step is missing or not an object.""" + return reduce(_child, keys, body) + + +def _unmapped_litellm(body: Json) -> Json: + return section(body, "unmapped", "litellm") + + +@dataclass(frozen=True, slots=True) +class ChainCheck: + """``valid`` is False only on tampering. ``gaps`` counts places where consecutive records are not consecutive + sequences, which is expected for a filtered or limited page and a dropped record for a complete export.""" + + valid: bool + gaps: int + + +def verify(records: tuple[AuditRecord, ...]) -> ChainCheck: + """Every record's hash must match its body, the first record of a process names no predecessor, and records + with adjacent sequences must link. Records separated by a gap are checked on their own only.""" + pairs: Final = tuple(zip(records, records[1:], strict=False)) + adjacent: Final = tuple(pair for pair in pairs if pair[1].sequence == pair[0].sequence + 1) + return ChainCheck( + valid=all(_sealed(record) for record in records) and all(_linked(prev, record) for prev, record in adjacent), + gaps=len(pairs) - len(adjacent), + ) + + +def _sealed(record: AuditRecord) -> bool: + litellm_block: Final = _unmapped_litellm(record.body) + unsealed: Final = { + **record.body, + "unmapped": {"litellm": {key: value for key, value in litellm_block.items() if key != "hash"}}, + } + hashed: Final = litellm_block.get("hash") == HASH_PREFIX + hashlib.sha256(canonical(unsealed)).hexdigest() + return hashed and (litellm_block.get("prev_hash") is None) == (record.sequence == 1) + + +def _linked(prev: AuditRecord, record: AuditRecord) -> bool: + return _unmapped_litellm(record.body).get("prev_hash") == prev.hash diff --git a/litellm/proxy/control_plane/audit/pipeline.py b/litellm/proxy/control_plane/audit/pipeline.py new file mode 100644 index 00000000000..ff4f16bc999 --- /dev/null +++ b/litellm/proxy/control_plane/audit/pipeline.py @@ -0,0 +1,228 @@ +"""Asynchronous delivery of audit records to configured sinks. One bounded queue and one worker per sink, so a sink +that is down never blocks another one or the request path. Full queue: drop the oldest and count it, or in strict +mode refuse the record so the caller can fail the request before it does anything.""" + +import asyncio +from collections.abc import Awaitable, Callable +from contextlib import suppress +from dataclasses import dataclass, field +from typing import Final, Literal, Protocol, TypeAlias + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.control_plane.audit.ocsf import AuditRecord + +Sleeper: TypeAlias = Callable[[float], Awaitable[None]] +Batch: TypeAlias = tuple[AuditRecord, ...] + + +@dataclass(frozen=True, slots=True) +class Delivered: + count: int + + +@dataclass(frozen=True, slots=True) +class Undelivered: + reason: str + + +DeliveryResult: TypeAlias = Delivered | Undelivered + + +class AuditSink(Protocol): + @property + def name(self) -> str: ... + + async def deliver(self, batch: Batch) -> DeliveryResult: ... + + async def close(self) -> None: ... + + +@dataclass(frozen=True, slots=True) +class Accepted: + pass + + +@dataclass(frozen=True, slots=True) +class Dropped: + sink: str + + +@dataclass(frozen=True, slots=True) +class Refused: + sink: str + + +Offer: TypeAlias = Accepted | Dropped | Refused +Mode: TypeAlias = Literal["drop_oldest", "strict"] +EXPONENT_CAP: Final = 30 + + +@dataclass(frozen=True, slots=True) +class Backoff: + base_seconds: float = 0.5 + max_seconds: float = 30.0 + + def delay(self, failures: int) -> float: + return min(self.max_seconds, self.base_seconds * (2.0 ** min(failures, EXPONENT_CAP))) + + +@dataclass(slots=True) +class SinkWorker: + sink: AuditSink + queue: asyncio.Queue[AuditRecord] + batch_size: int + mode: Mode + backoff: Backoff = field(default_factory=Backoff) + sleep: Sleeper = asyncio.sleep + dropped: int = 0 + failures: int = 0 + task: asyncio.Task[None] | None = None + inflight: Batch | None = None + reservations: frozenset[object] = frozenset() + + def offer(self, record: AuditRecord) -> Offer: + if self.mode == "strict" and self.full(): + return Refused(self.sink.name) + if not self.queue.full(): + self.queue.put_nowait(record) + return Accepted() + self.queue.get_nowait() + self.queue.put_nowait(record) + self.dropped += 1 + verbose_proxy_logger.warning( + "control_plane.audit sink %s queue full, dropped oldest (total %d)", self.sink.name, self.dropped + ) + return Dropped(self.sink.name) + + def full(self) -> bool: + """Queued plus reserved slots. Reservations only matter in strict mode, where a refusal is the contract.""" + return self.queue.qsize() + len(self.reservations) >= self.queue.maxsize + + def reserve(self, token: object) -> None: + self.reservations = self.reservations | {token} + + def release(self, token: object) -> None: + self.reservations = self.reservations - {token} + + def _take_batch(self, first: AuditRecord) -> Batch: + rest: Final = tuple(self.queue.get_nowait() for _ in range(min(self.batch_size - 1, self.queue.qsize()))) + return (first, *rest) + + async def _deliver_until_done(self, batch: Batch) -> None: + """Keeps the batch as ``inflight`` until the sink takes it, so a cancelled delivery can be redone.""" + self.inflight = batch + while isinstance(result := await self.sink.deliver(batch), Undelivered): + self.failures += 1 + verbose_proxy_logger.warning( + "control_plane.audit sink %s failed (%s), retry %d", self.sink.name, result.reason, self.failures + ) + await self.sleep(self.backoff.delay(self.failures - 1)) + self.failures = 0 + self.inflight = None + + async def flush_once(self) -> int: + """Deliver what is queued right now, returns how many records went out. Tests drive this directly.""" + if self.queue.empty(): + return 0 + batch: Final = self._take_batch(self.queue.get_nowait()) + await self._deliver_until_done(batch) + return len(batch) + + async def _next_batch(self) -> None: + first: Final = await self.queue.get() + await self._deliver_until_done(self._take_batch(first)) + + async def run(self) -> None: + while True: + await self._next_batch() + + def start(self) -> None: + if self.task is None: + self.task = asyncio.create_task(self.run(), name=f"control_plane.audit.{self.sink.name}") + + async def close(self, drain_seconds: float) -> None: + """Stops the worker, then delivers the batch it was holding and whatever is queued, within the deadline.""" + if self.task is not None: + self.task.cancel() + with suppress(asyncio.CancelledError): + await self.task + try: + await asyncio.wait_for(self._drain(), timeout=drain_seconds) + except TimeoutError: + verbose_proxy_logger.warning( + "control_plane.audit sink %s left %d records undelivered at shutdown", + self.sink.name, + self.queue.qsize() + len(self.inflight or ()), + ) + await self.sink.close() + + async def _drain(self) -> None: + if self.inflight is not None: + await self._deliver_until_done(self.inflight) + while await self.flush_once(): + pass + + +@dataclass(frozen=True, slots=True) +class Reservation: + """One held slot on every strict worker. Releasing twice is harmless, so the request can release in a + ``finally`` after the audit record consumed it.""" + + token: object + workers: tuple[SinkWorker, ...] + + def release(self) -> None: + for worker in self.workers: + worker.release(self.token) + + +@dataclass(frozen=True, slots=True) +class AuditPipeline: + workers: tuple[SinkWorker, ...] + drain_seconds: float = 5.0 + + def publish(self, record: AuditRecord) -> tuple[Offer, ...]: + return tuple(worker.offer(record) for worker in self.workers) + + def accepting(self) -> bool: + """False only in strict mode with a full queue, the signal to answer 503 before a write goes out.""" + return not any(worker.mode == "strict" and worker.full() for worker in self.workers) + + def reserve(self) -> Reservation | None: + """Atomically (no await) takes one slot on every strict worker so a write that passes the gate always has + room for its record, however many writes are in flight. None when any strict queue is full.""" + if not self.accepting(): + return None + token: Final = object() + strict: Final = tuple(worker for worker in self.workers if worker.mode == "strict") + for worker in strict: + worker.reserve(token) + return Reservation(token, strict) + + def start(self) -> None: + for worker in self.workers: + worker.start() + + async def close(self) -> None: + for worker in self.workers: + await worker.close(self.drain_seconds) + + +def build_pipeline( + sinks: tuple[AuditSink, ...], queue_size: int, batch_size: int, strict: bool, sleep: Sleeper = asyncio.sleep +) -> AuditPipeline: + return AuditPipeline( + workers=tuple( + SinkWorker( + sink=sink, + queue=asyncio.Queue(maxsize=queue_size), + batch_size=batch_size, + mode="strict" if strict else "drop_oldest", + sleep=sleep, + ) + for sink in sinks + ) + ) + + +NO_PIPELINE: Final = AuditPipeline(workers=()) diff --git a/litellm/proxy/control_plane/audit/router.py b/litellm/proxy/control_plane/audit/router.py index 06b2fb0bc8c..5f79cb0dfc9 100644 --- a/litellm/proxy/control_plane/audit/router.py +++ b/litellm/proxy/control_plane/audit/router.py @@ -3,22 +3,45 @@ from typing import Annotated, Final from fastapi import APIRouter, Depends, Query, Security from litellm.proxy.control_plane.audit.dependencies import get_audit_log +from litellm.proxy.control_plane.audit.event import CONTROL_PLANE_ID from litellm.proxy.control_plane.audit.log import AuditLog +from litellm.proxy.control_plane.audit.ocsf import verify from litellm.proxy.control_plane.auth.dependencies import require_scopes from litellm.proxy.control_plane.auth.rbac import DataplaneBinding, Principal, Scope from litellm.proxy.control_plane.common.exceptions import ControlPlaneException, forbidden -from litellm.types.proxy.control_plane.audit import AuditEventResponse, AuditListResponse +from litellm.types.proxy.control_plane.audit import AuditEventResponse, AuditListResponse, OcsfAuditListResponse router: Final = APIRouter(prefix="/audit", tags=["audit"]) def visible_dataplanes(binding: DataplaneBinding, requested: str | None) -> frozenset[str] | None: - """What this caller may read: everything when fleet-wide, else only bound ids, and a requested id must be bound.""" + """What this caller may read: everything when fleet-wide, else only bound ids, and a requested id must be bound. + Events against the control plane itself (logins, rejected tokens, denied scopes) are fleet-wide only.""" + if requested == CONTROL_PLANE_ID and not binding.all: + raise ControlPlaneException(forbidden("Control plane events are visible to fleet-wide auditors only.")) if requested is not None and not (binding.all or requested in binding.ids): raise ControlPlaneException(forbidden(f"Dataplane {requested!r} is not bound to this caller.")) if requested is not None: return frozenset({requested}) - return None if binding.all else binding.ids + return None if binding.all else binding.ids - {CONTROL_PLANE_ID} + + +@router.get("/ocsf", response_model=OcsfAuditListResponse) +async def list_audit_ocsf( + principal: Annotated[Principal, Security(require_scopes, scopes=[Scope.AUDIT_READ.value])], + audit: Annotated[AuditLog, Depends(get_audit_log)], + limit: Annotated[int, Query(ge=1, le=1000)] = 100, + dataplane: Annotated[str | None, Query(description="Only events against this dataplane id.")] = None, + actor: Annotated[str | None, Query(description="Only events by this control plane subject.")] = None, +) -> OcsfAuditListResponse: + """The same events as OCSF 1.5 objects, newest first. ``chain_valid`` is false only when a returned record was + altered, ``sequence_gaps`` counts records between the returned ones that the filters or the limit left out.""" + visible: Final = visible_dataplanes(principal.dataplanes, dataplane) + records: Final = audit.recent_records(limit, visible, actor) + check: Final = verify(tuple(reversed(records))) + return OcsfAuditListResponse( + items=tuple(record.body for record in records), chain_valid=check.valid, sequence_gaps=check.gaps + ) @router.get("", response_model=AuditListResponse) diff --git a/litellm/proxy/control_plane/audit/sinks.py b/litellm/proxy/control_plane/audit/sinks.py new file mode 100644 index 00000000000..53ea83bddce --- /dev/null +++ b/litellm/proxy/control_plane/audit/sinks.py @@ -0,0 +1,123 @@ +"""Built-in audit sinks: the stdout log line and an OTel log record over OTLP. Each is one `AuditSink`.""" + +import asyncio +import json +import logging +import time +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +from opentelemetry._logs import LogRecord, SeverityNumber +from opentelemetry.sdk._logs import ReadableLogRecord +from opentelemetry.sdk._logs.export import LogRecordExporter, LogRecordExportResult +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.util.instrumentation import InstrumentationScope +from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags, TraceState, set_span_in_context +from pydantic import TypeAdapter + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.control_plane.audit.ocsf import AuditRecord +from litellm.proxy.control_plane.audit.pipeline import Batch, Delivered, DeliveryResult, Undelivered + +EVENT_NAME: Final = "litellm.control_plane.audit" +SCOPE: Final = InstrumentationScope(EVENT_NAME) +_VERDICT: Final = TypeAdapter(LogRecordExportResult) +LOG_PREFIX: Final = "control_plane.audit" +AUDIT_LOGGER: Final = verbose_proxy_logger.getChild(LOG_PREFIX) +AUDIT_LOGGER.setLevel(logging.INFO) + + +def _line(record: AuditRecord) -> str: + return json.dumps(record.body, default=str, separators=(",", ":")) + + +@dataclass(frozen=True, slots=True) +class LogSink: + name: str = "log" + + async def deliver(self, batch: Batch) -> DeliveryResult: + for record in batch: + AUDIT_LOGGER.info("%s %s", LOG_PREFIX, _line(record)) + return Delivered(len(batch)) + + async def close(self) -> None: + return None + + +def _severity(record: AuditRecord) -> tuple[str, SeverityNumber]: + return ("INFO", SeverityNumber.INFO) if record.event.decision == "allowed" else ("WARN", SeverityNumber.WARN) + + +def _ids(record: AuditRecord) -> tuple[int, int]: + trace_id: Final = int(record.event.trace_id, 16) if record.event.trace_id else 0 + span_id: Final = int(record.event.span_id, 16) if record.event.span_id else 0 + return trace_id, span_id + + +def _attributes(record: AuditRecord) -> Mapping[str, str | int]: + return { + "event.name": EVENT_NAME, + "enduser.id": record.event.actor, + "http.request.method": record.event.method, + "url.path": record.event.path, + "litellm.control_plane.decision": record.event.decision, + "litellm.control_plane.dataplane": record.event.dataplane_id, + "litellm.control_plane.audit.sequence": record.sequence, + "litellm.control_plane.audit.hash": record.hash, + } + + +def _log_data(record: AuditRecord, resource: Resource) -> ReadableLogRecord: + trace_id, span_id = _ids(record) + severity_text, severity_number = _severity(record) + trace_flags: Final = TraceFlags(TraceFlags.SAMPLED if trace_id else TraceFlags.DEFAULT) + return ReadableLogRecord( + LogRecord( + timestamp=int(record.event.at.timestamp() * 1_000_000_000), + observed_timestamp=time.time_ns(), + context=set_span_in_context( + NonRecordingSpan( + SpanContext( + trace_id=trace_id, + span_id=span_id, + is_remote=False, + trace_flags=trace_flags, + trace_state=TraceState(), + ) + ) + ), + severity_text=severity_text, + severity_number=severity_number, + body=_line(record), + attributes=dict(_attributes(record)), + ), + resource, + SCOPE, + ) + + +@dataclass(frozen=True, slots=True) +class OtlpLogSink: + """One OTel log record per audit event, body is the OCSF object, trace and span ids point at the request span. + The batch goes straight to the exporter, off the event loop, so its verdict is the delivery result and a + collector that refuses the batch makes the worker retry instead of the SDK dropping it later.""" + + exporter: LogRecordExporter + resource: Resource + name: str = "otlp" + + async def deliver(self, batch: Batch) -> DeliveryResult: + data: Final = tuple(_log_data(record, self.resource) for record in batch) + try: + result: Final = await asyncio.to_thread(self.exporter.export, data) + except Exception as error: # noqa: BLE001 # an exporter that raises instead of returning FAILURE is still a failed delivery + return Undelivered(f"otlp export raised {type(error).__name__}") + match _VERDICT.validate_python(result): + case LogRecordExportResult.SUCCESS: + return Delivered(len(batch)) + case LogRecordExportResult.FAILURE: + return Undelivered("otlp export failed") + + async def close(self) -> None: + await asyncio.to_thread(self.exporter.shutdown) diff --git a/litellm/proxy/control_plane/auth/dependencies.py b/litellm/proxy/control_plane/auth/dependencies.py index bf4f9ffba90..437cc9d1e36 100644 --- a/litellm/proxy/control_plane/auth/dependencies.py +++ b/litellm/proxy/control_plane/auth/dependencies.py @@ -2,10 +2,13 @@ from typing import Annotated, Final -from fastapi import Depends +from fastapi import Depends, Request from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer, SecurityScopes from starlette.concurrency import run_in_threadpool +from litellm.proxy.control_plane.audit.dependencies import get_audit_log +from litellm.proxy.control_plane.audit.event import CONTROL_PLANE_ID +from litellm.proxy.control_plane.audit.log import AuditLog from litellm.proxy.control_plane.auth.jwt import Rejected, TokenVerifier from litellm.proxy.control_plane.auth.rbac import Principal, Scope from litellm.proxy.control_plane.common.exceptions import ControlPlaneException, forbidden, unauthorized @@ -23,26 +26,47 @@ def get_telemetry() -> Telemetry: return UNTRACED +ANONYMOUS: Final = "anonymous" + + async def get_principal( + request: Request, credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(bearer_scheme)], verifier: Annotated[TokenVerifier, Depends(get_token_verifier)], telemetry: Annotated[Telemetry, Depends(get_telemetry)], + audit: Annotated[AuditLog, Depends(get_audit_log)], ) -> Principal: if credentials is None: raise ControlPlaneException(unauthorized("Missing bearer credentials.")) verified: Final = await run_in_threadpool(verifier.verify, credentials.credentials) if isinstance(verified, Rejected): + audit.record( + ANONYMOUS, None, CONTROL_PLANE_ID, request.method, request.url.path, "denied", 401, detail=verified.reason + ) raise ControlPlaneException(unauthorized(verified.reason)) annotate_principal(verified, telemetry.identity_attributes) return verified async def require_scopes( + request: Request, security_scopes: SecurityScopes, principal: Annotated[Principal, Depends(get_principal)], + audit: Annotated[AuditLog, Depends(get_audit_log)], ) -> Principal: required: Final = frozenset(Scope(scope) for scope in security_scopes.scopes) missing: Final = principal.missing_scopes(required) if missing: + detail: Final = f"missing scope: {', '.join(sorted(scope.value for scope in missing))}" + audit.record( + principal.subject, + principal.email, + CONTROL_PLANE_ID, + request.method, + request.url.path, + "denied", + 403, + detail=detail, + ) raise ControlPlaneException(forbidden(f"Missing scope: {', '.join(sorted(scope.value for scope in missing))}.")) return principal diff --git a/litellm/proxy/control_plane/auth/router.py b/litellm/proxy/control_plane/auth/router.py index cf4c90db00b..bd723c3d691 100644 --- a/litellm/proxy/control_plane/auth/router.py +++ b/litellm/proxy/control_plane/auth/router.py @@ -2,6 +2,8 @@ from typing import Annotated, Final from fastapi import APIRouter, Depends, Security +from litellm.proxy.control_plane.audit.dependencies import get_audit_log +from litellm.proxy.control_plane.audit.log import AuditLog from litellm.proxy.control_plane.auth.dependencies import require_scopes from litellm.proxy.control_plane.auth.login import InvalidCredentials, LoginService, Throttled from litellm.proxy.control_plane.auth.rbac import Principal @@ -30,16 +32,20 @@ def principal_response(principal: Principal) -> PrincipalResponse: async def login( body: LoginRequest, service: Annotated[LoginService, Depends(get_login_service)], + audit: Annotated[AuditLog, Depends(get_audit_log)], ) -> SessionTokenResponse: match service.login(body.username, body.password): case SessionTokenResponse() as token: + audit.record_login(body.username, "allowed") return token case Throttled(retry_after_seconds=seconds): + audit.record_login(body.username, "denied", detail="throttled") raise ControlPlaneException( too_many_requests(f"Too many failed logins, retry in {seconds} seconds."), headers={"Retry-After": str(seconds)}, ) case InvalidCredentials(): + audit.record_login(body.username, "denied", detail="invalid credentials") raise ControlPlaneException(unauthorized("Invalid username or password.")) diff --git a/litellm/proxy/control_plane/common/tracing.py b/litellm/proxy/control_plane/common/tracing.py index 80055965bd4..62503985539 100644 --- a/litellm/proxy/control_plane/common/tracing.py +++ b/litellm/proxy/control_plane/common/tracing.py @@ -5,6 +5,8 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Final, Literal, Protocol from urllib.parse import urlsplit +from opentelemetry.sdk._logs.export import LogRecordExporter +from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SpanExporter from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE @@ -16,6 +18,8 @@ from pydantic import AliasChoices from litellm.integrations.otel.model.config import OpenTelemetryV2Config from litellm.integrations.otel.plumbing.context import inject_trace_context +from litellm.integrations.otel.plumbing.providers import build_log_exporter as _build_log_exporter +from litellm.integrations.otel.plumbing.providers import build_resource as _build_resource from litellm.integrations.otel.plumbing.providers import build_tracer_provider from litellm.proxy.control_plane.audit.log import redact_path from litellm.proxy.control_plane.auth.rbac import Principal @@ -237,9 +241,22 @@ def tracing_enabled(settings: ControlPlaneSettings, environ: Mapping[str, str] = ) +def _config(settings: ControlPlaneSettings) -> OpenTelemetryV2Config: + return OpenTelemetryV2Config.from_env().model_copy(update={"service_name": settings.otel_service_name}) + + def build_provider(settings: ControlPlaneSettings, exporter: SpanExporter | None = None) -> TracerProvider: - config: Final = OpenTelemetryV2Config.from_env().model_copy(update={"service_name": settings.otel_service_name}) - return build_tracer_provider(config, exporter=exporter, use_simple_processor=None if exporter is None else True) + return build_tracer_provider( + _config(settings), exporter=exporter, use_simple_processor=None if exporter is None else True + ) + + +def build_log_exporter(settings: ControlPlaneSettings) -> LogRecordExporter: + return _build_log_exporter(_config(settings)) + + +def build_resource(settings: ControlPlaneSettings) -> Resource: + return _build_resource(_config(settings)) def build_telemetry(settings: ControlPlaneSettings, provider: TracerProvider | None = None) -> Telemetry: diff --git a/litellm/proxy/control_plane/container.py b/litellm/proxy/control_plane/container.py index d8410cd1f0c..e2d6c2a7697 100644 --- a/litellm/proxy/control_plane/container.py +++ b/litellm/proxy/control_plane/container.py @@ -4,6 +4,7 @@ import time from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import timedelta +from importlib.metadata import PackageNotFoundError, version from types import MappingProxyType from typing import Final @@ -15,6 +16,9 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # params is an untyped dict upstream ) from litellm.proxy.control_plane.audit.log import AuditLog, build_audit_log +from litellm.proxy.control_plane.audit.ocsf import Product +from litellm.proxy.control_plane.audit.pipeline import AuditPipeline, AuditSink, build_pipeline +from litellm.proxy.control_plane.audit.sinks import LogSink, OtlpLogSink from litellm.proxy.control_plane.auth.jwt import JwksKeyLookup, SessionTokenCodec, TokenVerifier, build_token_verifier from litellm.proxy.control_plane.auth.login import LoginService, LoginThrottle from litellm.proxy.control_plane.auth.rbac import ( @@ -27,7 +31,7 @@ from litellm.proxy.control_plane.auth.rbac import ( parse_role_definition, ) from litellm.proxy.control_plane.common.fanout import Fanout -from litellm.proxy.control_plane.common.tracing import Telemetry, build_telemetry +from litellm.proxy.control_plane.common.tracing import Telemetry, build_log_exporter, build_resource, build_telemetry from litellm.proxy.control_plane.dataplanes.client import DataplaneClient, HttpTransport, RetryPolicy from litellm.proxy.control_plane.dataplanes.credentials import TokenExchanger, http_token_exchanger from litellm.proxy.control_plane.dataplanes.registry import DataplaneRegistry, build_registry @@ -109,12 +113,33 @@ def _shared_http_client(settings: ControlPlaneSettings) -> HttpTransport: return HandlerTransport(_shared_http_handler(settings)) +def _sink(kind: str, settings: ControlPlaneSettings) -> AuditSink: + match kind: + case "otlp": + return OtlpLogSink(build_log_exporter(settings), build_resource(settings)) + case _: + return LogSink() + + +def build_audit_pipeline(settings: ControlPlaneSettings, sinks: tuple[AuditSink, ...] | None = None) -> AuditPipeline: + chosen: Final = sinks if sinks is not None else tuple(_sink(kind, settings) for kind in settings.audit_sinks) + return build_pipeline(chosen, settings.audit_queue_size, settings.audit_batch_size, settings.audit_strict) + + +def litellm_product() -> Product: + try: + return Product(version=version("litellm")) + except PackageNotFoundError: + return Product(version="unknown") + + def build_container( settings: ControlPlaneSettings, jwks: JwksKeyLookup | None = None, transport: HttpTransport | None = None, exchanger: TokenExchanger | None = None, tracer_provider: TracerProvider | None = None, + audit_sinks: tuple[AuditSink, ...] | None = None, ) -> Container: master_key: Final = settings.master_key.get_secret_value() policy: Final = build_role_policy(settings) @@ -153,6 +178,8 @@ def build_container( tracing=telemetry.recorder, ), fanout=Fanout(concurrency=settings.fanout_concurrency), - audit_log=build_audit_log(settings.audit_capacity), + audit_log=build_audit_log( + settings.audit_capacity, pipeline=build_audit_pipeline(settings, audit_sinks), product=litellm_product() + ), telemetry=telemetry, ) diff --git a/litellm/proxy/control_plane/proxy/router.py b/litellm/proxy/control_plane/proxy/router.py index 0d3a60245ea..0d0ae899b33 100644 --- a/litellm/proxy/control_plane/proxy/router.py +++ b/litellm/proxy/control_plane/proxy/router.py @@ -8,7 +8,8 @@ from fastapi.responses import JSONResponse from pydantic import TypeAdapter, ValidationError from litellm.proxy.control_plane.audit.dependencies import get_audit_log -from litellm.proxy.control_plane.audit.log import AuditLog, redact_path +from litellm.proxy.control_plane.audit.event import redact_path +from litellm.proxy.control_plane.audit.log import AuditLog, Reservation from litellm.proxy.control_plane.auth.dependencies import require_scopes from litellm.proxy.control_plane.auth.rbac import Principal, Scope from litellm.proxy.control_plane.common.exceptions import ControlPlaneException, bad_gateway, error_body, forbidden @@ -23,6 +24,7 @@ from litellm.proxy.control_plane.dataplanes.dependencies import dataplane_access from litellm.proxy.control_plane.dataplanes.registry import Dataplane from litellm.proxy.control_plane.proxy.allowlist import Denied, decide +AUDIT_UNAVAILABLE: Final = "audit-unavailable" DATAPLANE_HEADER: Final = "X-LiteLLM-Dataplane" READ_METHODS: Final = ("GET",) WRITE_METHODS: Final = ("POST", "PUT", "PATCH", "DELETE") @@ -33,12 +35,17 @@ _BODY: Final = TypeAdapter(Mapping[str, object]) _METHOD: Final[TypeAdapter[Method]] = TypeAdapter(Method) +def audit_unavailable() -> ControlPlaneException: + return ControlPlaneException( + error_body(AUDIT_UNAVAILABLE, "Audit Unavailable", 503, "No audit sink can record this write right now.") + ) + + def bad_request(detail: str) -> ControlPlaneException: return ControlPlaneException(error_body("bad-request", "Bad Request", 400, detail)) -async def _json_body(request: Request) -> Mapping[str, object] | None: - raw: Final = await request.body() +def _json_body(raw: bytes) -> Mapping[str, object] | None: if not raw: return None try: @@ -57,10 +64,42 @@ async def _forward( audit: AuditLog, ) -> JSONResponse: method: Final = _METHOD.validate_python(request.method) + raw: Final = await request.body() + reservation: Final = audit.reserve() if method in WRITE_METHODS else None + if method in WRITE_METHODS and reservation is None: + raise audit_unavailable() + try: + return await _forward_reserved( + request, raw, path, method, principal, dataplane, caller, client, audit, reservation + ) + finally: + if reservation is not None: + reservation.release() + + +async def _forward_reserved( + request: Request, + raw: bytes, + path: str, + method: Method, + principal: Principal, + dataplane: Dataplane, + caller: Caller, + client: DataplaneClient, + audit: AuditLog, + reservation: Reservation | None, +) -> JSONResponse: decision: Final = decide(path) if isinstance(decision, Denied): audit.record( - principal.subject, principal.email, dataplane.id, method, decision.path, "denied", detail=decision.reason + principal.subject, + principal.email, + dataplane.id, + method, + decision.path, + "denied", + detail=decision.reason, + reservation=reservation, ) raise ControlPlaneException(forbidden(f"{redact_path(decision.path)}: {decision.reason}.")) result: Final = await client.request( @@ -69,7 +108,7 @@ async def _forward( decision.path, caller, params=tuple(request.query_params.multi_items()), - json=await _json_body(request), + json=_json_body(raw), ) match result: case DataplaneReply(status=status, credential=credential) if 300 <= status < 400: @@ -83,21 +122,44 @@ async def _forward( status, credential, detail="dataplane answered with a redirect, which the control plane never follows", + reservation=reservation, ) raise ControlPlaneException(bad_gateway("Dataplane answered with a redirect, which is not followed.")) case DataplaneReply(status=status, body=body, credential=credential): audit.record( - principal.subject, principal.email, dataplane.id, method, decision.path, "allowed", status, credential + principal.subject, + principal.email, + dataplane.id, + method, + decision.path, + "allowed", + status, + credential, + reservation=reservation, ) return JSONResponse(status_code=status, content=body, headers={DATAPLANE_HEADER: dataplane.id}) case DataplaneUnreachable(detail=detail, attempts=attempts): audit.record( - principal.subject, principal.email, dataplane.id, method, decision.path, "unreachable", detail=detail + principal.subject, + principal.email, + dataplane.id, + method, + decision.path, + "unreachable", + detail=detail, + reservation=reservation, ) raise ControlPlaneException(bad_gateway(f"Dataplane unreachable after {attempts} attempt(s): {detail}")) case CredentialUnavailable(reason=reason): audit.record( - principal.subject, principal.email, dataplane.id, method, decision.path, "unreachable", detail=reason + principal.subject, + principal.email, + dataplane.id, + method, + decision.path, + "unreachable", + detail=reason, + reservation=reservation, ) raise ControlPlaneException(bad_gateway(f"No usable credential for this dataplane: {reason}")) diff --git a/litellm/proxy/control_plane/settings.py b/litellm/proxy/control_plane/settings.py index 4f9b86c9f3b..740b8c8a49b 100644 --- a/litellm/proxy/control_plane/settings.py +++ b/litellm/proxy/control_plane/settings.py @@ -2,7 +2,7 @@ from collections.abc import Mapping from dataclasses import dataclass -from typing import Final, TypeAlias +from typing import Final, Literal, TypeAlias import yaml from pydantic import BaseModel, ConfigDict, Field, SecretStr, TypeAdapter, model_validator @@ -12,6 +12,7 @@ from litellm.litellm_core_utils.oauth_token_exchange import TokenExchangeProfile from litellm.types.proxy.control_plane.auth import PASSWORD_MAX_LENGTH, USERNAME_MAX_LENGTH DEFAULT_SESSION_TTL_SECONDS: Final = 12 * 60 * 60 +AuditSinkKind: TypeAlias = Literal["log", "otlp"] DEFAULT_DATAPLANE_TIMEOUT_SECONDS: Final = 10.0 @@ -112,6 +113,10 @@ class ControlPlaneSettings(BaseSettings): otel_enabled: bool | None = None otel_service_name: str = "litellm-control-plane" otel_identity_attributes: bool = True + audit_sinks: tuple[AuditSinkKind, ...] = ("log",) + audit_queue_size: int = 10_000 + audit_batch_size: int = 100 + audit_strict: bool = False @property def oidc(self) -> OidcSettings | None: diff --git a/litellm/types/proxy/control_plane/audit.py b/litellm/types/proxy/control_plane/audit.py index aab75106dd5..f95ebd57092 100644 --- a/litellm/types/proxy/control_plane/audit.py +++ b/litellm/types/proxy/control_plane/audit.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from datetime import datetime from typing import Literal, TypeAlias @@ -27,3 +28,11 @@ class AuditListResponse(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) items: tuple[AuditEventResponse, ...] + + +class OcsfAuditListResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + items: tuple[Mapping[str, object], ...] + chain_valid: bool + sequence_gaps: int diff --git a/tests/unit/proxy/control_plane/audit/test_ocsf.py b/tests/unit/proxy/control_plane/audit/test_ocsf.py new file mode 100644 index 00000000000..8b435b910c4 --- /dev/null +++ b/tests/unit/proxy/control_plane/audit/test_ocsf.py @@ -0,0 +1,97 @@ +from datetime import UTC, datetime + +from litellm.proxy.control_plane.audit.event import AuditEvent +from litellm.proxy.control_plane.audit.ocsf import ( + AuditRecord, + ChainCheck, + ChainState, + Product, + chain, + section, + to_ocsf, + verify, +) + +T0 = datetime(2026, 10, 13, 12, 0, tzinfo=UTC) +PRODUCT = Product(version="1.2.3") + + +def _event(**overrides: object) -> AuditEvent: + base: dict[str, object] = { + "at": T0, + "actor": "alice", + "actor_email": "alice@x.test", + "dataplane_id": "eu", + "method": "POST", + "path": "/key/generate", + "decision": "allowed", + "status": 200, + "credential": "service_token", + "detail": None, + "trace_id": "0af7651916cd43dd8448eb211c80319c", + "span_id": "b7ad6b7169203331", + } + return AuditEvent(**{**base, **overrides}) # pyright: ignore[reportArgumentType] # test literal + + +def test_an_allowed_write_is_an_api_activity_create_with_success_status() -> None: + body = to_ocsf(_event(), PRODUCT, ChainState()) + assert (body["class_uid"], body["category_uid"], body["type_uid"], body["activity_id"]) == (6003, 6, 600301, 1) + assert (body["status_id"], body["severity_id"], body["time"]) == (1, 1, int(T0.timestamp() * 1000)) + assert body["actor"] == { + "user": {"uid": "alice", "email_addr": "alice@x.test"}, + "session": {"credential_uid": "service_token"}, + "authorizations": [{"decision": "allowed"}], + } + assert body["api"] == { + "operation": "POST", + "service": {"name": "eu", "uid": "dataplane"}, + "response": {"code": 200}, + } + assert body["http_request"] == {"http_method": "POST", "url": {"path": "/key/generate"}} + assert body["metadata"] == { + "version": "1.5.0", + "product": {"name": "LiteLLM Control Plane", "vendor_name": "LiteLLM", "version": "1.2.3"}, + "sequence": 1, + "correlation_uid": "0af7651916cd43dd8448eb211c80319c", + } + + +def test_denied_and_failed_calls_are_failures_with_medium_severity() -> None: + denied = to_ocsf(_event(decision="denied", status=None, detail="no scope"), PRODUCT, ChainState()) + failed = to_ocsf(_event(status=400), PRODUCT, ChainState()) + unreachable = to_ocsf(_event(decision="unreachable", status=None, method="GET"), PRODUCT, ChainState()) + assert (denied["status_id"], denied["severity_id"], denied["status_detail"]) == (2, 3, "no scope") + assert (failed["status_id"], failed["status_code"]) == (2, "400") + assert (unreachable["status_id"], unreachable["activity_id"], unreachable["activity_name"]) == (2, 2, "Read") + + +def test_a_login_is_an_authentication_logon_event() -> None: + body = to_ocsf(_event(kind="authentication", method="POST", path="/auth/login", status=None), PRODUCT, ChainState()) + assert (body["class_uid"], body["category_uid"], body["type_uid"], body["activity_name"]) == ( + 3002, + 3, + 300201, + "Logon", + ) + assert body["user"] == {"uid": "alice"} + + +def test_records_chain_and_verify_detects_a_tampered_or_missing_link() -> None: + state0 = ChainState() + state1, first = chain(_event(), PRODUCT, state0) + state2, second = chain(_event(method="GET", path="/key/list"), PRODUCT, state1) + _, third = chain(_event(method="DELETE", path="/key/delete"), PRODUCT, state2) + assert (first.sequence, second.sequence, third.sequence) == (1, 2, 3) + assert section(first.body, "unmapped", "litellm")["prev_hash"] is None + assert section(second.body, "unmapped", "litellm")["prev_hash"] == first.hash + assert verify((first, second, third)) == ChainCheck(valid=True, gaps=0) + assert verify((first, third)) == ChainCheck(valid=True, gaps=1), "a missing record is a gap, not tampering" + assert verify((second, third)) == ChainCheck(valid=True, gaps=0), "a page may start after the first record" + tampered_body = {**second.body, "actor": {"user": {"uid": "mallory"}}} + tampered = AuditRecord(second.event, second.sequence, second.hash, tampered_body) + assert verify((first, tampered, third)).valid is False, "an edited record no longer matches its hash" + _, forged = chain(_event(method="GET", path="/key/list"), PRODUCT, ChainState(1, "sha256:" + "0" * 64)) + assert verify((first, forged)).valid is False, "a re-hashed record that points at the wrong predecessor" + assert verify((forged,)).valid is True, "alone it is consistent, only the link exposes it" + assert verify(()) == ChainCheck(valid=True, gaps=0) diff --git a/tests/unit/proxy/control_plane/audit/test_pipeline.py b/tests/unit/proxy/control_plane/audit/test_pipeline.py new file mode 100644 index 00000000000..0bca3da00f6 --- /dev/null +++ b/tests/unit/proxy/control_plane/audit/test_pipeline.py @@ -0,0 +1,174 @@ +import asyncio +from dataclasses import dataclass, field +from datetime import UTC, datetime +from typing import Final + +import pytest + +from litellm.proxy.control_plane.audit.event import AuditEvent +from litellm.proxy.control_plane.audit.ocsf import AuditRecord, ChainState, Product, chain +from litellm.proxy.control_plane.audit.pipeline import ( + Accepted, + Batch, + Delivered, + DeliveryResult, + Dropped, + Refused, + Undelivered, + build_pipeline, +) + +T0 = datetime(2026, 10, 13, 12, 0, tzinfo=UTC) + + +@dataclass +class FakeSink: # mutable-ok: records deliveries for assertions + name: str = "fake" + fail_first: int = 0 + batches: list[Batch] = field(default_factory=list) + closed: bool = False + + async def deliver(self, batch: Batch) -> DeliveryResult: + if self.fail_first > 0: + self.fail_first -= 1 + return Undelivered("sink down") + self.batches.append(batch) + return Delivered(len(batch)) + + async def close(self) -> None: + self.closed = True + + +@dataclass +class Clock: # mutable-ok: records sleeps for assertions + slept: list[float] = field(default_factory=list) + + async def sleep(self, seconds: float) -> None: + self.slept.append(seconds) + + +def _records(count: int) -> tuple[AuditRecord, ...]: + state: Final[list[ChainState]] = [ChainState()] + out: Final[list[AuditRecord]] = [] + for index in range(count): + event = AuditEvent(T0, "alice", None, "eu", "GET", f"/key/{index}", "allowed", 200, None, None) + state[0], record = chain(event, Product("t"), state[0]) + out.append(record) + return tuple(out) + + +@pytest.mark.asyncio +async def test_records_are_delivered_in_order_in_batches_of_the_configured_size() -> None: + sink, clock = FakeSink(), Clock() + pipeline = build_pipeline((sink,), queue_size=10, batch_size=2, strict=False, sleep=clock.sleep) + for record in _records(5): + assert pipeline.publish(record) == (Accepted(),) + (worker,) = pipeline.workers + assert [await worker.flush_once() for _ in range(4)] == [2, 2, 1, 0] + assert [[r.sequence for r in batch] for batch in sink.batches] == [[1, 2], [3, 4], [5]] + + +@pytest.mark.asyncio +async def test_a_full_queue_drops_the_oldest_and_counts_it_unless_strict() -> None: + sink = FakeSink() + lenient = build_pipeline((sink,), queue_size=2, batch_size=10, strict=False) + one, two, three = _records(3) + assert [lenient.publish(r) for r in (one, two, three)] == [(Accepted(),), (Accepted(),), (Dropped("fake"),)] + (worker,) = lenient.workers + assert worker.dropped == 1 and lenient.accepting() is True + await worker.flush_once() + assert [r.sequence for r in sink.batches[0]] == [2, 3] + + strict = build_pipeline((FakeSink(),), queue_size=2, batch_size=10, strict=True) + assert [strict.publish(r) for r in (one, two, three)] == [(Accepted(),), (Accepted(),), (Refused("fake"),)] + assert strict.accepting() is False and strict.workers[0].queue.qsize() == 2 + + +def test_a_strict_reservation_holds_the_slot_until_the_record_consumes_or_releases_it() -> None: + strict = build_pipeline((FakeSink(),), queue_size=1, batch_size=10, strict=True) + one, two = _records(2) + held = strict.reserve() + assert held is not None and strict.reserve() is None, "two in-flight writes cannot share one free slot" + assert strict.publish(one) == (Refused("fake"),), "an unreserved event does not steal a reserved slot" + held.release() + assert strict.publish(one) == (Accepted(),) + held.release() + assert strict.publish(two) == (Refused("fake"),), "a second release does not free a slot that is in use" + assert build_pipeline((FakeSink(),), queue_size=1, batch_size=10, strict=False).reserve() is not None + + +@pytest.mark.asyncio +async def test_a_failing_sink_is_retried_with_backoff_and_the_batch_is_never_lost() -> None: + sink, clock = FakeSink(fail_first=3), Clock() + pipeline = build_pipeline((sink,), queue_size=10, batch_size=10, strict=False, sleep=clock.sleep) + for record in _records(2): + pipeline.publish(record) + (worker,) = pipeline.workers + assert await worker.flush_once() == 2 + assert clock.slept == [0.5, 1.0, 2.0] and [r.sequence for r in sink.batches[0]] == [1, 2] + assert worker.failures == 0 + + +@pytest.mark.asyncio +async def test_a_sink_down_for_thousands_of_attempts_still_gets_the_batch_when_it_recovers() -> None: + sink, clock = FakeSink(fail_first=5_000), Clock() + pipeline = build_pipeline((sink,), queue_size=10, batch_size=10, strict=False, sleep=clock.sleep) + pipeline.publish(_records(1)[0]) + assert await pipeline.workers[0].flush_once() == 1 + assert len(clock.slept) == 5_000 and clock.slept[-1] == 30.0 and sink.batches + + +@pytest.mark.asyncio +async def test_each_sink_has_its_own_queue_so_one_down_sink_does_not_block_another() -> None: + down, up, clock = FakeSink("down", fail_first=1), FakeSink("up"), Clock() + pipeline = build_pipeline((down, up), queue_size=10, batch_size=10, strict=False, sleep=clock.sleep) + (record,) = _records(1) + assert pipeline.publish(record) == (Accepted(), Accepted()) + await pipeline.workers[1].flush_once() + assert up.batches and not down.batches + await pipeline.workers[0].flush_once() + assert down.batches + + +@pytest.mark.asyncio +async def test_start_runs_the_worker_and_close_drains_then_closes_the_sink() -> None: + sink = FakeSink() + pipeline = build_pipeline((sink,), queue_size=10, batch_size=10, strict=False) + pipeline.start() + for record in _records(3): + pipeline.publish(record) + await asyncio.sleep(0) + await pipeline.close() + assert sum(len(batch) for batch in sink.batches) == 3 and sink.closed is True + assert pipeline.workers[0].queue.empty() + + +@dataclass +class StallingSink: # mutable-ok: the first delivery hangs until shutdown interrupts it + stalled: asyncio.Event = field(default_factory=asyncio.Event) + calls: int = 0 + batches: list[Batch] = field(default_factory=list) + name: str = "stalling" + + async def deliver(self, batch: Batch) -> DeliveryResult: + self.calls += 1 + if self.calls == 1: + self.stalled.set() + await asyncio.Event().wait() + self.batches.append(batch) + return Delivered(len(batch)) + + async def close(self) -> None: + return None + + +@pytest.mark.asyncio +async def test_close_redelivers_the_batch_a_worker_was_holding_when_it_was_stopped() -> None: + sink = StallingSink() + pipeline = build_pipeline((sink,), queue_size=10, batch_size=10, strict=False) + pipeline.start() + for record in _records(2): + pipeline.publish(record) + await sink.stalled.wait() + await pipeline.close() + assert [[r.sequence for r in batch] for batch in sink.batches] == [[1, 2]] diff --git a/tests/unit/proxy/control_plane/audit/test_router.py b/tests/unit/proxy/control_plane/audit/test_router.py new file mode 100644 index 00000000000..becd5859872 --- /dev/null +++ b/tests/unit/proxy/control_plane/audit/test_router.py @@ -0,0 +1,57 @@ +from litellm.proxy.control_plane.auth.rbac import EDITOR_ROLE +from tests.unit.proxy.control_plane.support import response, rig + +EU = "https://eu.test" + + +def test_ocsf_listing_returns_chained_records_newest_first_with_a_valid_chain() -> None: + client, _, minter = rig({f"{EU}/key/list": [response(200, {})], f"{EU}/key/generate": [response(200, {})]}) + client.get("/api/v1/dataplanes/eu/proxy/key/list", headers=minter.headers("eu")) + client.post("/api/v1/dataplanes/eu/proxy/key/generate", json={}, headers=minter.headers("eu", role=EDITOR_ROLE)) + page = client.get("/api/v1/audit/ocsf?dataplane=eu", headers=minter.headers(role="admin")).json() + assert [(i["class_uid"], i["activity_name"], i["http_request"]["url"]["path"]) for i in page["items"]] == [ + (6003, "Create", "/key/generate"), + (6003, "Read", "/key/list"), + ] + assert (page["chain_valid"], page["sequence_gaps"]) == (True, 0) + assert page["items"][0]["unmapped"]["litellm"]["prev_hash"] == page["items"][1]["unmapped"]["litellm"]["hash"] + newest = client.get("/api/v1/audit/ocsf?limit=1", headers=minter.headers(role="admin")).json() + assert (newest["chain_valid"], newest["sequence_gaps"], len(newest["items"])) == (True, 0, 1) + reads = client.get("/api/v1/audit/ocsf?actor=alice&dataplane=eu", headers=minter.headers(role="admin")).json() + assert reads["chain_valid"] is True, "a filtered page is intact even though it skips records" + + +def test_logins_denied_scopes_and_rejected_tokens_are_audited_against_the_control_plane() -> None: + client, _, minter = rig({}, UI_USERNAME="admin", UI_PASSWORD="pw") + assert client.post("/api/v1/auth/login", json={"username": "admin", "password": "nope"}).status_code == 401 + assert client.post("/api/v1/auth/login", json={"username": "admin", "password": "pw"}).status_code == 200 + assert client.get("/api/v1/audit", headers={"Authorization": "Bearer garbage"}).status_code == 401 + assert client.get("/api/v1/audit", headers=minter.headers(role=EDITOR_ROLE)).status_code == 403 + trail = client.get("/api/v1/audit?dataplane=control-plane", headers=minter.headers(role="admin")).json()["items"] + assert [(e["actor"], e["path"], e["decision"], e["status"], e["detail"]) for e in trail] == [ + ("alice", "/api/v1/audit", "denied", 403, "missing scope: audit:read"), + ("anonymous", "/api/v1/audit", "denied", 401, trail[1]["detail"]), + ("admin", "/auth/login", "allowed", None, None), + ("admin", "/auth/login", "denied", None, "invalid credentials"), + ] + ocsf = client.get( + "/api/v1/audit/ocsf?dataplane=control-plane&limit=1&actor=admin", headers=minter.headers(role="admin") + ) + assert (ocsf.json()["items"][0]["class_uid"], ocsf.json()["chain_valid"]) == (3002, True) + bound = client.get("/api/v1/audit?dataplane=control-plane", headers=minter.headers("control-plane", role="admin")) + assert bound.status_code == 403, "control plane events are fleet-wide only" + implicit = client.get("/api/v1/audit", headers=minter.headers("control-plane", "eu", role="admin")).json() + assert implicit["items"] == [], "a binding that names the control plane does not grant its events either" + + +def test_strict_mode_refuses_writes_but_not_reads_when_no_sink_can_take_the_record() -> None: + client, transport, minter = rig( + {f"{EU}/key/list": [response(200, {}), response(200, {})]}, audit_strict=True, audit_queue_size=1 + ) + client.get("/api/v1/dataplanes/eu/proxy/key/list", headers=minter.headers("eu")) + refused = client.post( + "/api/v1/dataplanes/eu/proxy/key/generate", json={}, headers=minter.headers("eu", role=EDITOR_ROLE) + ) + assert (refused.status_code, refused.json()["type"]) == (503, "urn:litellm:error:audit-unavailable") + assert transport.writes == [], "a refused write never reaches the dataplane" + assert client.get("/api/v1/dataplanes/eu/proxy/key/list", headers=minter.headers("eu")).status_code != 503 diff --git a/tests/unit/proxy/control_plane/audit/test_sinks.py b/tests/unit/proxy/control_plane/audit/test_sinks.py new file mode 100644 index 00000000000..b82520691a6 --- /dev/null +++ b/tests/unit/proxy/control_plane/audit/test_sinks.py @@ -0,0 +1,83 @@ +import logging +from collections.abc import Mapping, Sequence +from datetime import UTC, datetime +from typing import Final + +import pytest +from opentelemetry.sdk._logs import ReadableLogRecord +from opentelemetry.sdk._logs.export import InMemoryLogRecordExporter, LogRecordExporter, LogRecordExportResult +from pydantic import TypeAdapter + +from litellm.proxy.control_plane.audit.event import AuditEvent +from litellm.proxy.control_plane.audit.ocsf import ChainState, Product, chain +from litellm.proxy.control_plane.audit.pipeline import Delivered, Undelivered +from litellm.proxy.control_plane.audit.sinks import LogSink, OtlpLogSink +from litellm.proxy.control_plane.common.tracing import build_resource +from litellm.proxy.control_plane.settings import ControlPlaneSettings + +T0 = datetime(2026, 10, 13, 12, 0, tzinfo=UTC) +TRACE = "0af7651916cd43dd8448eb211c80319c" +SPAN = "b7ad6b7169203331" + + +_JSON: Final = TypeAdapter(Mapping[str, object]) + + +def _record(): + event = AuditEvent( + T0, "alice", None, "eu", "POST", "/key/generate", "denied", 403, None, "no", "api_activity", TRACE, SPAN + ) + return chain(event, Product("t"), ChainState())[1] + + +@pytest.mark.asyncio +async def test_the_otlp_sink_emits_one_log_record_per_event_tied_to_the_request_span() -> None: + exporter = InMemoryLogRecordExporter() + settings = ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": "sk-m"}) + sink = OtlpLogSink(exporter, build_resource(settings)) + record = _record() + assert await sink.deliver((record,)) == Delivered(1) + (emitted,) = exporter.get_finished_logs() + log = emitted.log_record + assert (format(log.trace_id, "032x"), format(log.span_id, "016x")) == (TRACE, SPAN) + assert isinstance(log.body, str) and _JSON.validate_json(log.body) == record.body + assert log.severity_text == "WARN" + assert emitted.resource.attributes["service.name"] == settings.otel_service_name + assert log.attributes is not None and log.attributes["event.name"] == "litellm.control_plane.audit" + assert log.attributes["litellm.control_plane.audit.hash"] == record.hash + assert log.timestamp == int(T0.timestamp() * 1_000_000_000) + + +class RefusingExporter(LogRecordExporter): + def export(self, batch: Sequence[ReadableLogRecord]) -> LogRecordExportResult: + return LogRecordExportResult.FAILURE + + def shutdown(self) -> None: + return None + + def force_flush(self, timeout_millis: int = 30000) -> bool: + return True + + +@pytest.mark.asyncio +async def test_the_otlp_sink_reports_a_refused_export_so_the_worker_retries_it() -> None: + settings = ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": "sk-m"}) + assert await OtlpLogSink(RefusingExporter(), build_resource(settings)).deliver((_record(),)) == Undelivered( + "otlp export failed" + ) + + +@pytest.mark.asyncio +async def test_the_log_sink_writes_the_ocsf_object_as_one_line_even_when_proxy_logs_are_at_warning( + caplog: pytest.LogCaptureFixture, +) -> None: + record = _record() + proxy_logger = logging.getLogger("LiteLLM Proxy") + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy.control_plane.audit"): + proxy_logger.setLevel(logging.WARNING) + try: + assert await LogSink().deliver((record,)) == Delivered(1) + finally: + proxy_logger.setLevel(logging.NOTSET) + (line,) = [m for m in caplog.messages if m.startswith("control_plane.audit ")] + assert '"class_uid":6003' in line and record.hash in line and "\n" not in line diff --git a/tests/unit/proxy/control_plane/proxy/test_router.py b/tests/unit/proxy/control_plane/proxy/test_router.py index abe4934fbf7..4b131a7dc4d 100644 --- a/tests/unit/proxy/control_plane/proxy/test_router.py +++ b/tests/unit/proxy/control_plane/proxy/test_router.py @@ -1,7 +1,24 @@ -import httpx +import asyncio +from typing import Final -from litellm.proxy.control_plane.auth.rbac import EDITOR_ROLE, VIEWER_ROLE -from tests.unit.proxy.control_plane.support import response, rig +import httpx +import pytest +from starlette.requests import Request +from starlette.types import Message, Scope + +from litellm.proxy.control_plane.auth.rbac import EDITOR_ROLE, VIEWER_ROLE, build_principal +from litellm.proxy.control_plane.container import build_container +from litellm.proxy.control_plane.dataplanes.credentials import Caller +from litellm.proxy.control_plane.proxy.router import _forward +from litellm.proxy.control_plane.settings import ControlPlaneSettings +from tests.unit.proxy.control_plane.support import ( + DATAPLANES, + MASTER_KEY, + RoutedTransport, + no_exchange, + response, + rig, +) EU = "https://eu.test" @@ -45,7 +62,7 @@ def test_writes_need_proxy_write_go_out_once_and_are_audited() -> None: assert timed_out.status_code == 502 and "1 attempt" in timed_out.json()["detail"] failed = client.delete("/api/v1/dataplanes/eu/proxy/key/delete", headers=editor) assert (failed.status_code, failed.json()) == (400, {"error": {"message": "no such key"}}) - trail = client.get("/api/v1/audit", headers=minter.headers(role="admin")).json()["items"] + trail = client.get("/api/v1/audit?dataplane=eu", headers=minter.headers(role="admin")).json()["items"] assert [(e["method"], e["path"], e["decision"], e["status"]) for e in trail] == [ ("DELETE", "/key/delete", "allowed", 400), ("PATCH", "/team/update", "unreachable", None), @@ -57,7 +74,9 @@ def test_writes_need_proxy_write_go_out_once_and_are_audited() -> None: def test_denied_paths_are_403_audited_and_never_reach_the_dataplane() -> None: client, transport, minter = rig({}) blocked = client.post( - "/api/v1/dataplanes/eu/proxy/v1/chat/completions", json={"model": "x"}, headers=minter.headers(role=EDITOR_ROLE) + "/api/v1/dataplanes/eu/proxy/v1/chat/completions", + content=b"[", + headers={**minter.headers(role=EDITOR_ROLE), "Content-Type": "application/json"}, ) assert blocked.status_code == 403 and "inference" in blocked.json()["detail"] assert transport.calls == [] @@ -81,7 +100,10 @@ def test_audit_read_needs_its_scope_and_filters_by_actor() -> None: client.get("/api/v1/dataplanes/eu/proxy/key/list", headers=minter.headers("eu")) assert client.get("/api/v1/audit", headers=minter.headers(role=EDITOR_ROLE)).status_code == 403 assert client.get("/api/v1/audit?actor=bob", headers=minter.headers(role="admin")).json()["items"] == [] - assert len(client.get("/api/v1/audit?actor=alice", headers=minter.headers(role="admin")).json()["items"]) == 1 + assert ( + len(client.get("/api/v1/audit?actor=alice&dataplane=eu", headers=minter.headers(role="admin")).json()["items"]) + == 1 + ) def test_audit_reads_are_clamped_to_the_callers_dataplane_binding() -> None: @@ -126,3 +148,73 @@ def test_a_denied_key_bearing_path_is_redacted_in_the_403_too() -> None: client, _, minter = rig({}) denied = client.get("/api/v1/dataplanes/eu/proxy/key/sk-live-123/x%25y", headers=minter.headers("eu")) assert denied.status_code == 403 and "sk-live-123" not in denied.text and "sk-[redacted]" in denied.text + + +@pytest.mark.asyncio +async def test_pending_request_body_does_not_hold_a_strict_audit_reservation() -> None: + settings: Final = ControlPlaneSettings.model_validate( + { + "LITELLM_MASTER_KEY": MASTER_KEY, + "dataplanes": DATAPLANES, + "audit_queue_size": 1, + "audit_strict": True, + } + ) + transport: Final = RoutedTransport({f"{EU}/key/generate": iter([response(200, {"token_id": "t1"})])}) + container: Final = build_container(settings, transport=transport, exchanger=no_exchange) + principal: Final = build_principal( + container.policy, + "alice", + "alice@x.test", + frozenset({EDITOR_ROLE}), + None, + frozenset({"eu"}), + "session", + ) + caller: Final = Caller(principal, "session-token") + dataplane: Final = container.registry.get("eu") + assert dataplane is not None + body_started: Final = asyncio.Event() + release_body: Final = asyncio.Event() + + async def receive() -> Message: + body_started.set() + await release_body.wait() + return {"type": "http.request", "body": b'{"key_alias":"x"}', "more_body": False} + + scope: Final[Scope] = { + "type": "http", + "asgi": {"version": "3.0", "spec_version": "2.3"}, + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/key/generate", + "raw_path": b"/key/generate", + "query_string": b"", + "headers": [(b"content-type", b"application/json")], + "client": ("testclient", 123), + "server": ("testserver", 80), + } + request: Final = Request(scope, receive) + write: Final = asyncio.create_task( + _forward( + request, + "key/generate", + principal, + dataplane, + caller, + container.dataplane_client, + container.audit_log, + ) + ) + await body_started.wait() + await asyncio.sleep(0) + try: + reservation: Final = container.audit_log.reserve() + assert reservation is not None + reservation.release() + finally: + release_body.set() + reply: Final = await write + assert reply.status_code == 200 + assert [(event.path, event.decision) for event in container.audit_log.recent(1)] == [("/key/generate", "allowed")] diff --git a/tests/unit/proxy/control_plane/test_app.py b/tests/unit/proxy/control_plane/test_app.py index e93cc527d51..a60a94f55ae 100644 --- a/tests/unit/proxy/control_plane/test_app.py +++ b/tests/unit/proxy/control_plane/test_app.py @@ -4,9 +4,11 @@ from typing import Final import pytest from fastapi.testclient import TestClient +from opentelemetry.sdk.trace import SpanProcessor, TracerProvider from pydantic import SecretStr from litellm.proxy.control_plane.app import LicenseRequired, create_app +from litellm.proxy.control_plane.audit.pipeline import Batch, Delivered, DeliveryResult from litellm.proxy.control_plane.auth.jwt import SessionTokenCodec from litellm.proxy.control_plane.auth.login import LOCKOUT_FAILURES, LoginService, LoginThrottle, Throttled from litellm.proxy.control_plane.auth.rbac import VIEWER_ROLE, build_principal @@ -27,6 +29,24 @@ class Unlicensed: return False +class BrokenCloseSink: + name: Final = "broken" + + async def deliver(self, batch: Batch) -> DeliveryResult: + return Delivered(len(batch)) + + async def close(self) -> None: + raise RuntimeError("sink refused to close") + + +class ShutdownCounter(SpanProcessor): # mutable-ok: counts provider shutdowns for assertions + def __init__(self) -> None: + self.shutdowns = 0 + + def shutdown(self) -> None: + self.shutdowns += 1 + + def _settings(**overrides: object) -> ControlPlaneSettings: return ControlPlaneSettings.model_validate( {"LITELLM_MASTER_KEY": MASTER_KEY, "UI_USERNAME": "admin", "UI_PASSWORD": "pw", **overrides} @@ -188,3 +208,14 @@ def test_a_role_granting_an_unknown_scope_fails_startup() -> None: def test_a_role_inheriting_an_unknown_role_fails_startup() -> None: with pytest.raises(InvalidRolePolicy, match="inherits unknown role 'ghost'"): build_container(_settings(roles={"x": {"inherits": ["ghost"]}})) + + +def test_tracing_is_shut_down_even_when_an_audit_sink_fails_to_close() -> None: + counter = ShutdownCounter() + provider = TracerProvider() + provider.add_span_processor(counter) + container = build_container(_settings(), tracer_provider=provider, audit_sinks=(BrokenCloseSink(),)) + with pytest.raises(RuntimeError, match="sink refused to close"): + with TestClient(create_app(container, Licensed())): + pass + assert counter.shutdowns == 1