From dab2deb5ed2245eafad800a4ce0286d8b394a19c Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 28 Sep 2026 16:12:08 -0700 Subject: [PATCH] test(integration): credential canary suite harness (#43300) * test(integration): credential canary suite harness Adds tests/integration/security with canary generation and search, sweeps over the database, GET routes, client responses, sink doubles and Redis, an owned proxy rig, a sweep sensitivity self-test and the config deployment api_key slot. Registers the security group in run.py, the manifest and the CircleCI integration matrix. * test(integration): widen canary route sweep and harden the rig Enumerate lazily registered feature routers, call parameterized routes with placeholder ids, fail on routes that return no response, skip provider pass-through routes, add an explicit admin-only route allowance, let the sink double use a configurable token, inflate gzip members anywhere in a blob, sweep Redis before the route walk, and trap outbound connections from the owned proxy. * test(integration): descend into any decoded value that can still hold an encoded canary * test(integration): bound canary decoding by depth and decoded bytes * test(integration): scope log-table and spend-log reads to the scenario window * test(integration): sweep spend-log rows in the scenario date window * test(integration): keep spend-log date window summarized * test(integration): resolve deployment ids, scope paginated log lists, key allowances by slot * test(integration): expect 404 from the caller-scoped team membership route * test(integration): use the rig's own master key and expect 404 from submission lookups * test(integration): check the overridden rig key without assuming the default key is unknown --- .circleci/config.yml | 2 +- tests/integration/_support/manifest.py | 1 + tests/integration/run.py | 1 + tests/integration/security/_canary.py | 168 +++++ tests/integration/security/_sinks.py | 297 +++++++++ tests/integration/security/_sweeps.py | 615 ++++++++++++++++++ .../security/test_config_deployment_key.py | 86 +++ .../security/test_sweep_sensitivity.py | 180 +++++ 8 files changed, 1349 insertions(+), 1 deletion(-) create mode 100644 tests/integration/security/_canary.py create mode 100644 tests/integration/security/_sinks.py create mode 100644 tests/integration/security/_sweeps.py create mode 100644 tests/integration/security/test_config_deployment_key.py create mode 100644 tests/integration/security/test_sweep_sensitivity.py diff --git a/.circleci/config.yml b/.circleci/config.yml index 7d4e2e40769..7276da9877b 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3419,7 +3419,7 @@ workflows: name: integration-<< matrix.suite >> matrix: parameters: - suite: [management, accounting, database, providers, mcp, sdk, cost, browser] + suite: [management, accounting, database, providers, mcp, sdk, cost, security, browser] - integration_contracts: name: integration-extensions suite: extensions diff --git a/tests/integration/_support/manifest.py b/tests/integration/_support/manifest.py index 376c2a515b7..a4a86a21568 100644 --- a/tests/integration/_support/manifest.py +++ b/tests/integration/_support/manifest.py @@ -17,5 +17,6 @@ OWNED_DIRECTORIES: Final = frozenset( "compatibility", "sdk", "cost_calculation", + "security", } ) diff --git a/tests/integration/run.py b/tests/integration/run.py index 8f1ff1f4a92..19bce35f542 100644 --- a/tests/integration/run.py +++ b/tests/integration/run.py @@ -19,6 +19,7 @@ GROUPS: Final = MappingProxyType( "mcp": ("mcp",), "sdk": ("sdk",), "cost": ("cost_calculation",), + "security": ("security",), } ) diff --git a/tests/integration/security/_canary.py b/tests/integration/security/_canary.py new file mode 100644 index 00000000000..a5ee99b0a47 --- /dev/null +++ b/tests/integration/security/_canary.py @@ -0,0 +1,168 @@ +"""Canary values and the canary search used by every credential sweep. + +A canary is a unique fake credential planted in one slot (one place the proxy can hold a +credential). Its value is ``lkc--<32 lowercase hex core>``; the slot id +names the source when a sweep finds it, and the random core is what every sweep searches for. + +API: + +- ``SLOTS``: slot id -> ``Slot(identity, description, prefix)``. Stacked suites add their slots + here. ``MARKER`` is not a credential; it is the sensitivity marker sent in message content to + prove that a sweep can see the surface it walks. +- ``canary(slot_id) -> Canary``: a fresh value per call. Call it inside the test (or the fixture + that owns the config holding it), never at import time, so leftovers from earlier runs cannot + match. +- ``find_canary(blob, canaries, *, budget_bytes=DECODE_BUDGET_BYTES) -> tuple[Match, ...]``: + every canary whose core occurs in ``blob`` either raw, inside any base64-looking run after + decoding it (standard and URL-safe alphabets, padded or not, at every 4-character alignment), + or inside a gzip member wherever it starts in the blob. Decoding is applied recursively, so a + gzip body carrying a ``Basic`` header value is still searched. JSON and URL encoding leave a + hex core unchanged, so the raw search covers them. A properly masked value such as + ``sk-...e71b`` is not a match. The search is bounded (three nested layers and ``budget_bytes`` + of decoded output per blob) and raises ``DecodeBudgetExceeded`` rather than returning a + partial result. +""" + +from __future__ import annotations + +import binascii +import re +import uuid +import zlib +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final + +_BASE64_RUN: Final = re.compile(rb"[A-Za-z0-9+/_-]{24,}={0,2}") +_GZIP_MAGIC: Final = b"\x1f\x8b" +_TO_STANDARD: Final = bytes.maketrans(b"-_", b"+/") +_MAX_DEPTH: Final = 3 +DECODE_BUDGET_BYTES: Final = 512 * 1024 * 1024 + + +class DecodeBudgetExceeded(AssertionError): + """A blob needs more decoded bytes than the search budget; the sweep cannot vouch for it.""" + + +@dataclass(slots=True) +class _Budget: + remaining: int + + def spend(self, size: int) -> None: + self.remaining -= size + if self.remaining < 0: + raise DecodeBudgetExceeded("find_canary needed more decoded bytes than its budget for one blob") + + +@dataclass(frozen=True, slots=True) +class Slot: + identity: str + description: str + prefix: str = "" + + +@dataclass(frozen=True, slots=True) +class Canary: + slot: str + core: str + value: str + + +@dataclass(frozen=True, slots=True) +class Match: + slot: str + encoding: str + + +MARKER: Final = "M0" + +SLOTS: Final = MappingProxyType( + { + MARKER: Slot(MARKER, "Sensitivity marker in message content; must appear where prompts are stored"), + "B1": Slot("B1", "Deployment api_key declared in the proxy config.yaml model_list"), + } +) + + +def canary(slot_id: str) -> Canary: + slot: Final = SLOTS[slot_id] + core: Final = uuid.uuid4().hex + return Canary(slot_id, core, f"{slot.prefix}lkc-{slot_id}-{core}") + + +def _decoded_runs(blob: bytes) -> Iterable[tuple[str, bytes]]: + for text in dict.fromkeys(run.group().rstrip(b"=") for run in _BASE64_RUN.finditer(blob)): + for offset in range(4): + aligned = text[offset:] + aligned = aligned[: len(aligned) - len(aligned) % 4] if len(aligned) % 4 == 1 else aligned + padded = aligned + b"=" * (-len(aligned) % 4) + alphabets = (("base64", b"+/"), ("base64url", b"-_")) + for name, extra in alphabets if any(char in aligned for char in b"+/-_") else alphabets[:1]: + try: + yield ( + name, + binascii.a2b_base64( + padded.translate(_TO_STANDARD) if extra == b"-_" else padded, strict_mode=False + ), + ) + except (binascii.Error, ValueError): + continue + + +def _gunzipped(blob: bytes, budget: _Budget) -> Iterable[bytes]: + """Inflate every gzip member in ``blob``, wherever it starts, ignoring trailing bytes.""" + start = blob.find(_GZIP_MAGIC) + while start != -1: + inflater = zlib.decompressobj(16 + zlib.MAX_WBITS) + try: + inflated = inflater.decompress(blob[start:], budget.remaining + 1) + except zlib.error: + inflated = b"" + budget.spend(len(inflated)) + if inflated: + yield inflated + start = blob.find(_GZIP_MAGIC, start + 1) + + +def _matches(blob: bytes, canaries: Sequence[Canary], encoding: str, depth: int, budget: _Budget) -> Iterable[Match]: + lowered: Final = blob.lower() + for candidate in canaries: + if candidate.core.encode() in lowered: + yield Match(candidate.slot, encoding) + if depth >= _MAX_DEPTH: + return + for inflated in _gunzipped(blob, budget): + yield from _matches(inflated, canaries, f"{encoding}>gzip" if encoding != "raw" else "gzip", depth + 1, budget) + for name, decoded in _decoded_runs(blob): + budget.spend(len(decoded)) + label = f"{encoding}>{name}" if encoding != "raw" else name + if _worth_descending(decoded): + yield from _matches(decoded, canaries, label, depth + 1, budget) + else: + lowered_decoded = decoded.lower() + yield from (Match(c.slot, label) for c in canaries if c.core.encode() in lowered_decoded) + + +def _worth_descending(decoded: bytes) -> bool: + """Recursion can only find something through a gzip member or another base64 run. + + Skipping the rest is exact, not a heuristic: the core check has already run on ``decoded``. + """ + return _GZIP_MAGIC in decoded or _BASE64_RUN.search(decoded) is not None + + +def find_canary( + blob: bytes | str, canaries: Sequence[Canary], *, budget_bytes: int = DECODE_BUDGET_BYTES +) -> tuple[Match, ...]: + """Every canary found in ``blob``, one ``Match`` per slot with the shallowest encoding seen. + + Decoding is bounded: at most ``_MAX_DEPTH`` nested layers and ``DECODE_BUDGET_BYTES`` decoded or + inflated bytes per call (``budget_bytes``). Exceeding the byte budget raises ``DecodeBudgetExceeded`` (an + ``AssertionError``) instead of returning a partial, possibly clean, result. + """ + data: Final = blob.encode() if isinstance(blob, str) else blob + found: Final[dict[str, Match]] = {} # mutable-ok: first (shallowest) encoding per slot wins + for match in _matches(data, canaries, "raw", 0, _Budget(budget_bytes)): + found.setdefault(match.slot, match) + return tuple(found.values()) diff --git a/tests/integration/security/_sinks.py b/tests/integration/security/_sinks.py new file mode 100644 index 00000000000..1d2bc236239 --- /dev/null +++ b/tests/integration/security/_sinks.py @@ -0,0 +1,297 @@ +"""The owned proxy every canary scenario runs against, with its provider and sink doubles. + +``canary_rig(root)`` starts a provider double and a ``generic_api`` sink double, writes an owned +config derived from ``tests/integration/proxy_config.yaml`` and starts an owned proxy on it: + +- ``store_prompts_in_spend_logs`` is on, so the stored request body exists for every sweep; +- Redis response-cache entries live 600 s, longer than any scenario's sweeps; +- spend logs flush every second (``proxy_batch_write_at``) and callbacks flush every second + (``DEFAULT_FLUSH_INTERVAL_SECONDS``), so ``eventually`` converges quickly; +- provider-default routes (file, batch, container lists with no deployment) resolve to the + provider double through ``OPENAI_BASE_URL``, and the remote catalogs (cost map, blog posts, + beta headers, autorouter presets, policy templates) are read from the package, so a route + sweep never leaves the machine; +- ``HTTP(S)_PROXY`` points at an egress trap that answers every connection with 403 and + records its first line; the rig fails on exit if the proxy tried to reach any non-loopback + host (``Rig.egress()`` lists the attempts so far); +- the config ``model_list`` declares ``CONFIG_MODEL`` whose ``api_key`` is a fresh slot B1 + canary, reaching the provider double at ``/v1``. + +API: + +- ``canary_rig(root, *, configure=None, environment=None, upstream=None, sink_token=SINK_TOKEN) + -> Iterator[Rig]``. ``configure(config, provider_url)`` may edit the parsed config before it + is written (add deployments, settings, callbacks); ``environment`` adds or overrides proxy + environment variables; ``upstream`` replaces ``chat_upstream`` as the provider double's + handler. ``sink_token`` is the bearer the ``generic_api`` double requires and + ``GENERIC_LOGGER_HEADERS`` sends; pass a ``Canary`` (slot G1 style) to plant a sink credential, + and ``Rig.own_headers`` then allows that one header to carry it (pass it to ``sweep_all``). +- ``Rig.model_id``: the router's ``model_info.id`` for ``CONFIG_MODEL`` (read from + ``/model/info`` once the proxy is up). Pass it as ``ids["model_id"]`` so the + ``{model_id}`` routes (``/credentials/by_model/{model_id}``, ...) resolve the deployment. +- ``Rig.proxy``: the owned proxy ``Gateway`` with its master key (the ``LITELLM_MASTER_KEY`` + from ``environment`` when a scenario overrides it). ``Rig.canaries``: config-held + canaries by slot id. ``Rig.provider`` and ``Rig.sinks[name]``: ``Recorder`` objects whose + ``requests()`` returns every request received so far (the underlying queue is drained into a + list, so repeated polls keep earlier requests). +- ``chat_upstream(request)``: an OpenAI chat double that echoes the last user message and + answers HTTP 400 when the message contains ``PROVIDER_4XX``. +- ``SINK_TOKEN``: the default static bearer the ``generic_api`` sink authenticates with. +- ``team_caller(scenario) -> Caller``: a team, an ``internal_user`` on it and a virtual key for + that user on that team (allowed ``CONFIG_MODEL``). Scenarios send traffic with ``Caller.key`` + and pass ``Caller.callers(rig)`` to ``sweep_all`` so S2 reads every route as the admin and as + the internal user. +- ``settle(rig, request_id, marker)``: wait (bounded) until the spend row for ``request_id`` + is written and every sink double has received an event carrying ``marker``, so the sweeps + that follow read the finished state instead of racing the asynchronous writers. +""" + +from __future__ import annotations + +import json +import socket +import threading +import uuid +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass, field +from types import MappingProxyType +from pathlib import Path +from typing import Final + +import yaml +from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment +from integration._support.database import read_rows +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.security._canary import Canary, canary + +STOCK_CONFIG: Final = Path("tests/integration/proxy_config.yaml") +CONFIG_MODEL: Final = "canary-config-deployment" +PROVIDER_4XX: Final = "canary-provider-4xx" +SINK_TOKEN: Final = "synthetic-canary-sink-token" +GENERIC_SINK: Final = "generic_api" +LOCAL_CATALOGS: Final = MappingProxyType( + { + name: "True" + for name in ( + "LITELLM_LOCAL_MODEL_COST_MAP", + "LITELLM_LOCAL_BLOG_POSTS", + "LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", + "LITELLM_LOCAL_AUTOROUTER_PRESETS", + "LITELLM_LOCAL_POLICY_TEMPLATES", + ) + } +) + + +@dataclass(slots=True) +class Recorder: + wire: Wire + seen: list[Request] = field(default_factory=list) # mutable-ok: drain() consumes the queue + + @property + def url(self) -> str: + return self.wire.url + + def requests(self) -> tuple[Request, ...]: + self.seen.extend(self.wire.drain()) + return tuple(self.seen) + + def carrying(self, text: str) -> tuple[Request, ...]: + """Requests whose body contains ``text``.""" + return tuple(request for request in self.requests() if text.encode() in request.body) + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + owned: OwnedProxy + provider: Recorder + sinks: Mapping[str, Recorder] + canaries: Mapping[str, Canary] + own_headers: Mapping[str, tuple[str, str]] = field(default_factory=lambda: MappingProxyType({})) + egress: Callable[[], tuple[bytes, ...]] = field(default=lambda: ()) + model_id: str = "" + + +def chat_upstream(request: Request) -> Reply: + body: Final = json.loads(request.body or b"{}") + messages: Final = body.get("messages") or [{"content": ""}] + text: Final = str(messages[-1].get("content", "")) + if PROVIDER_4XX in text: + return Reply( + status=400, + body=json.dumps( + {"error": {"type": "invalid_request_error", "code": "canary_rejected", "message": "rejected"}} + ).encode(), + ) + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "echo " + text}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}, + } + ).encode() + ) + + +def _sink_for(token: str) -> Callable[[Request], Reply]: + def sink(request: Request) -> Reply: + assert request.headers.get("authorization") == f"Bearer {token}", "sink double got a foreign bearer" + return Reply() + + return sink + + +@contextmanager +def _egress_trap() -> Iterator[tuple[str, Callable[[], tuple[bytes, ...]]]]: + """A forward-proxy stand-in: records the first line of every connection, answers 403.""" + attempts: Final[list[bytes]] = [] # mutable-ok: appended by the accept thread + server: Final = socket.create_server(("127.0.0.1", 0)) + server.settimeout(0.2) + stopped: Final = threading.Event() + + def serve() -> None: + while not stopped.is_set(): + try: + connection, _ = server.accept() + except TimeoutError: + continue + except OSError: + return + with connection: + connection.settimeout(2) + try: + attempts.append(connection.recv(512).split(b"\r\n", 1)[0]) + connection.sendall(b"HTTP/1.1 403 Forbidden\r\ncontent-length: 0\r\nconnection: close\r\n\r\n") + except OSError: + pass + + thread: Final = threading.Thread(target=serve, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.getsockname()[1]}", lambda: tuple(attempts) + finally: + stopped.set() + thread.join(timeout=5) + server.close() + + +def _config( + root: Path, provider_url: str, b1: Canary, configure: Callable[[dict[str, object], str], None] | None +) -> Path: + config: Final = yaml.safe_load(STOCK_CONFIG.read_text()) + config["model_list"] = [ + { + "model_name": CONFIG_MODEL, + "litellm_params": {"model": "openai/gpt-4o-mini", "api_base": provider_url + "/v1", "api_key": b1.value}, + } + ] + config["general_settings"]["store_prompts_in_spend_logs"] = True + config["litellm_settings"]["cache_params"]["ttl"] = 600 + config["litellm_settings"].update({"callbacks": [GENERIC_SINK], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}) + if configure is not None: + configure(config, provider_url) + path: Final = root / f"canary-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@contextmanager +def canary_rig( + root: Path, + *, + configure: Callable[[dict[str, object], str], None] | None = None, + environment: Mapping[str, str] | None = None, + upstream: Callable[[Request], Reply] | None = None, + sink_token: str | Canary = SINK_TOKEN, +) -> Iterator[Rig]: + b1: Final = canary("B1") + token: Final = sink_token.value if isinstance(sink_token, Canary) else sink_token + own_headers: Final = MappingProxyType( + {GENERIC_SINK: ("authorization", sink_token.slot)} if isinstance(sink_token, Canary) else {} + ) + planted: Final = {"B1": b1, **({sink_token.slot: sink_token} if isinstance(sink_token, Canary) else {})} + with ( + gateway_from_environment() as gateway, + wire_server(upstream or chat_upstream) as provider, + wire_server(_sink_for(token)) as sink, + _egress_trap() as (trap_url, egress), + ): + config: Final = _config(root, provider.url, b1, configure) + overrides: Final = { + **LOCAL_CATALOGS, + **{name: trap_url for name in ("HTTP_PROXY", "HTTPS_PROXY", "http_proxy", "https_proxy")}, + **{name: "127.0.0.1,localhost" for name in ("NO_PROXY", "no_proxy")}, + "OPENAI_BASE_URL": provider.url + "/v1", + "OPENAI_API_BASE": provider.url + "/v1", + "GENERIC_LOGGER_ENDPOINT": sink.url, + "GENERIC_LOGGER_HEADERS": f"Authorization=Bearer {token}", + **(environment or {}), + } + with owned_proxy_process(gateway, root, overrides, config=config) as owned: + admin = Gateway( + owned.gateway.client, overrides.get("LITELLM_MASTER_KEY", owned.gateway.key), owned.gateway.upstream_url + ) + yield Rig( + admin, + owned, + Recorder(provider), + MappingProxyType({GENERIC_SINK: Recorder(sink)}), + MappingProxyType(planted), + own_headers, + egress, + config_model_id(admin), + ) + assert egress() == (), f"Owned proxy tried to reach external hosts: {sorted(set(egress()))}" + + +def config_model_id(gateway: Gateway) -> str: + """The router's ``model_info.id`` for the ``CONFIG_MODEL`` deployment.""" + data: Final = gateway.get("/model/info").get("data") + assert isinstance(data, list), data + found: Final = tuple( + info["id"] + for entry in data + if isinstance(entry, dict) + and entry.get("model_name") == CONFIG_MODEL + and isinstance(info := entry.get("model_info"), dict) + and isinstance(info.get("id"), str) + ) + assert len(found) == 1, f"expected one {CONFIG_MODEL} deployment in /model/info, got {found}" + return str(found[0]) + + +@dataclass(frozen=True, slots=True) +class Caller: + team_id: str + user_id: str + key: str + + def callers(self, rig: Rig) -> Mapping[str, str]: + return {"admin": rig.proxy.key, "internal_user": self.key} + + +def team_caller(scenario: Scenario) -> Caller: + team: Final = scenario.team() + user: Final = scenario.user(user_role="internal_user") + scenario.gateway.post("/team/member_add", {"team_id": team, "member": {"user_id": user, "role": "user"}}) + key: Final = scenario.key(team_id=team, user_id=user, models=[CONFIG_MODEL]) + return Caller(team, user, key) + + +def settle(rig: Rig, request_id: str, marker: Canary) -> None: + eventually( + lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)), + lambda rows: len(rows) == 1, + seconds=70, + ) + for sink in rig.sinks.values(): + eventually(lambda sink=sink: sink.carrying(marker.core), bool, seconds=30) diff --git a/tests/integration/security/_sweeps.py b/tests/integration/security/_sweeps.py new file mode 100644 index 00000000000..617bf4c9bae --- /dev/null +++ b/tests/integration/security/_sweeps.py @@ -0,0 +1,615 @@ +"""Sweeps: every place a canary must NOT appear, searched with ``find_canary``. + +Each sweep returns ``Hit(sweep, location, slot, encoding)`` records; ``assert_no_hits`` fails +with a table that names the slot, the sweep and the exact location, so the code path that copied it is +usually obvious from the failure alone. The sweeps are generic on purpose: a new table, a new +GET route or a new copy of the request body is covered without editing this module. + +API: + +- ``sweep_database(canaries, *, database_url=None) -> tuple[Hit, ...]`` (S1): every base table + of every non-system schema from ``information_schema.tables``, read as + ``SELECT to_jsonb(t)::text FROM ""."" t``. Location is ``table.column`` + (``schema.table.column`` outside ``public``); a table dropped mid-sweep is skipped. With + ``since``, the append-only log tables in ``TIME_SCOPED_TABLES`` are read from ``since`` on + (minus ``SCOPE_SLACK``), so the sweep stays fast on a database shared by many tests. +- ``get_routes() -> tuple[str, ...]`` and ``sweep_routes(gateway, canaries, ids, *, callers)`` + (S2): every GET route registered on the proxy app (``app.routes``, which includes the + routes hidden from the OpenAPI spec and every lazily registered feature router), enumerated + once per session by importing the app in a child interpreter. Path parameters are filled from + ``ids`` (parameter name -> value), then from ``DEFAULT_IDS``; any other parameter gets + ``PLACEHOLDER_ID`` so the route is still called and its (usually 404) response still searched. + A parameter in ``REAL_ID_REQUIRED`` is never given a placeholder (the proxy would call a public + provider); such a route is skipped unless ``ids`` supplies it. Routes called with a placeholder + or skipped for want of a real id are listed in ``RouteSweep.unfilled``; pass real ids to make + them return data. ``route_denied(route)`` names why a route is skipped: ``ROUTE_DENY_LIST`` + holds the routes that stream forever, redirect into an external flow or contact an external + service, and ``PROVIDER_PASSTHROUGH`` matches the ``//{endpoint:path}`` routes that + forward to the provider (swept by the pass-through slots, not by S2). Every response is searched + whatever its status; responses with status >= 500 are also listed in ``RouteSweep.errors``. + A call that got no response at all (timeout, reset) is listed in ``RouteSweep.unreachable``, + and ``sweep_all`` fails on it, since that route went unchecked. ``ADMIN_ONLY_ALLOWANCES`` + names exact ``(route, caller label)`` pairs allowed to return a credential by design, and + ``ALLOWANCE_SLOT_FAMILIES`` the slot families each pair may return; those hits land in + ``RouteSweep.allowed`` instead of ``hits``, while any other slot on that route, and every other + caller of it, is still a hit. A route whose path parameters all came from ``ids`` must not + answer the admin with 404 (an id the scenario passed is wrong, so the route saw no data); + such calls are listed in ``RouteSweep.not_found`` and ``sweep_all`` fails on them, except the + routes in ``NOT_FOUND_EXPECTED``. ``PARAMETER_ALIASES`` fills a parameter from another id for + the routes where the name misleads (``/v1/models/{model_id}`` takes the public model name, so + it is filled from ``ids["model"]``, while ``/credentials/by_model/{model_id}`` takes the + router's deployment id). + ``RouteSweep.statuses`` maps each call's location to its status code. ``record_route_sweep(routes, node)`` appends the report to + ``$INTEGRATION_RESULTS_DIR/security-route-sweep.jsonl`` (a CI artifact). With ``since``, + the log list routes (``SCENARIO_SCOPED_LIST_ROUTES``: ``/spend/logs``, ``/spend/logs/ui``, + ``/spend/logs/v2``) are called with this scenario's request id, user id and a date window + (summarized for ``/spend/logs``; ``since`` to ``since + LIST_WINDOW`` with ``LIST_PAGE_SIZE`` + rows for the paginated two) instead of unfiltered. A 4xx from one of those calls is listed in + ``RouteSweep.rejected`` and ``sweep_all`` fails on it, since the route then returned no rows. + ``scoped_queries(route, ids, since)`` returns the query strings S2 uses for a route. +- ``sweep_responses(responses, canaries) -> tuple[Hit, ...]`` (S3): body and headers of every + client-facing response the scenario received. +- ``sweep_sink(name, requests, canaries, *, own_header=None) -> tuple[Hit, ...]`` (S4): every + byte a sink double received (gzip bodies are inflated by ``find_canary``). ``own_header`` is + the ``(header name, slot)`` pair the sink legitimately authenticates with; that one header may + carry that one canary. +- ``sweep_redis(canaries, *, host, port) -> tuple[Hit, ...]`` (S5): ``SCAN`` of every key, with + strings, hashes, lists, sets and sorted sets dumped and searched along with the key name. +- ``sweep_all(gateway, canaries, *, responses, sinks, ids, callers=None, own_headers=None, + since=None) -> SweepReport``: S1 to S5 in one pass for a finished scenario. Search the scenario's marker + and its credential canaries together; ``SweepReport.credential_hits()`` is every hit that is not the + marker, and ``assert_marker_seen(report, expected)`` is the per-test sensitivity control + (``expected`` maps a sweep id to a location substring the marker must be reported at). +""" + +from __future__ import annotations + +import json +import os +import re +import subprocess +import sys +from collections.abc import Callable, Iterable, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from datetime import UTC, datetime, timedelta +from functools import cache +from pathlib import Path +from types import MappingProxyType +from typing import Final +from urllib.parse import quote, urlencode + +import httpx +import psycopg +from integration._support.client import Gateway +from integration._support.wire import Request +from integration.security._canary import MARKER, Canary, find_canary +from psycopg import sql +from redis import Redis + +_PATH_PARAMETER: Final = re.compile(r"{([^}:]+)(?::[^}]+)?}") +_ROUTE_TIMEOUT: Final = 20.0 + +ROUTE_DENY_LIST: Final = MappingProxyType( + { + "/mcp": "streamable HTTP GET opens a server-sent event stream that never ends", + "/mcp/proxy": "MCP transport endpoint, not a JSON read", + "/{mcp_server_name}/mcp": "MCP transport endpoint, not a JSON read", + "/toolset/{toolset_name}/mcp": "MCP transport endpoint, not a JSON read", + "/sso/key/generate": "starts an external SSO redirect flow", + "/sso/callback": "external SSO redirect target", + "/sso/saml/login": "starts an external SAML redirect flow", + "/sso/debug/login": "starts an external SSO redirect flow", + "/sso/debug/callback": "external SSO redirect target", + "/fallback/login": "HTML login page", + "/plugin-proxy/{plugin_name}/{path:path}": "reverse proxy to a plugin process", + "/openai_passthrough/{endpoint:path}": "forwards to a provider, not a proxy read", + "/get/latest_release_info": "fetches the latest release from api.github.com", + } +) + +PROVIDER_PASSTHROUGH: Final = re.compile(r"^(/[^/{}]+)+/\{endpoint:path\}$") +PROVIDER_PASSTHROUGH_REASON: Final = "provider pass-through: forwards to the provider, not a proxy read" + + +def route_denied(route: str) -> str | None: + """Why S2 skips ``route``, or None when it is swept.""" + if route in ROUTE_DENY_LIST: + return ROUTE_DENY_LIST[route] + return PROVIDER_PASSTHROUGH_REASON if PROVIDER_PASSTHROUGH.match(route) else None + + +DEFAULT_IDS: Final = MappingProxyType({"provider": "openai"}) +PLACEHOLDER_ID: Final = "canary-placeholder-id" +REAL_ID_REQUIRED: Final = MappingProxyType( + { + "video_id": "a video id encodes its provider; an unknown id falls back to the public OpenAI API", + "character_id": "a character id encodes its provider; an unknown id falls back to the public OpenAI API", + } +) + +PARAMETER_ALIASES: Final = MappingProxyType( + { + "/models/{model_id}": {"model_id": "model"}, + "/v1/models/{model_id}": {"model_id": "model"}, + } +) +NOT_FOUND_EXPECTED: Final = MappingProxyType( + { + "/fallback/{model}": "answers 404 when the model has no fallbacks configured", + "/team/{team_id}/members/me": "answers 404 when the caller is not a member, which the admin is not", + "/guardrails/submissions/{guardrail_id}": "answers 404 for a guardrail no team submitted for review", + } +) + +ADMIN_ONLY_ALLOWANCES: Final = MappingProxyType( + { + ("/get/config/callbacks", "admin"): ( + "proxy admin holds the master key and edits these env values in the config UI" + ), + } +) + + +ALLOWANCE_SLOT_FAMILIES: Final = MappingProxyType({("/get/config/callbacks", "admin"): ("G",)}) + + +def route_allowance(route: str, caller: str, slot: str | None = None) -> str | None: + """The documented reason ``caller`` may read a credential from ``route``, or None. + + With ``slot``, the allowance also has to cover that slot: its id must start with one of the + families in ``ALLOWANCE_SLOT_FAMILIES`` for the pair (``/get/config/callbacks`` serves the + callback env values, so only the G-family sink credentials), so any other slot found there + is still a hit. + """ + reason: Final = ADMIN_ONLY_ALLOWANCES.get((route, caller)) + if reason is None or slot is None: + return reason + return reason if slot.startswith(ALLOWANCE_SLOT_FAMILIES.get((route, caller), ())) else None + + +@dataclass(frozen=True, slots=True) +class Hit: + sweep: str + location: str + slot: str + encoding: str + + +@dataclass(frozen=True, slots=True) +class RouteSweep: + hits: tuple[Hit, ...] + called: tuple[str, ...] + unfilled: tuple[str, ...] + errors: tuple[str, ...] = field(default=()) + unreachable: tuple[str, ...] = field(default=()) + allowed: tuple[Hit, ...] = field(default=()) + not_found: tuple[str, ...] = field(default=()) + rejected: tuple[str, ...] = field(default=()) + statuses: Mapping[str, int] = field(default_factory=lambda: MappingProxyType({})) + + +def format_hits(hits: Iterable[Hit]) -> str: + rows: Final = tuple((hit.slot, hit.sweep, hit.encoding, hit.location) for hit in hits) + header: Final = ("slot", "sweep", "encoding", "location") + widths: Final = tuple(max(len(row[index]) for row in (header, *rows)) for index in range(3)) + return "\n".join( + f"{slot:<{widths[0]}} {sweep:<{widths[1]}} {encoding:<{widths[2]}} {location}" + for slot, sweep, encoding, location in (header, *rows) + ) + + +def assert_no_hits(hits: Sequence[Hit], context: str) -> None: + assert not hits, f"Credential canary found outside its destination ({context}):\n{format_hits(hits)}" + + +def _hits(sweep: str, location: str, blob: bytes | str, canaries: Sequence[Canary]) -> tuple[Hit, ...]: + return tuple(Hit(sweep, location, match.slot, match.encoding) for match in find_canary(blob, canaries)) + + +def sweep_database( + canaries: Sequence[Canary], *, database_url: str | None = None, since: datetime | None = None +) -> tuple[Hit, ...]: + """S1: every row of every base table, as ``to_jsonb``, attributed to the column that holds it. + + With ``since``, the append-only log tables in ``TIME_SCOPED_TABLES`` are read only for rows + written or changed at or after it; every other table is still read in full. + """ + found: Final[list[Hit]] = [] # mutable-ok: accumulated across tables + with psycopg.connect(database_url or os.environ["DATABASE_URL"], autocommit=True) as connection: + tables: Final = connection.execute( + "SELECT table_schema, table_name FROM information_schema.tables " + "WHERE table_type = 'BASE TABLE' AND table_schema NOT IN ('pg_catalog', 'information_schema') " + "ORDER BY table_schema, table_name" + ).fetchall() + for schema, table in tables: + query = sql.SQL("SELECT to_jsonb(t)::text FROM {}.{} t").format( + sql.Identifier(schema), sql.Identifier(table) + ) + scoped = TIME_SCOPED_TABLES.get(table) if since is not None else None + if scoped is not None: + query = sql.SQL("{} WHERE {}").format( + query, + sql.SQL(" OR ").join( + sql.SQL("t.{} >= {}").format(sql.Identifier(column), sql.Literal(_naive_utc(since))) + for column in scoped + ), + ) + where = table if schema == "public" else f"{schema}.{table}" + try: + rows = connection.execute(query).fetchall() + except psycopg.errors.UndefinedTable: + continue + for (row,) in rows: + if not find_canary(row, canaries): + continue + for column, value in json.loads(row).items(): + found.extend(_hits("S1", f"{where}.{column}", json.dumps(value), canaries)) + return tuple(found) + + +TIME_SCOPED_TABLES: Final = MappingProxyType( + { + "LiteLLM_SpendLogs": ("startTime", "updated_at"), + "LiteLLM_ErrorLogs": ("startTime", "endTime"), + "LiteLLM_AuditLog": ("updated_at",), + "LiteLLM_DeletedTeamTable": ("deleted_at",), + "LiteLLM_DeletedVerificationToken": ("deleted_at",), + } +) +SCOPE_SLACK: Final = timedelta(seconds=5) + + +def _naive_utc(moment: datetime) -> datetime: + """Prisma writes these columns as naive UTC; compare with a little slack for clock skew.""" + aware: Final = moment if moment.tzinfo is not None else moment.replace(tzinfo=UTC) + return (aware - SCOPE_SLACK).astimezone(UTC).replace(tzinfo=None) + + +def _route_queries(route: str, ids: Mapping[str, str], since: datetime | None) -> tuple[str, ...]: + """Query strings a route is called with; unbounded list routes are narrowed to this scenario.""" + if route not in SCENARIO_SCOPED_LIST_ROUTES or since is None: + return ("",) + aware: Final = since if since.tzinfo is not None else since.replace(tzinfo=UTC) + return tuple("?" + urlencode(query) for query in SCENARIO_SCOPED_LIST_ROUTES[route](ids, aware.astimezone(UTC))) + + +def scoped_queries(route: str, ids: Mapping[str, str], since: datetime | None) -> tuple[str, ...]: + """The query strings S2 calls ``route`` with (``("",)`` unless it is a scoped list route).""" + return _route_queries(route, ids, since) + + +def _scenario_filters(ids: Mapping[str, str]) -> tuple[Mapping[str, str], ...]: + return ( + *(({"request_id": ids["request_id"]},) if "request_id" in ids else ()), + *(({"user_id": ids["user_id"]},) if "user_id" in ids else ()), + ) + + +def _spend_logs_queries(ids: Mapping[str, str], since: datetime) -> tuple[Mapping[str, str], ...]: + window: Final = { + "start_date": since.date().isoformat(), + "end_date": (datetime.now(UTC).date() + timedelta(days=1)).isoformat(), + } + return (*_scenario_filters(ids), window) + + +LIST_PAGE_SIZE: Final = 50 +LIST_WINDOW: Final = timedelta(hours=1) + + +def _spend_logs_page_queries(ids: Mapping[str, str], since: datetime) -> tuple[Mapping[str, str], ...]: + """``/spend/logs/ui`` and ``/spend/logs/v2`` require a window; keep it to this scenario.""" + window: Final = { + "start_date": (since - SCOPE_SLACK).strftime("%Y-%m-%d %H:%M:%S"), + "end_date": (since + LIST_WINDOW).strftime("%Y-%m-%d %H:%M:%S"), + "page_size": str(LIST_PAGE_SIZE), + } + return (*({**window, **query} for query in _scenario_filters(ids)), window) + + +SCENARIO_SCOPED_LIST_ROUTES: Final[ + Mapping[str, Callable[[Mapping[str, str], datetime], tuple[Mapping[str, str], ...]]] +] = MappingProxyType( + { + "/spend/logs": _spend_logs_queries, + "/spend/logs/ui": _spend_logs_page_queries, + "/spend/logs/v2": _spend_logs_page_queries, + } +) + + +@cache +def get_routes() -> tuple[str, ...]: + """Every GET route path on the proxy app, including routes hidden from the OpenAPI spec. + + The child imports the same source tree the owned proxy runs from (``INTEGRATION_PROXY_ROOT`` + or this checkout), without reading the database. Lazily registered feature routers + (``LAZY_FEATURES``) are loaded first, so their GET routes are enumerated too; on the running + proxy the first request to such a path registers the router before it is served. Mounted + ASGI sub-apps (the MCP server) have no methods and are out of scope for S2. + """ + script: Final = ( + "import asyncio, json\n" + "from litellm.proxy._lazy_features import LAZY_FEATURES, _force_load\n" + "from litellm.proxy.proxy_server import app\n" + "async def load():\n" + " for feature in LAZY_FEATURES:\n" + " await _force_load(app, feature)\n" + "asyncio.run(load())\n" + "paths = [getattr(r, 'path', '') for r in app.routes]\n" + "missing = sorted(f.name for f in LAZY_FEATURES if not any(f.matches(p) for p in paths))\n" + "print('MISSING=' + json.dumps(missing))\n" + "print('ROUTES=' + json.dumps(sorted({r.path for r in app.routes " + "if 'GET' in (getattr(r, 'methods', None) or ())})))\n" + ) + root: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3]) + inherited: Final = {name: value for name, value in os.environ.items() if name != "DATABASE_URL"} + completed: Final = subprocess.run( + [sys.executable, "-P", "-c", script], + cwd=root, + env={**inherited, "PYTHONPATH": os.pathsep.join((str(root), inherited.get("PYTHONPATH", "")))}, + capture_output=True, + text=True, + timeout=120, + check=True, + ) + lines: Final = completed.stdout.splitlines() + missing: Final = json.loads(next(line for line in lines if line.startswith("MISSING=")).removeprefix("MISSING=")) + routes: Final = tuple( + json.loads(next(line for line in lines if line.startswith("ROUTES=")).removeprefix("ROUTES=")) + ) + assert missing == [], f"Lazy features registered no route, so S2 cannot sweep them: {missing}" + assert "/spend/logs/ui/{request_id}" in routes, "Route enumeration missed hidden routes" + assert "/guardrails/list" in routes, "Route enumeration missed lazily registered feature routes" + return routes + + +def _route_ids(route: str, ids: Mapping[str, str]) -> Mapping[str, str]: + """``ids`` with the route's ``PARAMETER_ALIASES`` applied (``/v1/models/{model_id}`` takes a model name).""" + aliases: Final = PARAMETER_ALIASES.get(route, {}) + return {**ids, **{name: ids[source] for name, source in aliases.items() if source in ids}} + + +def _filled(route: str, ids: Mapping[str, str]) -> tuple[str, bool]: + """The concrete path, and whether any parameter fell back to ``PLACEHOLDER_ID``.""" + known: Final = {**DEFAULT_IDS, **_route_ids(route, ids)} + names: Final = _PATH_PARAMETER.findall(route) + path: Final = _PATH_PARAMETER.sub(lambda match: quote(known.get(match.group(1), PLACEHOLDER_ID), safe=""), route) + return path, any(name not in known for name in names) + + +@dataclass(frozen=True, slots=True) +class _RouteCall: + hits: tuple[Hit, ...] + allowed: tuple[Hit, ...] + error: str | None + unreachable: str | None + location: str = "" + status: int = 0 + + +def sweep_routes( + gateway: Gateway, + canaries: Sequence[Canary], + ids: Mapping[str, str], + *, + callers: Mapping[str, str] | None = None, + since: datetime | None = None, +) -> RouteSweep: + """S2: call every GET route as each caller (label -> bearer key; default the master key). + + With ``since``, the log list routes in ``SCENARIO_SCOPED_LIST_ROUTES`` are called with this + scenario's filters (its request id, its user, and a date window from ``since``) instead of + unfiltered, which on a shared database returns every row ever written or no rows at all. + """ + routes: Final = tuple(route for route in get_routes() if route_denied(route) is None) + targets: Final = tuple( + (route, *_filled(route, ids)) + for route in routes + if all(name in ids for name in _PATH_PARAMETER.findall(route) if name in REAL_ID_REQUIRED) + ) + who: Final = callers if callers is not None else {"admin": gateway.key} + base_url: Final = str(gateway.client.base_url) + + def call(route: str, label: str, key: str, path: str) -> _RouteCall: + location: Final = f"GET {path} as {label}" + try: + with httpx.Client(base_url=base_url, timeout=_ROUTE_TIMEOUT, trust_env=False) as client: + response = client.get(path, headers={"Authorization": f"Bearer {key}"}) + except httpx.HTTPError as error: + return _RouteCall((), (), None, f"{location}: {type(error).__name__}", location) + headers = "\n".join(f"{name}: {value}" for name, value in response.headers.items()) + found = _hits( + "S2", f"{location} -> {response.status_code}", response.content + b"\n" + headers.encode(), canaries + ) + return _RouteCall( + tuple(hit for hit in found if route_allowance(route, label, hit.slot) is None), + tuple(hit for hit in found if route_allowance(route, label, hit.slot) is not None), + f"{location}: {response.status_code}" if response.status_code >= 500 else None, + None, + location, + response.status_code, + ) + + jobs: Final = tuple( + (route, label, key, path + query) + for label, key in who.items() + for route, path, _ in targets + for query in _route_queries(route, ids, since) + ) + with ThreadPoolExecutor(max_workers=8) as pool: + results: Final = tuple(pool.map(lambda job: call(*job), jobs)) + supplied: Final = { + route + for route, _, _ in targets + if route not in NOT_FOUND_EXPECTED + and _PATH_PARAMETER.findall(route) + and all(name in _route_ids(route, ids) for name in _PATH_PARAMETER.findall(route)) + } + scoped: Final = {route for route in SCENARIO_SCOPED_LIST_ROUTES if since is not None} + return RouteSweep( + hits=tuple(hit for result in results for hit in result.hits), + called=tuple(f"{label} {path}" for _, label, _, path in jobs), + unfilled=( + *(route for route, _, placeholder in targets if placeholder), + *(route for route in routes if route not in {target for target, _, _ in targets}), + ), + errors=tuple(result.error for result in results if result.error is not None), + unreachable=tuple(result.unreachable for result in results if result.unreachable is not None), + allowed=tuple(hit for result in results for hit in result.allowed), + not_found=tuple( + f"{result.location} -> 404" + for (route, label, _, _), result in zip(jobs, results, strict=True) + if route in supplied and label == "admin" and result.status == 404 + ), + rejected=tuple( + f"{result.location} -> {result.status}" + for (route, _, _, _), result in zip(jobs, results, strict=True) + if route in scoped and 400 <= result.status < 500 + ), + statuses=MappingProxyType({result.location: result.status for result in results}), + ) + + +def record_route_sweep(routes: RouteSweep, node: str) -> None: + """Append the route sweep's errors and unfilled routes to the results directory, when set.""" + destination: Final = os.environ.get("INTEGRATION_RESULTS_DIR") + if not destination: + return + entry: Final = { + "node": node, + "called": len(routes.called), + "errors": routes.errors, + "unreachable": routes.unreachable, + "unfilled": routes.unfilled, + "allowed": [f"{hit.slot} {hit.location}" for hit in routes.allowed], + "not_found": routes.not_found, + "rejected": routes.rejected, + } + with (Path(destination) / "security-route-sweep.jsonl").open("a") as report: + report.write(json.dumps(entry) + "\n") + + +def sweep_responses(responses: Sequence[httpx.Response], canaries: Sequence[Canary]) -> tuple[Hit, ...]: + """S3: body and headers of each client-facing response.""" + found: Final[list[Hit]] = [] # mutable-ok: accumulated across responses + for index, response in enumerate(responses): + where = f"response[{index}] {response.request.method} {response.request.url.path} -> {response.status_code}" + found.extend(_hits("S3", where + " body", response.content, canaries)) + for name, value in response.headers.items(): + found.extend(_hits("S3", f"{where} header {name}", value, canaries)) + return tuple(found) + + +def sweep_sink( + name: str, + requests: Sequence[Request], + canaries: Sequence[Canary], + *, + own_header: tuple[str, str] | None = None, +) -> tuple[Hit, ...]: + """S4: every request a sink double received; ``own_header`` may carry its own canary only.""" + found: Final[list[Hit]] = [] # mutable-ok: accumulated across requests + for index, request in enumerate(requests): + where = f"{name}[{index}] {request.method} {request.target}" + found.extend(_hits("S4", where + " body", request.body, canaries)) + for header, value in request.headers.items(): + found.extend( + hit + for hit in _hits("S4", f"{where} header {header}", value, canaries) + if own_header is None or (header, hit.slot) != own_header + ) + return tuple(found) + + +def _redis_values(cache: Redis, key: bytes) -> Iterable[bytes]: + kind: Final = cache.type(key) + readers: Final[Mapping[bytes, Callable[[], Iterable[bytes]]]] = { + b"string": lambda: (cache.get(key) or b"",), + b"hash": lambda: (part for pair in cache.hgetall(key).items() for part in pair), + b"list": lambda: cache.lrange(key, 0, -1), + b"set": lambda: cache.smembers(key), + b"zset": lambda: cache.zrange(key, 0, -1), + } + reader: Final = readers.get(kind) + return reader() if reader is not None else () + + +def sweep_redis(canaries: Sequence[Canary], *, host: str | None = None, port: int | None = None) -> tuple[Hit, ...]: + """S5: every key name and value in the Redis database the proxy uses.""" + found: Final[list[Hit]] = [] # mutable-ok: accumulated across keys + with Redis( + host=host or os.environ["REDIS_HOST"], port=port or int(os.environ["REDIS_PORT"]), decode_responses=False + ) as cache: + for key in cache.scan_iter(count=500): + found.extend(_hits("S5", f"redis key {key!r}", key, canaries)) + for value in _redis_values(cache, key): + found.extend(_hits("S5", f"redis value {key!r}", value, canaries)) + return tuple(found) + + +@dataclass(frozen=True, slots=True) +class SweepReport: + hits: tuple[Hit, ...] + routes: RouteSweep + + def credential_hits(self) -> tuple[Hit, ...]: + return tuple(hit for hit in self.hits if hit.slot != MARKER) + + def marker_locations(self) -> tuple[tuple[str, str], ...]: + return tuple((hit.sweep, hit.location) for hit in self.hits if hit.slot == MARKER) + + +def sweep_all( + gateway: Gateway, + canaries: Sequence[Canary], + *, + responses: Sequence[httpx.Response], + sinks: Mapping[str, Sequence[Request]], + ids: Mapping[str, str], + callers: Mapping[str, str] | None = None, + own_headers: Mapping[str, tuple[str, str]] | None = None, + since: datetime | None = None, +) -> SweepReport: + """S1 to S5 for one finished scenario; fails if any GET route returned no response. + + Redis goes first: it holds entries with a TTL, and the route walk is the slow sweep. Pass + ``since`` (taken before the scenario's first request) to scope the append-only log tables + and the unpaginated log list routes to this scenario; the sensitivity marker's own spend-log + row must then still be found, which ``assert_marker_seen`` checks. + """ + redis: Final = sweep_redis(canaries) + routes: Final = sweep_routes(gateway, canaries, ids, callers=callers, since=since) + assert not routes.unreachable, f"GET routes returned no response, so S2 did not check them: {routes.unreachable}" + assert not routes.rejected, ( + f"Scoped list routes rejected the scenario's query, so S2 saw no rows: {routes.rejected}" + ) + assert not routes.not_found, ( + f"GET routes whose ids were all supplied answered 404 to the admin, so an id is wrong: {routes.not_found}" + ) + hits: Final = ( + *sweep_database(canaries, since=since), + *routes.hits, + *sweep_responses(responses, canaries), + *( + hit + for name, received in sinks.items() + for hit in sweep_sink(name, received, canaries, own_header=(own_headers or {}).get(name)) + ), + *redis, + ) + return SweepReport(hits, routes) + + +def assert_marker_seen(report: SweepReport, expected: Mapping[str, str]) -> None: + """Sensitivity control: the marker must be reported by each sweep at the expected location.""" + seen: Final = report.marker_locations() + missing: Final = tuple( + f"{sweep} at *{where}*" + for sweep, where in expected.items() + if not any(found_sweep == sweep and where in location for found_sweep, location in seen) + ) + assert not missing, f"Sweep could not see its surface, missing marker {missing}; marker seen at:\n" + "\n".join( + f" {sweep} {location}" for sweep, location in seen + ) diff --git a/tests/integration/security/test_config_deployment_key.py b/tests/integration/security/test_config_deployment_key.py new file mode 100644 index 00000000000..1b95d90e43b --- /dev/null +++ b/tests/integration/security/test_config_deployment_key.py @@ -0,0 +1,86 @@ +"""Slot B1: a deployment ``api_key`` declared in the proxy config reaches only the provider. + +Positive control: the provider double must receive ``Authorization: Bearer `` for +the scenario's request, or the test fails before sweeping. Sensitivity control: the marker sent +in the same request must be reported by the sweeps where stored prompts belong. Then no sweep +may find the B1 canary anywhere. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from datetime import UTC, datetime +from pathlib import Path +from typing import Final + +import pytest +from integration._support.client import string_value +from integration.security._canary import MARKER, canary +from integration.security._sinks import CONFIG_MODEL, GENERIC_SINK, PROVIDER_4XX, Rig, canary_rig, settle, team_caller +from integration.security._sweeps import assert_marker_seen, assert_no_hits, record_route_sweep, sweep_all + + +@pytest.fixture +def rig(tmp_path: Path) -> Iterator[Rig]: + """One owned proxy per test: B1 lives in the config, so a fresh core needs a fresh proxy.""" + with canary_rig(tmp_path) as value: + yield value + + +@pytest.mark.timeout(240) # full S1/S2 walk: every table and ~400 GET routes as two callers +@pytest.mark.parametrize("outcome", ["success", "provider_4xx"]) +def test_config_deployment_api_key_reaches_only_the_provider( + rig: Rig, outcome: str, request: pytest.FixtureRequest +) -> None: + b1: Final = rig.canaries["B1"] + marker: Final = canary(MARKER) + started: Final = datetime.now(UTC) + text: Final = f"slot B1 {marker.value}" + (f" {PROVIDER_4XX}" if outcome == "provider_4xx" else "") + with rig.proxy.scenario() as scenario: + caller: Final = team_caller(scenario) + response: Final = rig.proxy.request( + "POST", + "/v1/chat/completions", + {"model": CONFIG_MODEL, "messages": [{"role": "user", "content": text}]}, + key=caller.key, + ) + assert response.status_code == (200 if outcome == "success" else 400), response.text + delivered: Final = rig.provider.carrying(marker.value) + assert [request.headers.get("authorization") for request in delivered] == [f"Bearer {b1.value}"], ( + "Positive control: the provider double never received the B1 canary" + ) + request_id: Final = ( + string_value(response.json()["id"]) if outcome == "success" else response.headers["x-litellm-call-id"] + ) + settle(rig, request_id, marker) + + report: Final = sweep_all( + rig.proxy, + (marker, b1), + responses=(response,), + sinks={name: sink.requests() for name, sink in rig.sinks.items()}, + ids={ + "request_id": request_id, + "team_id": caller.team_id, + "user_id": caller.user_id, + "model_id": rig.model_id, + "model": CONFIG_MODEL, + }, + callers=caller.callers(rig), + own_headers=rig.own_headers, + since=started, + ) + record_route_sweep(report.routes, request.node.nodeid) + assert_marker_seen( + report, + { + "S1": "LiteLLM_SpendLogs.proxy_server_request", + "S2": f"GET /spend/logs/ui/{request_id} as admin -> 200", + "S4": f"{GENERIC_SINK}[", + }, + ) + by_model: Final = f"GET /credentials/by_model/{rig.model_id} as admin" + assert report.routes.statuses.get(by_model) == 200, ( + f"{by_model} must resolve the config deployment: {report.routes.statuses.get(by_model)}" + ) + assert_no_hits(report.credential_hits(), f"slot B1, {outcome}") diff --git a/tests/integration/security/test_sweep_sensitivity.py b/tests/integration/security/test_sweep_sensitivity.py new file mode 100644 index 00000000000..2a1340dc358 --- /dev/null +++ b/tests/integration/security/test_sweep_sensitivity.py @@ -0,0 +1,180 @@ +"""Sensitivity controls: every sweep must find a marker where prompts are legitimately stored. + +A sweep that cannot see its surface would pass every credential slot vacuously. Each test here +sends a fresh marker in message content with ``store_prompts_in_spend_logs`` on and requires +each sweep to report it at the place it belongs. +""" + +from __future__ import annotations + +import base64 +import gzip +import uuid +from collections.abc import Iterator +from datetime import UTC, datetime +from pathlib import Path +from typing import Final + +import pytest +from integration._support.client import eventually, string_value +from integration.security._canary import DECODE_BUDGET_BYTES, MARKER, SLOTS, DecodeBudgetExceeded, canary, find_canary +from integration.security._sinks import CONFIG_MODEL, GENERIC_SINK, Rig, canary_rig, settle, team_caller +from integration._support.wire import Request +from integration.security._sweeps import ( + ADMIN_ONLY_ALLOWANCES, + ALLOWANCE_SLOT_FAMILIES, + PROVIDER_PASSTHROUGH_REASON, + assert_marker_seen, + get_routes, + record_route_sweep, + route_allowance, + route_denied, + scoped_queries, + sweep_all, + sweep_redis, + sweep_sink, +) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with canary_rig(tmp_path_factory.mktemp("canary-sensitivity")) as value: + yield value + + +@pytest.mark.parametrize("prefix", ["", "u:", "us:", "use:"], ids=["align0", "align1", "align2", "align3"]) +def test_find_canary_decodes_base64_at_every_alignment_and_gzip(prefix: str) -> None: + marker: Final = canary(MARKER) + basic: Final = base64.b64encode(f"{prefix}{marker.value}".encode()).decode() + urlsafe: Final = base64.urlsafe_b64encode(f"{prefix}{marker.value}".encode()).decode().rstrip("=") + assert [match.slot for match in find_canary(f"Authorization: Basic {basic}", (marker,))] == [MARKER] + assert [match.slot for match in find_canary(f'{{"token":"{urlsafe}"}}', (marker,))] == [MARKER] + assert [match.slot for match in find_canary(gzip.compress(f"Basic {basic}".encode()), (marker,))] == [MARKER] + embedded: Final = b"prefix:" + gzip.compress(f"Basic {basic}".encode()) + b":suffix" + assert [match.slot for match in find_canary(embedded, (marker,))] == [MARKER] + members: Final = gzip.compress(b"first member") + gzip.compress(f"Basic {basic}".encode()) + assert [match.slot for match in find_canary(members, (marker,))] == [MARKER] + binary_wrapper: Final = bytes(range(256)) + f" Basic {basic} ".encode() + bytes(range(256)) + assert [match.slot for match in find_canary(base64.b64encode(binary_wrapper), (marker,))] == [MARKER] + assert find_canary(f"Basic {basic}".replace(basic[10:20], "A" * 10), (marker,)) == () + assert find_canary(f"sk-...{marker.core[-4:]}", (marker,)) == () + + +def test_find_canary_fails_loudly_past_its_decode_budget() -> None: + marker: Final = canary(MARKER) + bomb: Final = gzip.compress(b"\0" * (1024 * 1024 + 1)) + with pytest.raises(DecodeBudgetExceeded): + find_canary(bomb, (marker,), budget_bytes=1024 * 1024) + assert find_canary(gzip.compress(b"\0" * 1024) + marker.value.encode(), (marker,), budget_bytes=1024 * 1024) + assert DECODE_BUDGET_BYTES >= 256 * 1024 * 1024 + + +def test_rig_with_an_overridden_master_key_resolves_the_config_deployment(tmp_path: Path) -> None: + master_key: Final = f"sk-canary-override-{uuid.uuid4().hex}" + with canary_rig(tmp_path, environment={"LITELLM_MASTER_KEY": master_key}) as overridden: + assert overridden.proxy.key == master_key + assert overridden.model_id + assert overridden.proxy.request("GET", "/model/info").status_code == 200 + + +def test_route_allowances_match_only_their_exact_route_and_caller() -> None: + routes: Final = get_routes() + callers: Final = ("admin", "internal_user", "Admin", "admin ", "") + for route, caller in ADMIN_ONLY_ALLOWANCES: + assert route in routes, f"Allowance names a route the proxy no longer registers: {route}" + for variant in (route + "/", route.upper(), route.rstrip("s"), "/v1" + route): + assert route_allowance(variant, caller) is None, variant + allowed: Final = {(route, caller) for route in routes for caller in callers if route_allowance(route, caller)} + assert allowed == set(ADMIN_ONLY_ALLOWANCES), allowed + assert all(route_denied(route) is None for route, _ in ADMIN_ONLY_ALLOWANCES) + assert set(ALLOWANCE_SLOT_FAMILIES) == set(ADMIN_ONLY_ALLOWANCES) + for (route, caller), families in ALLOWANCE_SLOT_FAMILIES.items(): + for family in families: + assert route_allowance(route, caller, family + "1") is not None + for slot in SLOTS: + if not slot.startswith(families): + assert route_allowance(route, caller, slot) is None, (route, caller, slot) + + +def test_only_provider_passthrough_routes_match_the_passthrough_deny_rule() -> None: + denied: Final = {route for route in get_routes() if route_denied(route) == PROVIDER_PASSTHROUGH_REASON} + assert "/openai/{endpoint:path}" in denied and "/langfuse/{endpoint:path}" in denied + assert all(route.endswith("/{endpoint:path}") and route.count("{") == 1 for route in denied), denied + for swept in ("/v1/files/{file_id:path}", "/spend/logs/ui/{request_id}", "/v1/memory/{key:path}"): + assert route_denied(swept) is None, swept + + +def test_sink_own_header_allows_only_that_header_and_slot() -> None: + own: Final = canary("B1") + other: Final = canary(MARKER) + request: Final = Request( + "POST", + "/", + {"authorization": f"Bearer {own.value}", "x-extra": f"Bearer {own.value}", "x-other": other.value}, + f'{{"copied": "{own.value}"}}'.encode(), + ) + hits: Final = sweep_sink("double", (request,), (own, other), own_header=("authorization", own.slot)) + assert {(hit.slot, hit.location) for hit in hits} == { + (MARKER, "double[0] POST / header x-other"), + ("B1", "double[0] POST / body"), + ("B1", "double[0] POST / header x-extra"), + } + + +@pytest.mark.timeout(240) # full S1/S2 walk: every table and ~400 GET routes as two callers +def test_every_sweep_finds_the_stored_prompt_marker(rig: Rig, request: pytest.FixtureRequest) -> None: + marker: Final = canary(MARKER) + started: Final = datetime.now(UTC) + with rig.proxy.scenario() as scenario: + caller: Final = team_caller(scenario) + response: Final = rig.proxy.request( + "POST", + "/v1/chat/completions", + {"model": CONFIG_MODEL, "messages": [{"role": "user", "content": f"sensitivity {marker.value}"}]}, + key=caller.key, + ) + assert response.status_code == 200, response.text + assert len(rig.provider.carrying(marker.value)) == 1 + request_id: Final = string_value(response.json()["id"]) + settle(rig, request_id, marker) + eventually(lambda: sweep_redis((marker,)), bool, seconds=10) + + ids: Final = { + "request_id": request_id, + "team_id": caller.team_id, + "user_id": caller.user_id, + "model_id": rig.model_id, + "model": CONFIG_MODEL, + } + report: Final = sweep_all( + rig.proxy, + (marker,), + responses=(response,), + sinks={name: sink.requests() for name, sink in rig.sinks.items()}, + ids=ids, + callers=caller.callers(rig), + own_headers=rig.own_headers, + since=started, + ) + record_route_sweep(report.routes, request.node.nodeid) + assert_marker_seen( + report, + { + "S1": "LiteLLM_SpendLogs.proxy_server_request", + "S2": f"GET /spend/logs/ui/{request_id} as admin -> 200", + "S3": "response[0] POST /v1/chat/completions -> 200 body", + "S4": f"{GENERIC_SINK}[", + "S5": "redis value", + }, + ) + assert_marker_seen(report, {"S2": f"GET /spend/logs?request_id={request_id} as admin -> 200"}) + assert_marker_seen(report, {"S2": f"GET /spend/logs?user_id={caller.user_id} as admin -> 200"}) + assert_marker_seen(report, {"S2": f"GET /spend/logs/ui/{request_id} as internal_user -> 200"}) + for route in ("/spend/logs/ui", "/spend/logs/v2"): + filtered = tuple(query for query in scoped_queries(route, ids, started) if "_id=" in query) + assert len(filtered) == 2, filtered + for query in filtered: + assert report.routes.statuses.get(f"GET {route}{query} as admin") == 200, (route, query) + listed = rig.proxy.request("GET", route + query) + assert request_id in listed.text, f"{route}{query} does not list the scenario's row" + assert report.credential_hits() == ()