feat(control_plane): s3 object lock audit sink

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-10-14 04:21:00 +00:00
parent 1fda6c85f3
commit 23ccb66d85
7 changed files with 394 additions and 7 deletions

View file

@ -0,0 +1,144 @@
"""Append-only audit objects on S3: one NDJSON object per delivered batch, written once with ``If-None-Match: *`` and,
when configured, held under Object Lock so nobody can shorten or delete it inside the retention period."""
import base64
import hashlib
import re
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import timedelta
from typing import TYPE_CHECKING, Final, Protocol
from urllib.parse import quote
import httpx
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy.control_plane.audit.pipeline import Batch, Delivered, DeliveryResult, Undelivered
from litellm.proxy.control_plane.audit.sinks import ocsf_line
from litellm.proxy.control_plane.settings import S3SinkSettings
if TYPE_CHECKING:
from botocore.credentials import Credentials
CONTENT_TYPE: Final = "application/x-ndjson"
RETAIN_UNTIL_FORMAT: Final = "%Y-%m-%dT%H:%M:%SZ"
_ERROR_CODE: Final = re.compile(r"<Code>([^<]+)</Code>")
class ObjectStore(Protocol):
"""The one call the sink makes. Raises only httpx errors and never retries on its own."""
async def put(self, url: str, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: ...
class Signer(Protocol):
def __call__(self, url: str, content: bytes, headers: Mapping[str, str]) -> Mapping[str, str]: ...
@dataclass(frozen=True, slots=True)
class HttpxObjectStore:
"""Holds the handler, not its client: the shared cache may close and replace the client between batches."""
handler: AsyncHTTPHandler
async def put(self, url: str, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response:
return await self.handler.client.put(url, content=content, headers=dict(headers), timeout=timeout)
@dataclass(frozen=True, slots=True)
class SigV4Signer:
"""SigV4 over the usual credential chain (env, profile, IRSA, IMDS), the same way the proxy's s3 logger signs.
Runs in the signing executor because the credential chain may block on STS or IMDS."""
region: str
aws: BaseAWSLLM
def __call__(self, url: str, content: bytes, headers: Mapping[str, str]) -> Mapping[str, str]:
from botocore.auth import S3SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.credentials import RefreshableCredentials
credentials: Final[Credentials] = self.aws.get_credentials(aws_region_name=self.region)
frozen: Final = (
credentials.get_frozen_credentials() if isinstance(credentials, RefreshableCredentials) else credentials
)
request: Final = AWSRequest(method="PUT", url=url, data=content, headers=dict(headers))
S3SigV4Auth(frozen, "s3", self.region).add_auth(request)
return dict(request.headers.items())
def ndjson(batch: Batch) -> bytes:
return "".join(f"{ocsf_line(record)}\n" for record in batch).encode()
def object_key(prefix: str, batch: Batch) -> str:
"""Dated, ordered by sequence, suffixed with the last hash so an identical retry hits the same key."""
first: Final = batch[0]
last: Final = batch[-1]
return f"{prefix}/{first.event.at:%Y/%m/%d/%H}/{first.sequence:012d}-{last.sequence:012d}-{last.hash[-12:]}.ndjson"
def object_url(settings: S3SinkSettings, key: str) -> str:
if settings.endpoint_url is not None:
return f"{settings.endpoint_url.rstrip('/')}/{settings.bucket}/{quote(key)}"
suffix: Final = get_aws_dns_suffix(settings.region)
host: Final = f"s3.{settings.region}.{suffix}" if settings.region else f"s3.{suffix}"
return f"https://{settings.bucket}.{host}/{quote(key)}"
def put_headers(settings: S3SinkSettings, batch: Batch, content: bytes) -> Mapping[str, str]:
"""Integrity headers, the never-overwrite condition and, when a lock mode is set, retention counted from the
newest event in the batch so the clock in the record is the clock of the lock."""
base: Final = {
"Content-Type": CONTENT_TYPE,
"Content-MD5": base64.b64encode(hashlib.md5(content, usedforsecurity=False).digest()).decode(),
"x-amz-content-sha256": hashlib.sha256(content).hexdigest(),
"If-None-Match": "*",
}
if settings.object_lock_mode is None:
return base
retain_until: Final = max(record.event.at for record in batch) + timedelta(days=settings.retention_days)
return {
**base,
"x-amz-object-lock-mode": settings.object_lock_mode,
"x-amz-object-lock-retain-until-date": retain_until.strftime(RETAIN_UNTIL_FORMAT),
}
def _error_code(response: httpx.Response) -> str:
found: Final = _ERROR_CODE.search(response.text)
return found.group(1) if found else "no error code"
@dataclass(frozen=True, slots=True)
class S3ObjectLockSink:
"""A 412 means the key already exists from an earlier attempt whose acknowledgement was lost. The key carries the
last record's hash, which covers every record in the batch, so the stored object is this batch and it counts
as delivered."""
settings: S3SinkSettings
store: ObjectStore
sign: Signer
name: str = "s3"
async def deliver(self, batch: Batch) -> DeliveryResult:
key: Final = object_key(self.settings.prefix, batch)
url: Final = object_url(self.settings, key)
content: Final = ndjson(batch)
try:
signed: Final = await run_aws_signing(self.sign, url, content, put_headers(self.settings, batch, content))
response: Final = await self.store.put(url, content, signed, self.settings.timeout_seconds)
except httpx.HTTPError as error:
return Undelivered(f"{type(error).__name__} putting {key}")
except Exception as error: # noqa: BLE001 # a failing credential chain is a failed delivery, not a dead worker
return Undelivered(f"{type(error).__name__} signing {key}")
match response.status_code:
case 200 | 412:
return Delivered(len(batch))
case status:
return Undelivered(f"HTTP {status} {_error_code(response)} putting {key}")
async def close(self) -> None:
return None

View file

@ -28,7 +28,7 @@ AUDIT_LOGGER: Final = verbose_proxy_logger.getChild(LOG_PREFIX)
AUDIT_LOGGER.setLevel(logging.INFO)
def _line(record: AuditRecord) -> str:
def ocsf_line(record: AuditRecord) -> str:
return json.dumps(record.body, default=str, separators=(",", ":"))
@ -38,7 +38,7 @@ class LogSink:
async def deliver(self, batch: Batch) -> DeliveryResult:
for record in batch:
AUDIT_LOGGER.info("%s %s", LOG_PREFIX, _line(record))
AUDIT_LOGGER.info("%s %s", LOG_PREFIX, ocsf_line(record))
return Delivered(len(batch))
async def close(self) -> None:
@ -89,7 +89,7 @@ def _log_data(record: AuditRecord, resource: Resource) -> ReadableLogRecord:
),
severity_text=severity_text,
severity_number=severity_number,
body=_line(record),
body=ocsf_line(record),
attributes=dict(_attributes(record)),
),
resource,

