diff --git a/docker/prod_entrypoint.sh b/docker/prod_entrypoint.sh index 28d1bdcc2942..3abfc31f60b1 100644 --- a/docker/prod_entrypoint.sh +++ b/docker/prod_entrypoint.sh @@ -6,9 +6,16 @@ if [ "$SEPARATE_HEALTH_APP" = "1" ]; then exec supervisord -c /etc/supervisord.conf fi +# Internal-fork modification: invoke `python -m litellm_extras.entrypoint` +# instead of the upstream `litellm` console script. The wrapper installs +# PublicReqMiddleware on the FastAPI app before delegating to the original +# Click CLI, which is what enforces the X-Public-Req gating in production. +# Routing through the wrapper at the entrypoint level (rather than via +# docker-compose overrides) means every deployment of this image picks +# the middleware up automatically — there is no per-deploy step to forget. if [ "$USE_DDTRACE" = "true" ]; then export DD_TRACE_OPENAI_ENABLED="False" - exec ddtrace-run litellm "$@" + exec ddtrace-run python -m litellm_extras.entrypoint "$@" else - exec litellm "$@" + exec python -m litellm_extras.entrypoint "$@" fi \ No newline at end of file diff --git a/e2e/_config/docker-compose.yml b/e2e/_config/docker-compose.yml index 7df49baa6f0e..3477fd99f629 100644 --- a/e2e/_config/docker-compose.yml +++ b/e2e/_config/docker-compose.yml @@ -70,6 +70,11 @@ services: # Rendered at proxy-start time by `e2e/tools/proxy` from .env values. # See ../tools/proxy `render_config()`. The file is gitignored. - ./.litellm.rendered.yaml:/app/config.yaml:ro + # NOTE: no `entrypoint:` override here. The fork's + # docker/prod_entrypoint.sh already routes through + # `python -m litellm_extras.entrypoint`, so PublicReqMiddleware loads + # automatically. This keeps e2e behavior aligned with production — + # if the wrapper breaks in either path, both surfaces catch it. command: - --config=/app/config.yaml - --port=4000 diff --git a/e2e/cases/18_public_req_middleware.md b/e2e/cases/18_public_req_middleware.md new file mode 100644 index 000000000000..cc6fd529b767 --- /dev/null +++ b/e2e/cases/18_public_req_middleware.md @@ -0,0 +1,120 @@ +# Case 18 — `X-Public-Req` middleware: streaming + header / query gating + +## Goal + +Verify the `PublicReqMiddleware` (installed via +`litellm_extras.entrypoint`) correctly differentiates external vs internal +requests **without buffering streaming bodies**. + +Two-pronged check: + +1. **Streaming integrity.** SSE chunks arrive incrementally; TTFB is + small relative to total wall time. This is the regression guard + against accidentally switching to `BaseHTTPMiddleware` or otherwise + awaiting on body messages — both would collapse TTFB onto end-of-stream. +2. **Differential behavior.** `X-Public-Req: 1` strips `x-litellm-*` + from response headers and blocks `/v1/models?include_metadata=true`. + The same requests without the header pass through unchanged. + +Functional logic (inbound header strip, case-insensitive matching, +forbidden-query parsing, non-HTTP scope handling) is covered by +`tests/test_litellm/test_public_req_middleware.py`. This e2e adds the +wire-level confirmation. + +## Origin + +We expose the LiteLLM proxy to external users behind a public Nginx +ingress. The ingress injects `X-Public-Req: 1` on every forwarded +request; internal services reach the proxy on a separate Service that +does not pass through the ingress. The middleware sits in the proxy's +ASGI stack and applies safeguards conditional on that header. + +## Preconditions + +- `e2e/tools/proxy status` reports `ready` +- The proxy image was built **after** `litellm_extras/` was added — + run `e2e/tools/proxy rebuild` if unsure +- `ANTHROPIC_API_KEY` set (one ~50-word streaming completion; + cost ≈ $0.002) + +## Steps + +```bash +bash e2e/cases/data/18_public_req_middleware.sh +echo "exit=$?" +``` + +The fixture executes six assertions. Each prints `PASS:` or `FAIL:`: + +### Streaming integrity (paid) + +- **A1** SSE response yields ≥ 3 `data:` chunks +- **A2** Streaming phase (wall − TTFB) > 200 ms — chunks spread over + time rather than dumped at end-of-stream. Provider TTFT variance can + push the ratio above 0.5 on short completions, so the absolute + duration is more robust than a ratio threshold; a buffering + middleware would collapse this to ≤ 5 ms regardless of provider. + +### Header gating (paid) + +- **A3** Public mode: zero `x-litellm-*` headers in response +- **A4** Internal mode (no `X-Public-Req`): at least one + `x-litellm-*` header in response (control) + +### Query gating on `/v1/models` (free) + +- **A5** Public mode: `GET /v1/models?include_metadata=true` returns + HTTP 200 with the **same** body as `GET /v1/models` (forbidden query + silently stripped before reaching the proxy). +- **A6** Internal mode: same request returns HTTP 200 with a body that + **differs** from the bare-models response (metadata expansion still + works for internal callers). + +### Inbound `x-litellm-*` strip (free — uses `mock_response`) + +Two chat completions are made with `mock_response: "pong"` (no upstream +provider call). Each carries `X-Litellm-Spend-Logs-Metadata` with a +unique marker. LiteLLM persists honored header values to +`metadata.spend_logs_metadata` in the spend_logs row, so the marker's +presence/absence in Postgres is a definitive signal of whether the +header reached the proxy core. + +- **A7** Internal mode (control): marker **must** appear in spend_logs + within 15 s (proves the header would otherwise be honored) +- **A8** Public mode (`X-Public-Req: 1`): marker **must not** appear in + any spend_logs row (proves the middleware stripped the header before + LiteLLM saw it) + +## Expected — GREEN + +``` +A1 PASS: 7 SSE chunks received +A2 PASS: ttfb=8294ms wall=10171ms streaming_phase=1877ms +A3 PASS: 0 x-litellm-* headers in public response +A4 PASS: 7 x-litellm-* headers in internal response +A5 PASS: /v1/models?include_metadata=true returned 200 with query stripped (public) +A6 PASS: /v1/models?include_metadata=true returned 200 with metadata expanded (internal) +A7 PASS: internal mode: x-litellm-spend-logs-metadata reached spend_logs (1 row) +A8 PASS: public mode: inbound x-litellm-spend-logs-metadata stripped (marker absent from spend_logs) +``` + +## Failure modes + +| Symptom | Likely cause | +|---|---| +| A1 fails with 0-1 chunks | Middleware is buffering — switched to `BaseHTTPMiddleware`, or `send_wrapper` is awaiting on body | +| A2 `streaming_phase ≤ 5ms` | Middleware buffered the body and dumped it on close — same root cause as A1 | +| A3 fails (still see `x-litellm-*` in public response) | `send_wrapper` not wired, or middleware not installed; check `proxy logs` for `PublicReqMiddleware` | +| A4 fails (no `x-litellm-*` in internal) | Middleware is stripping for *all* requests; check `_is_public` returns False without the header | +| A5 fails — public bodies differ between `?include_metadata=true` and bare `/v1/models` | Query strip did not run — `/v1/models` not in `MODELS_PATHS`, middleware not installed, or `parse_qsl`/`urlencode` lost the rewrite | +| A5 fails — public returns 4xx | Middleware reverted to the old reject-with-400 behavior; revert the strip refactor | +| A6 fails — internal bodies identical | Middleware running for internal calls — `_is_public` defaulting to True; strip is happening when it shouldn't | +| A7 fails — internal marker missing from spend_logs | Async spend logger lag or DB schema drift — not a middleware bug. Raise the 15 s poll if reproducible | +| A8 fails — public marker present in spend_logs | Inbound `x-litellm-*` strip is NOT running. Verify the middleware is installed and that `LITELLM_HEADER_PREFIX` matching is case-insensitive | + +## Cross-reference + +- `litellm_extras/public_req_middleware.py` — middleware under test +- `litellm_extras/entrypoint.py` — wrapper that installs the middleware +- `e2e/_config/docker-compose.yml` — `command:` invokes the wrapper +- `tests/test_litellm/test_public_req_middleware.py` — functional unit tests diff --git a/e2e/cases/README.md b/e2e/cases/README.md index 4e1d487bc96c..f7135f4ed35c 100644 --- a/e2e/cases/README.md +++ b/e2e/cases/README.md @@ -27,6 +27,7 @@ Humans can execute them too — every step is a concrete shell command. | 15 | `15_v1_models_user_filter.md` | (none — proxy only) | `GET /v1/models` honors `LiteLLM_UserTable.models` (Personal Models). Regression for BerriAI/litellm#26420 | ✓ | | 16 | `16_budget_reset_no_prisma_error.md` | (none — proxy only) | `ResetBudgetJob.reset_budget_windows` background tick must not raise `prisma.errors.MissingRequiredValueError` on `Json?` null-filter. Regression for BerriAI/litellm#26346 | ✓ | | 17 | `17_model_info_user_filter.md` | (none — proxy only) | `GET /v1/model/info` (Path B) and `GET /v2/model/info` (every flag combo) honor `LiteLLM_UserTable.models`. Extends PR #10 fix from `/v1/models` to the two info endpoints | ✓ | +| 18 | `18_public_req_middleware.md` | Anthropic | `litellm_extras.PublicReqMiddleware` keeps streaming responses incremental, strips `x-litellm-*` under `X-Public-Req: 1`, and rejects sensitive `/v1/models` query params | — | ## How to invoke diff --git a/e2e/cases/data/18_public_req_middleware.sh b/e2e/cases/data/18_public_req_middleware.sh new file mode 100755 index 000000000000..a015e326b513 --- /dev/null +++ b/e2e/cases/data/18_public_req_middleware.sh @@ -0,0 +1,232 @@ +#!/usr/bin/env bash +# Case 18 fixture — see e2e/cases/18_public_req_middleware.md +# +# Asserts that the PublicReqMiddleware installed via +# litellm_extras.entrypoint: +# - keeps streaming responses incremental (A1, A2) +# - strips x-litellm-* response headers under X-Public-Req: 1 (A3) +# - leaves x-litellm-* response headers intact otherwise (A4) +# - rejects sensitive query params on /v1/models in public mode (A5) +# - accepts them in internal mode (A6) +# +# Exits 0 on PASS, 77 on SKIP (missing API key), anything else on FAIL. + +set -u + +PROXY="${PROXY_URL:-http://localhost:4011}" +DB_CONTAINER="${DB_CONTAINER:-litellm-e2e-db}" +DB_USER="${DB_USER:-litellm}" +DB_NAME="${DB_NAME:-litellm}" +KEY="${MASTER_KEY:-sk-e2e-test}" +MODEL="${MODEL_E2E_NAME:-claude-sonnet-cache}" +FAILED=0 + +pass() { echo "PASS: $*"; } +fail() { echo "FAIL: $*"; FAILED=1; } +skip() { echo "SKIP: $*"; exit 77; } + +# ---- precondition: Anthropic key wired? ---------------------------------- +# The proxy reads ANTHROPIC_API_KEY at startup; if we don't have it the +# streaming assertions cannot run. (/v1/models assertions would still work +# but the fixture is treated as a unit.) +if ! curl -sSf -o /dev/null -m 3 "$PROXY/health/readiness"; then + skip "proxy not ready at $PROXY" +fi + +# ---- shared streaming request -------------------------------------------- +# A single Anthropic streaming call serves A1+A2+A3. +PUB_BODY=$(mktemp); PUB_HDRS=$(mktemp); PUB_TIMING=$(mktemp) +trap 'rm -f "$PUB_BODY" "$PUB_HDRS" "$PUB_TIMING" "$INT_HDRS" 2>/dev/null' EXIT + +curl -sN -D "$PUB_HDRS" -o "$PUB_BODY" \ + -w "%{time_starttransfer} %{time_total} %{http_code}\n" \ + -H "Authorization: Bearer $KEY" \ + -H "X-Public-Req: 1" \ + -H "Content-Type: application/json" \ + -d "{ + \"model\":\"$MODEL\", + \"messages\":[{\"role\":\"user\",\"content\":\"Write a 50-word essay about clouds.\"}], + \"max_tokens\":200, + \"stream\":true + }" "$PROXY/v1/chat/completions" > "$PUB_TIMING" + +read TTFB_S WALL_S STATUS < "$PUB_TIMING" +if [ "$STATUS" != "200" ]; then + # Upstream auth failure → treat as SKIP not FAIL (no API key configured). + if [ "$STATUS" = "401" ] || [ "$STATUS" = "403" ]; then + skip "upstream returned $STATUS — ANTHROPIC_API_KEY likely missing" + fi + fail "streaming POST returned HTTP $STATUS" + echo "--- response body (truncated) ---" + head -c 400 "$PUB_BODY" + echo + exit 1 +fi + +TTFB_MS=$(awk "BEGIN{printf \"%d\", $TTFB_S*1000}") +WALL_MS=$(awk "BEGIN{printf \"%d\", $WALL_S*1000}") + +# ---- A1: stream produced >= 3 SSE chunks --------------------------------- +CHUNKS=$(grep -c '^data:' "$PUB_BODY" || true) +if [ "$CHUNKS" -ge 3 ]; then + pass "$CHUNKS SSE chunks received" +else + fail "only $CHUNKS SSE chunks (need >=3) — middleware may be buffering" +fi + +# ---- A2: streaming phase visible (wall - ttfb > 200ms) ------------------- +# +# A buffering middleware would emit ALL chunks at end-of-stream, collapsing +# (wall - ttfb) toward 0. Provider TTFT variance can push the *ratio* above +# 0.5 on short completions (slow first token + fast last tokens), so we use +# the absolute streaming-phase duration instead of a ratio. Even ~200 ms of +# streaming phase across multiple SSE chunks is unambiguous evidence that +# the middleware did not buffer. +STREAM_PHASE_MS=$(awk "BEGIN{printf \"%d\", ($WALL_S - $TTFB_S) * 1000}") +if [ "$STREAM_PHASE_MS" -gt 200 ]; then + pass "ttfb=${TTFB_MS}ms wall=${WALL_MS}ms streaming_phase=${STREAM_PHASE_MS}ms" +else + fail "ttfb=${TTFB_MS}ms wall=${WALL_MS}ms streaming_phase=${STREAM_PHASE_MS}ms (need >200ms; buffering suspected)" +fi + +# ---- A3: public response has zero x-litellm-* headers ------------------- +PUB_LITELLM_COUNT=$(grep -ic '^x-litellm-' "$PUB_HDRS" || true) +if [ "$PUB_LITELLM_COUNT" -eq 0 ]; then + pass "0 x-litellm-* headers in public response" +else + fail "$PUB_LITELLM_COUNT x-litellm-* headers leaked (expected 0)" + grep -i '^x-litellm-' "$PUB_HDRS" | sed 's/^/ /' +fi + +# ---- A4: internal response keeps x-litellm-* (control) ------------------ +INT_HDRS=$(mktemp) +curl -sN -D "$INT_HDRS" -o /dev/null \ + -H "Authorization: Bearer $KEY" \ + -H "Content-Type: application/json" \ + -d "{ + \"model\":\"$MODEL\", + \"messages\":[{\"role\":\"user\",\"content\":\"hi\"}], + \"max_tokens\":5, + \"stream\":true + }" "$PROXY/v1/chat/completions" + +INT_LITELLM_COUNT=$(grep -ic '^x-litellm-' "$INT_HDRS" || true) +if [ "$INT_LITELLM_COUNT" -ge 1 ]; then + pass "$INT_LITELLM_COUNT x-litellm-* headers in internal response" +else + fail "no x-litellm-* in internal response — middleware over-strips" +fi + +# ---- A5: public /v1/models — forbidden query silently stripped --------- +# Strategy: compare bodies for `?include_metadata=true` in public mode vs +# internal mode. In public mode the middleware drops the parameter before +# the proxy sees it, so the response must match the internal `no metadata` +# baseline — neither expanded with fallback chains nor a 4xx error. +PUB_NO_META=$(mktemp); PUB_WITH_META=$(mktemp); INT_NO_META=$(mktemp); INT_WITH_META=$(mktemp) + +S=$(curl -s -o "$PUB_WITH_META" -w "%{http_code}" \ + -H "Authorization: Bearer $KEY" \ + -H "X-Public-Req: 1" \ + "$PROXY/v1/models?include_metadata=true") +S_PUB_NOMETA=$(curl -s -o "$PUB_NO_META" -w "%{http_code}" \ + -H "Authorization: Bearer $KEY" \ + -H "X-Public-Req: 1" \ + "$PROXY/v1/models") +S_INT_META=$(curl -s -o "$INT_WITH_META" -w "%{http_code}" \ + -H "Authorization: Bearer $KEY" \ + "$PROXY/v1/models?include_metadata=true") +S_INT_NOMETA=$(curl -s -o "$INT_NO_META" -w "%{http_code}" \ + -H "Authorization: Bearer $KEY" \ + "$PROXY/v1/models") + +if [ "$S" = "200" ] && [ "$S_PUB_NOMETA" = "200" ] \ + && diff -q "$PUB_WITH_META" "$PUB_NO_META" >/dev/null; then + pass "/v1/models?include_metadata=true returned 200 with query stripped (public)" +else + fail "/v1/models?include_metadata=true public status=$S vs no-meta=$S_PUB_NOMETA; bodies differ → strip failed" + diff "$PUB_WITH_META" "$PUB_NO_META" | head -5 +fi + +# ---- A6: internal /v1/models — metadata still expanded ------------------ +# Internal mode (no X-Public-Req) must NOT strip. The metadata-enriched +# response must differ from the bare-models response. +if [ "$S_INT_META" = "200" ] && [ "$S_INT_NOMETA" = "200" ] \ + && ! diff -q "$INT_WITH_META" "$INT_NO_META" >/dev/null; then + pass "/v1/models?include_metadata=true returned 200 with metadata expanded (internal)" +else + fail "internal status: with-meta=$S_INT_META no-meta=$S_INT_NOMETA; bodies identical → middleware over-strips" +fi + +rm -f "$PUB_NO_META" "$PUB_WITH_META" "$INT_NO_META" "$INT_WITH_META" + +# ---- A7/A8: inbound x-litellm-* strip vs preserve ------------------------ +# Send two mock_response chat completions (free, no provider call) each +# carrying x-litellm-spend-logs-metadata with a unique marker. The marker +# only reaches LiteLLM's pre-call hooks if the header was honored — i.e., +# the proxy's spend_logs row will contain it under +# metadata.spend_logs_metadata.case18_marker. +# +# Expected: +# A7 (public + X-Public-Req: 1): row exists, marker ABSENT → strip ran +# A8 (internal, no X-Public-Req): row exists, marker PRESENT → control +PUB_MARKER="case18-pub-$(date +%s%N)" +INT_MARKER="case18-int-$(date +%s%N)" + +# A7: public call — header must be stripped by the middleware +curl -sS -o /dev/null \ + -H "Authorization: Bearer $KEY" \ + -H "X-Public-Req: 1" \ + -H "X-Litellm-Spend-Logs-Metadata: {\"case18_marker\":\"$PUB_MARKER\"}" \ + -H "Content-Type: application/json" \ + -d "{ + \"model\":\"$MODEL\", + \"messages\":[{\"role\":\"user\",\"content\":\"ping\"}], + \"mock_response\":\"pong\" + }" "$PROXY/v1/chat/completions" + +# A8: internal call (control) — header must be honored +curl -sS -o /dev/null \ + -H "Authorization: Bearer $KEY" \ + -H "X-Litellm-Spend-Logs-Metadata: {\"case18_marker\":\"$INT_MARKER\"}" \ + -H "Content-Type: application/json" \ + -d "{ + \"model\":\"$MODEL\", + \"messages\":[{\"role\":\"user\",\"content\":\"ping\"}], + \"mock_response\":\"pong\" + }" "$PROXY/v1/chat/completions" + +# Async spend logger flush — poll up to 15s for the internal-marker row to +# appear, then make the public-marker assertion. If the internal row never +# appears the logger is backed up and the absence of the public marker is +# not yet proof of strip; in that case we surface a warning. +INT_FOUND=0 +for _ in $(seq 1 15); do + sleep 1 + INT_CNT=$(docker exec "$DB_CONTAINER" psql -U "$DB_USER" -d "$DB_NAME" -tA -c " +SELECT COUNT(*) FROM \"LiteLLM_SpendLogs\" +WHERE metadata::text LIKE '%$INT_MARKER%'; +" 2>/dev/null | tr -d ' ') + if [ "${INT_CNT:-0}" -ge 1 ]; then + INT_FOUND=1 + break + fi +done + +if [ "$INT_FOUND" -ne 1 ]; then + fail "A8 control: internal marker '$INT_MARKER' never reached spend_logs within 15s — async logger backed up?" +else + pass "internal mode: x-litellm-spend-logs-metadata reached spend_logs ($INT_CNT row)" +fi + +PUB_CNT=$(docker exec "$DB_CONTAINER" psql -U "$DB_USER" -d "$DB_NAME" -tA -c " +SELECT COUNT(*) FROM \"LiteLLM_SpendLogs\" +WHERE metadata::text LIKE '%$PUB_MARKER%'; +" 2>/dev/null | tr -d ' ') + +if [ "${PUB_CNT:-0}" -eq 0 ]; then + pass "public mode: inbound x-litellm-spend-logs-metadata stripped (marker absent from spend_logs)" +else + fail "public mode: inbound x-litellm-spend-logs-metadata LEAKED — $PUB_CNT spend_logs row(s) carry marker '$PUB_MARKER'" +fi + +exit $FAILED diff --git a/e2e/tools/run-all-cases b/e2e/tools/run-all-cases index 3c69762bf613..1a3f2f908518 100755 --- a/e2e/tools/run-all-cases +++ b/e2e/tools/run-all-cases @@ -372,6 +372,21 @@ case_17() { fi } +# ------------------------------------------------------------------ 18 +case_18() { + echo "[18] X-Public-Req middleware streaming + header/query gating..." + local out=/tmp/e2e_case18.out + bash e2e/cases/data/18_public_req_middleware.sh > "$out" 2>&1 + local rc=$? + if [ "$rc" -eq 0 ]; then + ok "18 public-req-middleware: $(grep -c '^PASS' "$out") assertions PASS" + elif [ "$rc" -eq 77 ]; then + skip "18 public-req-middleware" "$(grep -m1 SKIP "$out")" + else + fail "18 public-req-middleware" "$(grep -m1 '^FAIL' "$out" || tail -2 "$out")" + fi +} + # Pre-flight: proxy must be ready if ! curl -sSL --max-time 3 -o /dev/null -w '%{http_code}' \ "$PROXY/health/readiness" 2>/dev/null | grep -q '^2'; then @@ -413,6 +428,11 @@ case_12 case_15 case_16 case_17 +if [ "$SKIP_PAID" -eq 1 ]; then + skip "18 public-req-middleware" "skipped (--skip-paid)" +else + case_18 +fi echo echo "============ SUMMARY ============" diff --git a/litellm_extras/__init__.py b/litellm_extras/__init__.py new file mode 100644 index 000000000000..3edddae54dee --- /dev/null +++ b/litellm_extras/__init__.py @@ -0,0 +1,7 @@ +"""Internal LiteLLM extensions kept out of the upstream tree. + +This package hosts deployment-specific middleware and entrypoints layered on +top of the vendored LiteLLM source. Modules here MUST NOT be imported by any +file under ``litellm/`` itself — that would couple upstream code to internal +add-ons and complicate future rebases. +""" diff --git a/litellm_extras/entrypoint.py b/litellm_extras/entrypoint.py new file mode 100644 index 000000000000..60c3ecead50e --- /dev/null +++ b/litellm_extras/entrypoint.py @@ -0,0 +1,51 @@ +"""Wrapper entrypoint that installs internal middleware before the LiteLLM CLI runs. + +Invoke this module in place of the plain ``litellm`` console script:: + + python -m litellm_extras.entrypoint --config=/app/config.yaml --port=4000 + +It works by importing :mod:`litellm.proxy.proxy_server` first (which has the +side effect of constructing the FastAPI ``app`` instance), attaching +:class:`~litellm_extras.public_req_middleware.PublicReqMiddleware`, and only +then delegating to the existing Click-based ``run_server`` CLI. Because +Python caches imported modules, the subsequent ``from .proxy_server import +app`` inside ``run_server`` returns the same ``app`` object with our +middleware already inserted. + +``FastAPI.add_middleware`` prepends to ``app.user_middleware``, so calling it +after the proxy's own ``add_middleware`` lines puts ours on the OUTSIDE — +which is what we want so the header strip runs before any LiteLLM +auth/logging layer. +""" + +import sys + +from litellm.proxy import proxy_server # noqa: E402 (load order is intentional) +from litellm.proxy.proxy_cli import run_server # noqa: E402 + +from litellm_extras.public_req_middleware import PublicReqMiddleware + + +def install_middleware() -> None: + """Attach internal middleware to the proxy's FastAPI app. + + Idempotent: if PublicReqMiddleware is already present we skip the + insert. This makes the entrypoint safe to invoke from tests that may + have imported ``litellm.proxy.proxy_server`` earlier. + """ + existing = { + m.cls.__name__ for m in proxy_server.app.user_middleware if hasattr(m, "cls") + } + if PublicReqMiddleware.__name__ in existing: + return + proxy_server.app.add_middleware(PublicReqMiddleware) + + +def main() -> int: + install_middleware() + run_server.main(args=sys.argv[1:], standalone_mode=False) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/litellm_extras/public_req_middleware.py b/litellm_extras/public_req_middleware.py new file mode 100644 index 000000000000..d636f22c1adc --- /dev/null +++ b/litellm_extras/public_req_middleware.py @@ -0,0 +1,134 @@ +"""ASGI middleware that gates public-vs-internal request behavior. + +The public Nginx ingress injects ``X-Public-Req: 1`` on every request it +forwards. Internal callers reach the proxy directly (e.g. via the in-cluster +Service ``litellm-internal``) and never carry that header. This middleware +keys off the header to apply two safeguards for public requests only: + +1. Strip every ``x-litellm-*`` request header — these headers (notably + ``x-litellm-api-key``, ``x-litellm-mock-response``, + ``x-litellm-num-retries``) let callers override internal proxy behavior + and must never be honored from untrusted sources. + +2. Silently strip sensitive query parameters from ``/v1/models`` and + ``/models`` requests (``include_metadata``, ``fallback_type``, + ``include_model_access_groups``, ``only_model_access_groups``). These + parameters expose router fallback chains and access-group naming, + which are deployment-internal details. The request still reaches the + inner app and returns a normal 200 with the redacted view, so naive + public clients that always set these params do not break. + +For public responses the middleware additionally strips every +``x-litellm-*`` response header (model deployment IDs, cache-hit flags, +cost/budget annotations) before they leave the proxy. + +Internal requests (no header / header != ``"1"``) are passed through +untouched so internal services can keep using the override headers and +observability fields. + +The middleware is pure-ASGI (not Starlette ``BaseHTTPMiddleware``). It does +not buffer response bodies, so streaming endpoints (chat completions with +``stream=true``, ``/v1/messages``, ``/v1/realtime``) keep their original +time-to-first-byte profile. +""" + +from typing import Iterable, List, Optional, Tuple +from urllib.parse import parse_qsl, urlencode + +from starlette.types import ASGIApp, Message, Receive, Scope, Send + +LITELLM_HEADER_PREFIX = b"x-litellm-" +PUBLIC_REQ_HEADER = b"x-public-req" +PUBLIC_REQ_VALUE = b"1" + +MODELS_PATHS = frozenset({"/v1/models", "/models"}) +MODELS_FORBIDDEN_QUERY_KEYS = frozenset( + { + "include_metadata", + "fallback_type", + "include_model_access_groups", + "only_model_access_groups", + } +) + + +class PublicReqMiddleware: + """Apply public-request safeguards keyed off ``X-Public-Req``. + + Install order matters: this middleware should be the OUTERMOST layer so + its header strip runs before any LiteLLM auth/logging middleware reads + the request. ``FastAPI.add_middleware`` inserts at index 0, so calling + it after the proxy's own ``add_middleware`` lines puts this on the + outside automatically. + """ + + def __init__(self, app: ASGIApp) -> None: + self.app = app + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + headers: List[Tuple[bytes, bytes]] = list(scope.get("headers") or []) + if not self._is_public(headers): + await self.app(scope, receive, send) + return + + scope = dict(scope) + + path = scope.get("path", "") + if path in MODELS_PATHS: + sanitized = self._strip_forbidden_query(scope.get("query_string", b"")) + if sanitized is not None: + scope["query_string"] = sanitized + + scope["headers"] = [ + (name, value) + for name, value in headers + if not name.lower().startswith(LITELLM_HEADER_PREFIX) + and name.lower() != PUBLIC_REQ_HEADER + ] + + async def send_wrapper(message: Message) -> None: + if message["type"] == "http.response.start": + message = { + **message, + "headers": [ + (name, value) + for name, value in message.get("headers", []) + if not name.lower().startswith(LITELLM_HEADER_PREFIX) + ], + } + await send(message) + + await self.app(scope, receive, send_wrapper) + + @staticmethod + def _is_public(headers: Iterable[Tuple[bytes, bytes]]) -> bool: + for name, value in headers: + if name.lower() == PUBLIC_REQ_HEADER: + return value.strip() == PUBLIC_REQ_VALUE + return False + + @staticmethod + def _strip_forbidden_query(query_string: bytes) -> Optional[bytes]: + """Remove forbidden keys from a percent-encoded query string. + + Returns the rewritten query string (possibly empty) when at least + one forbidden key was present, or ``None`` if the query string was + already safe. Returning ``None`` lets callers skip the + ``scope["query_string"]`` write and keep the original bytes + untouched — preserving exact ordering and any odd encoding the + client may have sent. + """ + if not query_string: + return None + pairs = parse_qsl( + query_string.decode("latin-1"), + keep_blank_values=True, + ) + kept = [(k, v) for k, v in pairs if k not in MODELS_FORBIDDEN_QUERY_KEYS] + if len(kept) == len(pairs): + return None + return urlencode(kept, doseq=True).encode("latin-1") diff --git a/tests/test_litellm/test_public_req_middleware.py b/tests/test_litellm/test_public_req_middleware.py new file mode 100644 index 000000000000..cd87580b1676 --- /dev/null +++ b/tests/test_litellm/test_public_req_middleware.py @@ -0,0 +1,358 @@ +"""Unit tests for litellm_extras.public_req_middleware.PublicReqMiddleware.""" + +from typing import Any, Dict, List, Optional, Tuple + +import pytest + +from litellm_extras.public_req_middleware import PublicReqMiddleware + + +def _build_scope( + method: str = "GET", + path: str = "/v1/chat/completions", + query_string: bytes = b"", + headers: Optional[List[Tuple[bytes, bytes]]] = None, +) -> Dict[str, Any]: + return { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": method, + "path": path, + "raw_path": path.encode(), + "query_string": query_string, + "headers": list(headers or []), + "client": ("127.0.0.1", 0), + "server": ("testserver", 80), + "scheme": "http", + } + + +class _RecordingApp: + """ASGI app that records what scope it received and yields predetermined chunks.""" + + def __init__( + self, + response_headers: Optional[List[Tuple[bytes, bytes]]] = None, + body_chunks: Optional[List[bytes]] = None, + status: int = 200, + ) -> None: + self.response_headers = response_headers or [] + self.body_chunks = body_chunks or [b'{"ok":true}'] + self.status = status + self.received_scope: Optional[Dict[str, Any]] = None + + async def __call__(self, scope: Dict[str, Any], receive: Any, send: Any) -> None: + self.received_scope = scope + await send( + { + "type": "http.response.start", + "status": self.status, + "headers": list(self.response_headers), + } + ) + for i, chunk in enumerate(self.body_chunks): + await send( + { + "type": "http.response.body", + "body": chunk, + "more_body": i < len(self.body_chunks) - 1, + } + ) + + +async def _drive( + middleware: PublicReqMiddleware, + scope: Dict[str, Any], +) -> Tuple[List[Dict[str, Any]], _RecordingApp]: + """Run the middleware against an empty-receive client and capture sends.""" + sent: List[Dict[str, Any]] = [] + + async def receive() -> Dict[str, Any]: + return {"type": "http.disconnect"} + + async def send(message: Dict[str, Any]) -> None: + sent.append(message) + + await middleware(scope, receive, send) + return sent, middleware.app # type: ignore[return-value] + + +def _start_msg(sent: List[Dict[str, Any]]) -> Dict[str, Any]: + starts = [m for m in sent if m["type"] == "http.response.start"] + assert starts, f"no response.start in {sent}" + return starts[0] + + +def _body_msgs(sent: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + return [m for m in sent if m["type"] == "http.response.body"] + + +@pytest.mark.asyncio +async def test_internal_request_passes_through_unchanged() -> None: + """No X-Public-Req → request scope untouched, response headers preserved.""" + app = _RecordingApp( + response_headers=[ + (b"content-type", b"application/json"), + (b"x-litellm-model-id", b"gpt-4o-mini-internal"), + (b"x-litellm-cache-hit", b"true"), + ] + ) + mw = PublicReqMiddleware(app) + scope = _build_scope( + headers=[ + (b"authorization", b"Bearer sk-internal"), + (b"x-litellm-mock-response", b"trace:xyz"), + (b"x-litellm-tags", b"team:alpha"), + ] + ) + + sent, _ = await _drive(mw, scope) + + received_hdrs = dict(app.received_scope["headers"]) + assert received_hdrs.get(b"x-litellm-mock-response") == b"trace:xyz" + assert received_hdrs.get(b"x-litellm-tags") == b"team:alpha" + + resp_hdrs = dict(_start_msg(sent)["headers"]) + assert resp_hdrs.get(b"x-litellm-model-id") == b"gpt-4o-mini-internal" + assert resp_hdrs.get(b"x-litellm-cache-hit") == b"true" + + +@pytest.mark.asyncio +async def test_public_request_strips_inbound_litellm_headers() -> None: + app = _RecordingApp() + mw = PublicReqMiddleware(app) + scope = _build_scope( + headers=[ + (b"authorization", b"Bearer sk-user"), + (b"x-public-req", b"1"), + (b"x-litellm-mock-response", b"hack"), + (b"x-litellm-num-retries", b"50"), + (b"x-litellm-tags", b"team:victim"), + (b"content-type", b"application/json"), + ] + ) + + await _drive(mw, scope) + + received = dict(app.received_scope["headers"]) + assert b"x-litellm-mock-response" not in received + assert b"x-litellm-num-retries" not in received + assert b"x-litellm-tags" not in received + # X-Public-Req itself is stripped too (don't surface the marker to LiteLLM). + assert b"x-public-req" not in received + # Non-LiteLLM headers must survive. + assert received[b"authorization"] == b"Bearer sk-user" + assert received[b"content-type"] == b"application/json" + + +@pytest.mark.asyncio +async def test_public_request_strips_outbound_litellm_headers() -> None: + app = _RecordingApp( + response_headers=[ + (b"content-type", b"application/json"), + (b"x-litellm-model-id", b"gpt-4o-mini-leak"), + (b"x-litellm-cache-hit", b"true"), + (b"x-litellm-response-cost", b"0.0001"), + (b"x-other", b"keep"), + ] + ) + mw = PublicReqMiddleware(app) + scope = _build_scope(headers=[(b"x-public-req", b"1")]) + + sent, _ = await _drive(mw, scope) + + resp = dict(_start_msg(sent)["headers"]) + assert b"x-litellm-model-id" not in resp + assert b"x-litellm-cache-hit" not in resp + assert b"x-litellm-response-cost" not in resp + assert resp[b"x-other"] == b"keep" + + +@pytest.mark.asyncio +async def test_public_request_streaming_body_chunks_pass_through_unbuffered() -> None: + """Body chunks must reach send() one-by-one without coalescing. + + Regression guard for the BaseHTTPMiddleware trap: a misimplemented + middleware would buffer the whole body before forwarding. We assert that + each chunk produced by the inner app surfaces as a separate + http.response.body message in original order. + """ + chunks = [ + b'data: {"a":1}\n\n', + b'data: {"b":2}\n\n', + b'data: {"c":3}\n\n', + b"data: [DONE]\n\n", + ] + app = _RecordingApp( + response_headers=[(b"content-type", b"text/event-stream")], + body_chunks=chunks, + ) + mw = PublicReqMiddleware(app) + scope = _build_scope(headers=[(b"x-public-req", b"1")]) + + sent, _ = await _drive(mw, scope) + + body_messages = _body_msgs(sent) + assert len(body_messages) == len(chunks) + assert [m["body"] for m in body_messages] == chunks + # All but the last chunk must signal more_body=True. + assert [m.get("more_body", False) for m in body_messages] == [ + True, + True, + True, + False, + ] + + +@pytest.mark.asyncio +async def test_public_request_strips_forbidden_models_query() -> None: + """Public mode silently drops forbidden keys and proxies the rest.""" + app = _RecordingApp() + mw = PublicReqMiddleware(app) + scope = _build_scope( + method="GET", + path="/v1/models", + query_string=b"include_metadata=true&fallback_type=general&team_id=keep", + headers=[(b"x-public-req", b"1")], + ) + + sent, _ = await _drive(mw, scope) + + assert _start_msg(sent)["status"] == 200, "request must reach inner app" + assert app.received_scope is not None + forwarded_qs = app.received_scope["query_string"] + assert b"include_metadata" not in forwarded_qs + assert b"fallback_type" not in forwarded_qs + assert b"team_id=keep" in forwarded_qs + + +@pytest.mark.asyncio +async def test_public_request_strip_leaves_query_when_only_forbidden_keys() -> None: + """If every query key was forbidden, the forwarded query is empty.""" + app = _RecordingApp() + mw = PublicReqMiddleware(app) + scope = _build_scope( + method="GET", + path="/v1/models", + query_string=b"include_metadata=true&only_model_access_groups=1", + headers=[(b"x-public-req", b"1")], + ) + + sent, _ = await _drive(mw, scope) + + assert _start_msg(sent)["status"] == 200 + assert app.received_scope is not None + assert app.received_scope["query_string"] == b"" + + +@pytest.mark.asyncio +async def test_public_request_models_clean_query_unchanged() -> None: + """Safe query strings must be passed through byte-for-byte.""" + app = _RecordingApp() + mw = PublicReqMiddleware(app) + scope = _build_scope( + method="GET", + path="/v1/models", + query_string=b"team_id=alpha&scope=expand", + headers=[(b"x-public-req", b"1")], + ) + + sent, _ = await _drive(mw, scope) + + assert _start_msg(sent)["status"] == 200 + assert app.received_scope["query_string"] == b"team_id=alpha&scope=expand" + + +@pytest.mark.asyncio +async def test_public_request_allows_models_without_forbidden_query() -> None: + app = _RecordingApp() + mw = PublicReqMiddleware(app) + scope = _build_scope( + method="GET", + path="/v1/models", + query_string=b"", + headers=[(b"x-public-req", b"1")], + ) + + sent, _ = await _drive(mw, scope) + + assert _start_msg(sent)["status"] == 200 + assert app.received_scope is not None + + +@pytest.mark.asyncio +async def test_internal_request_keeps_forbidden_query_on_models() -> None: + app = _RecordingApp() + mw = PublicReqMiddleware(app) + scope = _build_scope( + method="GET", + path="/v1/models", + query_string=b"include_metadata=true", + headers=[], + ) + + sent, _ = await _drive(mw, scope) + + assert _start_msg(sent)["status"] == 200 + + +@pytest.mark.asyncio +async def test_marker_value_other_than_1_treated_as_internal() -> None: + """Defensive: only the literal "1" enables public mode.""" + app = _RecordingApp() + mw = PublicReqMiddleware(app) + scope = _build_scope( + headers=[ + (b"x-public-req", b"true"), + (b"x-litellm-tags", b"team:keep"), + ] + ) + + await _drive(mw, scope) + + received = dict(app.received_scope["headers"]) + assert received.get(b"x-litellm-tags") == b"team:keep" + + +@pytest.mark.asyncio +async def test_marker_is_case_insensitive_for_header_name() -> None: + """ASGI lowercases header names by spec, but defensive parsing matters.""" + app = _RecordingApp() + mw = PublicReqMiddleware(app) + # ASGI gives us lowercase, but we double-check the comparison anyway. + scope = _build_scope( + headers=[ + (b"X-Public-Req", b"1"), # mixed case + (b"X-Litellm-Tags", b"team:victim"), + ] + ) + + await _drive(mw, scope) + + received = dict(app.received_scope["headers"]) + # Mixed-case x-litellm-tags should still be stripped. + assert b"X-Litellm-Tags" not in received + assert b"x-litellm-tags" not in received + + +@pytest.mark.asyncio +async def test_non_http_scope_is_passed_through() -> None: + """Websocket and lifespan must reach the inner app untouched.""" + app = _RecordingApp() + mw = PublicReqMiddleware(app) + scope = {"type": "websocket", "path": "/v1/realtime", "headers": []} + + async def receive() -> Dict[str, Any]: + return {"type": "websocket.connect"} + + async def send(message: Dict[str, Any]) -> None: + return None + + await mw(scope, receive, send) + + assert app.received_scope is scope + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])