diff --git a/litellm/proxy/control_plane/audit/s3.py b/litellm/proxy/control_plane/audit/s3.py new file mode 100644 index 00000000000..aab07b5f44e --- /dev/null +++ b/litellm/proxy/control_plane/audit/s3.py @@ -0,0 +1,159 @@ +import base64 +import hashlib +import re +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING, Final, Protocol, TypeAlias +from urllib.parse import quote +from uuid import uuid4 + +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"([^<]+)") +Clock: TypeAlias = Callable[[], datetime] + + +def _utc_now() -> datetime: + return datetime.now(UTC) + + +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, follow_redirects=False + ) + + +@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, run_id: str, batch: Batch) -> str: + normalized_prefix: Final = prefix.strip("/") + first: Final = batch[0] + last: Final = batch[-1] + suffix: Final = ( + f"{run_id}/{first.event.at:%Y/%m/%d/%H}/{first.sequence:012d}-{last.sequence:012d}-{last.hash[-12:]}.ndjson" + ) + return f"{normalized_prefix}/{suffix}" if normalized_prefix else suffix + + +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}" + if "." in settings.bucket: + return f"https://{host}/{quote(settings.bucket, safe='')}/{quote(key)}" + return f"https://{settings.bucket}.{host}/{quote(key)}" + + +def put_headers(settings: S3SinkSettings, batch: Batch, content: bytes, now: datetime) -> 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) + if retain_until <= now: + return base + 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" + run_id: str = field(default_factory=lambda: uuid4().hex) + clock: Clock = _utc_now + + async def deliver(self, batch: Batch) -> DeliveryResult: + key: Final = object_key(self.settings.prefix, self.run_id, batch) + url: Final = object_url(self.settings, key) + content: Final = ndjson(batch) + try: + headers: Final = put_headers(self.settings, batch, content, self.clock()) + signed: Final = await run_aws_signing(self.sign, url, content, headers) + 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 diff --git a/litellm/proxy/control_plane/audit/sinks.py b/litellm/proxy/control_plane/audit/sinks.py index df85bca6496..19db1efb647 100644 --- a/litellm/proxy/control_plane/audit/sinks.py +++ b/litellm/proxy/control_plane/audit/sinks.py @@ -26,7 +26,7 @@ EVENT_NAME: Final = "litellm.control_plane.audit" LOG_PREFIX: Final = "control_plane.audit" -def _line(record: AuditRecord) -> str: +def ocsf_line(record: AuditRecord) -> str: return json.dumps(record.body, default=str, separators=(",", ":")) @@ -41,7 +41,7 @@ class LogSink: async def deliver(self, batch: Batch) -> DeliveryResult: for record in batch: - self._logger.info("%s %s", LOG_PREFIX, _line(record)) + self._logger.info("%s %s", LOG_PREFIX, ocsf_line(record)) return Delivered(len(batch)) async def close(self) -> None: @@ -95,7 +95,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, diff --git a/litellm/proxy/control_plane/container.py b/litellm/proxy/control_plane/container.py index 2fed38b5b21..9cac4e2335e 100644 --- a/litellm/proxy/control_plane/container.py +++ b/litellm/proxy/control_plane/container.py @@ -2,17 +2,20 @@ from __future__ import annotations +import os import time from collections.abc import Mapping, Sequence from contextlib import AbstractAsyncContextManager -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import timedelta from importlib.metadata import PackageNotFoundError, version from types import MappingProxyType from typing import TYPE_CHECKING, Final import httpx +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 @@ -20,6 +23,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 @@ -119,10 +123,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() diff --git a/litellm/proxy/control_plane/settings.py b/litellm/proxy/control_plane/settings.py index ec706c59b66..c1cb8b97d17 100644 --- a/litellm/proxy/control_plane/settings.py +++ b/litellm/proxy/control_plane/settings.py @@ -4,6 +4,7 @@ import os 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, field_validator, model_validator @@ -15,10 +16,24 @@ from litellm.proxy.control_plane.auth.rbac import ADMIN_ROLE 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.""" @@ -124,6 +139,51 @@ class ControlPlaneSettings(BaseSettings): audit_batch_size: int = Field(default=100, ge=1) audit_drain_seconds: float = Field(default=5.0, gt=0) 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 + + @field_validator("audit_s3_bucket") + @classmethod + def _s3_bucket_is_not_blank(cls, value: str | None) -> str | None: + if value is not None and not value.strip(): + raise ValueError("CONTROL_PLANE_AUDIT_S3_BUCKET must not be empty or whitespace.") + return value + + @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, + ) @model_validator(mode="after") def _strict_audit_requires_a_sink(self) -> "ControlPlaneSettings": diff --git a/tests/unit/proxy/control_plane/audit/test_s3.py b/tests/unit/proxy/control_plane/audit/test_s3.py new file mode 100644 index 00000000000..5dab6a75c94 --- /dev/null +++ b/tests/unit/proxy/control_plane/audit/test_s3.py @@ -0,0 +1,193 @@ +import base64 +import hashlib +import json +from collections.abc import Mapping +from dataclasses import dataclass, field, replace +from datetime import UTC, datetime, timedelta +from functools import reduce +from typing import Final +from uuid import UUID + +import httpx +import pytest + +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +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 HttpxObjectStore, S3ObjectLockSink, object_key +from litellm.proxy.control_plane.settings import S3SinkSettings + +T0 = datetime(2026, 10, 13, 12, 0, tzinfo=UTC) + + +def test_empty_s3_prefix_does_not_add_a_leading_slash() -> None: + assert not object_key("", "run", _batch(1)).startswith("/") +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: Final = self.outcomes.pop(0) + if isinstance(outcome, Exception): + raise outcome + return outcome + + +@pytest.mark.asyncio +async def test_s3_put_does_not_follow_redirects_from_a_following_http_client() -> None: + transport: Final = httpx.MockTransport( + lambda request: httpx.Response( + 307 if request.url.host == "bucket.example" else 200, + headers={"Location": "https://redirected.example/object"} if request.url.host == "bucket.example" else {}, + ) + ) + handler: Final = AsyncHTTPHandler(transport=transport, follow_redirects=True) + + response: Final = await HttpxObjectStore(handler).put("https://bucket.example/object", b"body", {}, 1.0) + await handler.close() + + assert (response.status_code, str(response.request.url)) == (307, "https://bucket.example/object") + + +def _sign(url: str, content: bytes, headers: Mapping[str, str]) -> Mapping[str, str]: + return {**headers, "Authorization": f"AWS4-HMAC-SHA256 test {len(content)}"} + + +def _append_record(state: tuple[ChainState, Batch], index: int) -> tuple[ChainState, Batch]: + previous, records = state + event: Final = AuditEvent(T0, "alice", None, "eu", "POST", f"/key/{index}", "allowed", 200, None, None) + next_state, record = chain(event, Product("t"), previous) + return next_state, (*records, record) + + +def _batch(count: int) -> Batch: + return reduce(_append_record, range(count), (ChainState(), ()))[1] + + +@pytest.mark.asyncio +async def test_a_batch_is_one_locked_ndjson_object_under_a_dated_key_that_is_never_overwritten() -> None: + store: Final = FakeStore([httpx.Response(200)]) + batch: Final = _batch(3) + sink: Final = S3ObjectLockSink(SETTINGS, store, _sign, clock=lambda: T0) + assert await sink.deliver(batch) == Delivered(3) + ((url, content, headers, timeout),) = store.puts + suffix: Final = batch[-1].hash[-12:] + assert ( + url == f"https://audit.s3.eu-west-1.amazonaws.com/cp/{sink.run_id}/2026/10/13/12/" + f"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: Final = FakeStore([httpx.Response(200)]) + settings: Final = replace(SETTINGS, object_lock_mode=None, endpoint_url="https://minio.internal:9000/") + sink: Final = S3ObjectLockSink(settings, store, _sign) + assert await sink.deliver(_batch(1)) == Delivered(1) + ((url, _, headers, _),) = store.puts + assert url.startswith(f"https://minio.internal:9000/audit/cp/{sink.run_id}/2026/10/13/12/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: Final = FakeStore([httpx.Response(200)]) + sink: Final = S3ObjectLockSink(replace(SETTINGS, region="cn-north-1"), store, _sign) + await sink.deliver(_batch(1)) + ((url, _, _, _),) = store.puts + assert url.startswith(f"https://audit.s3.cn-north-1.amazonaws.com.cn/cp/{sink.run_id}/") + + +@pytest.mark.asyncio +async def test_sink_keys_are_unique_per_sink_and_stable_across_retries() -> None: + batch: Final = _batch(1) + first_store: Final = FakeStore([httpx.Response(200), httpx.Response(200)]) + second_store: Final = FakeStore([httpx.Response(200)]) + first: Final = S3ObjectLockSink(SETTINGS, first_store, _sign) + second: Final = S3ObjectLockSink(SETTINGS, second_store, _sign) + + await first.deliver(batch) + await first.deliver(batch) + await second.deliver(batch) + + first_urls: Final = tuple(url for url, _, _, _ in first_store.puts) + ((second_url, _, _, _),) = second_store.puts + assert UUID(first.run_id).version == UUID(second.run_id).version == 4 + assert first_urls[0] == first_urls[1] and first_urls[0] != second_url + + +@pytest.mark.asyncio +async def test_expired_retention_omits_lock_headers() -> None: + store: Final = FakeStore([httpx.Response(200)]) + sink: Final = S3ObjectLockSink(SETTINGS, store, _sign, clock=lambda: T0 + timedelta(days=30)) + + await sink.deliver(_batch(1)) + + ((_, _, headers, _),) = store.puts + assert not [key for key in headers if key.startswith("x-amz-object-lock")] + + +@pytest.mark.asyncio +async def test_dotted_bucket_uses_regional_path_style_url() -> None: + store: Final = FakeStore([httpx.Response(200)]) + sink: Final = S3ObjectLockSink(replace(SETTINGS, bucket="audit.logs"), store, _sign) + + await sink.deliver(_batch(1)) + + ((url, _, _, _),) = store.puts + assert url.startswith(f"https://s3.eu-west-1.amazonaws.com/audit.logs/cp/{sink.run_id}/") + + +@pytest.mark.asyncio +async def test_an_existing_key_is_delivered_while_s3_errors_and_network_errors_are_retried() -> None: + store: Final = FakeStore( + [ + httpx.Response(412, text="PreconditionFailed"), + httpx.Response(503, text="SlowDown"), + httpx.ConnectError("boom"), + ] + ) + sink: Final = S3ObjectLockSink(SETTINGS, store, _sign) + batch: Final = _batch(2) + assert await sink.deliver(batch) == Delivered(2), "the object is already there from an unacknowledged attempt" + slow: Final = await sink.deliver(batch) + assert isinstance(slow, Undelivered) and slow.reason.startswith( + f"HTTP 503 SlowDown putting cp/{sink.run_id}/2026/10/13/12/" + ) + down: Final = await sink.deliver(batch) + assert isinstance(down, Undelivered) and down.reason.startswith(f"ConnectError putting cp/{sink.run_id}/") + + +@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: Final = FakeStore([]) + sink: Final = S3ObjectLockSink(SETTINGS, store, no_credentials) + result: Final = await sink.deliver(_batch(1)) + assert isinstance(result, Undelivered) and result.reason.startswith(f"RuntimeError signing cp/{sink.run_id}/") + assert store.puts == [] diff --git a/tests/unit/proxy/control_plane/test_container.py b/tests/unit/proxy/control_plane/test_container.py index 4735b6af098..8c6eee4f0e9 100644 --- a/tests/unit/proxy/control_plane/test_container.py +++ b/tests/unit/proxy/control_plane/test_container.py @@ -5,8 +5,9 @@ import httpx import pytest from pydantic import TypeAdapter +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, _shared_http_handler +from litellm.proxy.control_plane.container import HandlerTransport, _environment_region, _shared_http_handler from litellm.proxy.control_plane.settings import ControlPlaneSettings _JSON_COUNT: Final = TypeAdapter(dict[str, int]) @@ -67,3 +68,14 @@ async def test_transport_never_follows_a_dataplane_redirect_even_on_a_following_ async with transport.stream("GET", "https://dp.test/v1/mcp/server/oauth/x/authorize", timeout=1.0) as response: status: Final = response.status_code assert status == 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" diff --git a/tests/unit/proxy/control_plane/test_settings.py b/tests/unit/proxy/control_plane/test_settings.py index ee043acb2b0..5acab7b3179 100644 --- a/tests/unit/proxy/control_plane/test_settings.py +++ b/tests/unit/proxy/control_plane/test_settings.py @@ -4,7 +4,7 @@ import pytest from pydantic import ValidationError from litellm.proxy.control_plane.audit.event import CONTROL_PLANE_ID -from litellm.proxy.control_plane.settings import ControlPlaneSettings, load_settings +from litellm.proxy.control_plane.settings import ControlPlaneSettings, S3SinkSettings, load_settings from litellm.types.proxy.control_plane.auth import PASSWORD_MAX_LENGTH, USERNAME_MAX_LENGTH MASTER_KEY: Final = "01234567890123456789012345678901" @@ -43,6 +43,45 @@ def test_password_login_without_default_roles_requires_a_custom_admin_role() -> assert not settings.include_default_roles and "admin" in settings.roles +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": MASTER_KEY, "audit_sinks": ["log", "s3"]}) + settings: Final = ControlPlaneSettings.model_validate( + { + "LITELLM_MASTER_KEY": MASTER_KEY, + "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": MASTER_KEY}).audit_s3 is None + + +@pytest.mark.parametrize("bucket", ["", " \t "]) +def test_the_s3_sink_rejects_a_blank_bucket(bucket: str) -> None: + with pytest.raises(ValidationError, match="must not be empty or whitespace"): + ControlPlaneSettings.model_validate( + {"LITELLM_MASTER_KEY": MASTER_KEY, "audit_sinks": ["s3"], "audit_s3_bucket": bucket} + ) + + +def test_a_plain_http_s3_endpoint_is_refused_unless_explicitly_allowed() -> None: + base: Final = { + "LITELLM_MASTER_KEY": MASTER_KEY, + "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: Final = 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: Final = ControlPlaneSettings.model_validate( {