"""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() == ()