View file

@ -1,8 +1,9 @@
"""Long-lived objects built once per process and handed to the app. Tests build their own."""
import os
import time
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from dataclasses import dataclass, replace
from datetime import timedelta
from importlib.metadata import PackageNotFoundError, version
from types import MappingProxyType
@ -10,7 +11,9 @@ from typing import Final
import httpx
from opentelemetry.sdk.trace import TracerProvider
from pydantic import TypeAdapter
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # params is an untyped dict upstream
@ -18,6 +21,7 @@ from litellm.llms.custom_httpx.http_handler import (
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.s3 import HttpxObjectStore, S3ObjectLockSink, SigV4Signer
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
@ -113,10 +117,37 @@ def _shared_http_client(settings: ControlPlaneSettings) -> HttpTransport:
return HandlerTransport(_shared_http_handler(settings))
_REGION: Final = TypeAdapter(str)
_OPTIONAL_REGION: Final[TypeAdapter[str | None]] = TypeAdapter(str | None)
def _s3_sink(settings: ControlPlaneSettings) -> AuditSink:
s3: Final = settings.audit_s3
if s3 is None:
raise ValueError("CONTROL_PLANE_AUDIT_SINKS names s3 but CONTROL_PLANE_AUDIT_S3_BUCKET is not set.")
aws: Final = BaseAWSLLM()
region: Final = s3.region or _environment_region(aws)
return S3ObjectLockSink(
replace(s3, region=region), HttpxObjectStore(_shared_http_handler(settings)), SigV4Signer(region, aws)
)
def _environment_region(aws: BaseAWSLLM) -> str:
"""The AWS SDK order: ``AWS_REGION``, then botocore's chain (``AWS_DEFAULT_REGION``, profile, config), then
litellm's resolver (``AWS_REGION_NAME`` and its default)."""
from botocore.session import Session
explicit: Final = os.environ.get("AWS_REGION")
configured: Final = explicit or _OPTIONAL_REGION.validate_python(Session().get_config_variable("region"))
return configured or _REGION.validate_python(aws.get_aws_region_name_for_non_llm_api_calls())
def _sink(kind: str, settings: ControlPlaneSettings) -> AuditSink:
match kind:
case "otlp":
return OtlpLogSink(build_log_exporter(settings), build_resource(settings))
case "s3":
return _s3_sink(settings)
case _:
return LogSink()

View file

@ -3,6 +3,7 @@
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Final, Literal, TypeAlias
from urllib.parse import urlsplit
import yaml
from pydantic import BaseModel, ConfigDict, Field, SecretStr, TypeAdapter, model_validator
@ -12,10 +13,24 @@ 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"]
AuditSinkKind: TypeAlias = Literal["log", "otlp", "s3"]
ObjectLockMode: TypeAlias = Literal["COMPLIANCE", "GOVERNANCE"]
DEFAULT_DATAPLANE_TIMEOUT_SECONDS: Final = 10.0
@dataclass(frozen=True, slots=True)
class S3SinkSettings:
"""Where the s3 audit sink writes. ``region`` is None until the container resolves it from the AWS environment."""
bucket: str
prefix: str
region: str | None
endpoint_url: str | None
object_lock_mode: ObjectLockMode | None
retention_days: int
timeout_seconds: float
@dataclass(frozen=True, slots=True)
class OidcSettings:
"""External IdP used to verify caller JWTs and, later, to exchange them for dataplane tokens."""
@ -117,6 +132,44 @@ class ControlPlaneSettings(BaseSettings):
audit_queue_size: int = 10_000
audit_batch_size: int = 100
audit_strict: bool = False
audit_s3_bucket: str | None = None
audit_s3_prefix: str = "control-plane-audit"
audit_s3_region: str | None = None
audit_s3_endpoint_url: str | None = None
audit_s3_allow_insecure_http: bool = False
audit_s3_object_lock_mode: ObjectLockMode | None = None
audit_s3_retention_days: int = Field(default=365, ge=1)
audit_s3_timeout_seconds: float = 30.0
@model_validator(mode="after")
def _s3_sink_needs_a_bucket(self) -> "ControlPlaneSettings":
if "s3" in self.audit_sinks and self.audit_s3_bucket is None:
raise ValueError("CONTROL_PLANE_AUDIT_SINKS names s3 but CONTROL_PLANE_AUDIT_S3_BUCKET is not set.")
return self
@model_validator(mode="after")
def _s3_endpoint_is_https_unless_opted_out(self) -> "ControlPlaneSettings":
endpoint: Final = self.audit_s3_endpoint_url
if endpoint is None or self.audit_s3_allow_insecure_http or urlsplit(endpoint).scheme.lower() == "https":
return self
raise ValueError(
f"audit s3 endpoint {endpoint!r} is not https, audit records and the SigV4 authorization would travel "
"unencrypted. Use https or set CONTROL_PLANE_AUDIT_S3_ALLOW_INSECURE_HTTP=true for an in-cluster store."
)
@property
def audit_s3(self) -> S3SinkSettings | None:
if self.audit_s3_bucket is None:
return None
return S3SinkSettings(
bucket=self.audit_s3_bucket,
prefix=self.audit_s3_prefix.strip("/"),
region=self.audit_s3_region,
endpoint_url=self.audit_s3_endpoint_url,
object_lock_mode=self.audit_s3_object_lock_mode,
retention_days=self.audit_s3_retention_days,
timeout_seconds=self.audit_s3_timeout_seconds,
)
@property
def oidc(self) -> OidcSettings | None:

View file

@ -0,0 +1,120 @@
import base64
import hashlib
import json
from collections.abc import Mapping
from dataclasses import dataclass, field, replace
from datetime import UTC, datetime
from typing import Final
import httpx
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 Batch, Delivered, Undelivered
from litellm.proxy.control_plane.audit.s3 import S3ObjectLockSink
from litellm.proxy.control_plane.settings import S3SinkSettings
T0 = datetime(2026, 10, 13, 12, 0, tzinfo=UTC)
SETTINGS: Final = S3SinkSettings(
bucket="audit",
prefix="cp",
region="eu-west-1",
endpoint_url=None,
object_lock_mode="COMPLIANCE",
retention_days=30,
timeout_seconds=5.0,
)
@dataclass
class FakeStore: # mutable-ok: scripts responses and records puts for assertions
outcomes: list[httpx.Response | Exception]
puts: list[tuple[str, bytes, Mapping[str, str], float]] = field(default_factory=list)
async def put(self, url: str, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response:
self.puts.append((url, content, headers, timeout))
outcome = self.outcomes.pop(0)
if isinstance(outcome, Exception):
raise outcome
return outcome
def _sign(url: str, content: bytes, headers: Mapping[str, str]) -> Mapping[str, str]:
return {**headers, "Authorization": f"AWS4-HMAC-SHA256 test {len(content)}"}
def _batch(count: int) -> Batch:
state: Final[list[ChainState]] = [ChainState()]
out: Final[list[AuditRecord]] = []
for index in range(count):
event = AuditEvent(T0, "alice", None, "eu", "POST", 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_a_batch_is_one_locked_ndjson_object_under_a_dated_key_that_is_never_overwritten() -> None:
store = FakeStore([httpx.Response(200)])
batch = _batch(3)
assert await S3ObjectLockSink(SETTINGS, store, _sign).deliver(batch) == Delivered(3)
((url, content, headers, timeout),) = store.puts
suffix = batch[-1].hash[-12:]
assert url == f"https://audit.s3.eu-west-1.amazonaws.com/cp/2026/10/13/12/000000000001-000000000003-{suffix}.ndjson"
assert [json.loads(line) for line in content.decode().splitlines()] == [record.body for record in batch]
assert headers["Authorization"] == f"AWS4-HMAC-SHA256 test {len(content)}" and timeout == 5.0
assert headers["If-None-Match"] == "*" and headers["Content-Type"] == "application/x-ndjson"
assert headers["Content-MD5"] == base64.b64encode(hashlib.md5(content, usedforsecurity=False).digest()).decode()
assert headers["x-amz-content-sha256"] == hashlib.sha256(content).hexdigest()
assert (headers["x-amz-object-lock-mode"], headers["x-amz-object-lock-retain-until-date"]) == (
"COMPLIANCE",
"2026-11-12T12:00:00Z",
)
@pytest.mark.asyncio
async def test_without_a_lock_mode_there_are_no_lock_headers_and_a_custom_endpoint_is_path_style() -> None:
store = FakeStore([httpx.Response(200)])
settings = replace(SETTINGS, object_lock_mode=None, endpoint_url="https://minio.internal:9000/")
assert await S3ObjectLockSink(settings, store, _sign).deliver(_batch(1)) == Delivered(1)
((url, _, headers, _),) = store.puts
assert url.startswith("https://minio.internal:9000/audit/cp/2026/10/13/12/000000000001-000000000001-")
assert not [key for key in headers if key.startswith("x-amz-object-lock")]
@pytest.mark.asyncio
async def test_a_china_partition_region_gets_its_own_dns_suffix() -> None:
store = FakeStore([httpx.Response(200)])
await S3ObjectLockSink(replace(SETTINGS, region="cn-north-1"), store, _sign).deliver(_batch(1))
((url, _, _, _),) = store.puts
assert url.startswith("https://audit.s3.cn-north-1.amazonaws.com.cn/cp/")
@pytest.mark.asyncio
async def test_an_existing_key_is_delivered_while_s3_errors_and_network_errors_are_retried() -> None:
store = FakeStore(
[
httpx.Response(412, text="<Error><Code>PreconditionFailed</Code></Error>"),
httpx.Response(503, text="<Error><Code>SlowDown</Code></Error>"),
httpx.ConnectError("boom"),
]
)
sink = S3ObjectLockSink(SETTINGS, store, _sign)
batch = _batch(2)
assert await sink.deliver(batch) == Delivered(2), "the object is already there from an unacknowledged attempt"
slow = await sink.deliver(batch)
assert isinstance(slow, Undelivered) and slow.reason.startswith("HTTP 503 SlowDown putting cp/2026/10/13/12/")
down = await sink.deliver(batch)
assert isinstance(down, Undelivered) and down.reason.startswith("ConnectError putting cp/")
@pytest.mark.asyncio
async def test_a_signer_that_cannot_get_credentials_is_a_failed_delivery_not_a_dead_worker() -> None:
def no_credentials(url: str, content: bytes, headers: Mapping[str, str]) -> Mapping[str, str]:
raise RuntimeError("no credential chain")
store = FakeStore([])
result = await S3ObjectLockSink(SETTINGS, store, no_credentials).deliver(_batch(1))
assert isinstance(result, Undelivered) and result.reason.startswith("RuntimeError signing cp/")
assert store.puts == []

View file

@ -4,8 +4,9 @@ from typing import cast
import httpx
import pytest
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy.control_plane.container import HandlerTransport
from litellm.proxy.control_plane.container import HandlerTransport, _environment_region
@dataclass
@ -53,3 +54,14 @@ async def test_transport_never_follows_a_dataplane_redirect_even_on_a_following_
transport = HandlerTransport(cast(AsyncHTTPHandler, handler)) # cast-ok: only the client property is used
reply = await transport.request("GET", "https://dp.test/v1/mcp/server/oauth/x/authorize", timeout=1.0)
assert reply.status_code == 302 and handler.hops == ["https://dp.test/v1/mcp/server/oauth/x/authorize"]
def test_the_s3_region_follows_the_aws_sdk_order_aws_region_then_default_region(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("AWS_REGION", "us-east-1")
monkeypatch.setenv("AWS_DEFAULT_REGION", "eu-west-1")
monkeypatch.delenv("AWS_REGION_NAME", raising=False)
assert _environment_region(BaseAWSLLM()) == "us-east-1"
monkeypatch.delenv("AWS_REGION")
assert _environment_region(BaseAWSLLM()) == "eu-west-1"

View file

@ -1,7 +1,7 @@
import pytest
from pydantic import ValidationError
from litellm.proxy.control_plane.settings import ControlPlaneSettings
from litellm.proxy.control_plane.settings import ControlPlaneSettings, S3SinkSettings
from litellm.types.proxy.control_plane.auth import PASSWORD_MAX_LENGTH, USERNAME_MAX_LENGTH
@ -16,6 +16,33 @@ def test_configured_admin_credentials_share_the_login_request_bounds(username: s
)
def test_the_s3_sink_refuses_to_start_without_a_bucket_and_reads_its_block_with_one() -> None:
with pytest.raises(ValidationError, match="CONTROL_PLANE_AUDIT_S3_BUCKET"):
ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": "sk-m", "audit_sinks": ["log", "s3"]})
settings = ControlPlaneSettings.model_validate(
{
"LITELLM_MASTER_KEY": "sk-m",
"audit_sinks": ["s3"],
"audit_s3_bucket": "audit",
"audit_s3_prefix": "/cp/",
"audit_s3_object_lock_mode": "COMPLIANCE",
"audit_s3_retention_days": 7,
}
)
assert settings.audit_s3 == S3SinkSettings("audit", "cp", None, None, "COMPLIANCE", 7, 30.0)
assert ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": "sk-m"}).audit_s3 is None
def test_a_plain_http_s3_endpoint_is_refused_unless_explicitly_allowed() -> None:
base = {"LITELLM_MASTER_KEY": "sk-m", "audit_s3_bucket": "audit", "audit_s3_endpoint_url": "http://minio:9000"}
with pytest.raises(ValidationError, match="not https"):
ControlPlaneSettings.model_validate(base)
with pytest.raises(ValidationError, match="not https"):
ControlPlaneSettings.model_validate({**base, "audit_s3_endpoint_url": "HTTP://minio:9000"})
allowed = ControlPlaneSettings.model_validate({**base, "audit_s3_allow_insecure_http": True}).audit_s3
assert allowed is not None and allowed.endpoint_url == "http://minio:9000"
def test_credentials_at_the_bound_are_accepted() -> None:
settings = ControlPlaneSettings.model_validate(
{