mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
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
This commit is contained in:
parent
5df502b360
commit
dab2deb5ed
8 changed files with 1349 additions and 1 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -17,5 +17,6 @@ OWNED_DIRECTORIES: Final = frozenset(
|
|||
"compatibility",
|
||||
"sdk",
|
||||
"cost_calculation",
|
||||
"security",
|
||||
}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ GROUPS: Final = MappingProxyType(
|
|||
"mcp": ("mcp",),
|
||||
"sdk": ("sdk",),
|
||||
"cost": ("cost_calculation",),
|
||||
"security": ("security",),
|
||||
}
|
||||
)
|
||||
|
||||
|
|
|
|||
168
tests/integration/security/_canary.py
Normal file
168
tests/integration/security/_canary.py
Normal file
|
|
@ -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 ``<field prefix>lkc-<slot id>-<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())
|
||||
297
tests/integration/security/_sinks.py
Normal file
297
tests/integration/security/_sinks.py
Normal file
|
|
@ -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 ``<provider>/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)
|
||||
615
tests/integration/security/_sweeps.py
Normal file
615
tests/integration/security/_sweeps.py
Normal file
|
|
@ -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 "<schema>"."<table>" 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 ``/<provider>/{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
|
||||
)
|
||||
86
tests/integration/security/test_config_deployment_key.py
Normal file
86
tests/integration/security/test_config_deployment_key.py
Normal file
|
|
@ -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 <B1 canary>`` 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}")
|
||||
180
tests/integration/security/test_sweep_sensitivity.py
Normal file
180
tests/integration/security/test_sweep_sensitivity.py
Normal file
|
|
@ -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() == ()
|
||||
Loading…
Add table
Reference in a new issue