From b7f37f3ff00e016046a7e4b604e241cebb8308c7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 24 Jun 2026 15:27:25 -0700 Subject: [PATCH] fix(proxy): move gateway auth under rust control plane --- litellm-rust/crates/ai-gateway/README.md | 10 +++--- litellm-rust/crates/ai-gateway/render.yaml | 4 +-- litellm-rust/crates/ai-gateway/src/main.rs | 2 +- litellm/proxy/proxy_server.py | 8 ++--- litellm/proxy/rust_control_plane/__init__.py | 1 + .../auth_endpoints.py} | 14 ++++---- .../proxy/rust_control_plane/__init__.py | 1 + .../test_auth_endpoints.py} | 36 +++++++++++++++++-- 8 files changed, 55 insertions(+), 21 deletions(-) create mode 100644 litellm/proxy/rust_control_plane/__init__.py rename litellm/proxy/{auth/internal_auth_endpoints.py => rust_control_plane/auth_endpoints.py} (94%) create mode 100644 tests/test_litellm/proxy/rust_control_plane/__init__.py rename tests/test_litellm/proxy/{auth/test_internal_auth_endpoints.py => rust_control_plane/test_auth_endpoints.py} (85%) diff --git a/litellm-rust/crates/ai-gateway/README.md b/litellm-rust/crates/ai-gateway/README.md index c1ed69d103b..e50fb433178 100644 --- a/litellm-rust/crates/ai-gateway/README.md +++ b/litellm-rust/crates/ai-gateway/README.md @@ -33,7 +33,7 @@ route/model permissions are enforced in exactly one place. A client presents - **Master key** (`LITELLM_MASTER_KEY`) → treated as proxy admin, checked locally (constant-time compare). No network call. - **Virtual key** (`sk-…`) → the gateway POSTs `{api_key, route, model}` to the proxy's - **`POST /internal/v1/auth/verify`**, which runs the proxy's real `user_api_key_auth` + **`POST /v1/rust_control_plane/authentication`**, which runs the proxy's real `user_api_key_auth` (key lookup, expiry, budget, rate-limit, and route **+ model** permissions via `can_key_call_model`) and returns the resolved identity. **By default every connection re-verifies** (no caching), so budget/block/rate-limit are enforced @@ -50,7 +50,7 @@ Data plane → control plane is itself authenticated with a **dedicated data-pla without it. Set the **same** secret on both sides. ```text -client ──Bearer sk-…──▶ ai-gateway ──POST /internal/v1/auth/verify──▶ LiteLLM proxy +client ──Bearer sk-…──▶ ai-gateway ──POST /v1/rust_control_plane/authentication──▶ LiteLLM proxy (X-LiteLLM-Data-Plane-Key) user_api_key_auth() re-verify per connection (cache opt-in) → UserAPIKeyAuth | 401 ``` @@ -61,14 +61,14 @@ On the **LiteLLM proxy** (control plane) — expose the verify endpoint: ```bash export LITELLM_DATA_PLANE_KEY= # a NEW secret, not the master key -litellm --config proxy_config.yaml # serves POST /internal/v1/auth/verify +litellm --config proxy_config.yaml # serves POST /v1/rust_control_plane/authentication ``` On the **gateway** (data plane) — point it at the proxy: ```bash export LITELLM_DATA_PLANE_KEY= -export LITELLM_AUTH_VERIFY_URL=https:///internal/v1/auth/verify +export LITELLM_AUTH_VERIFY_URL=https:///v1/rust_control_plane/authentication ./litellm-ai-gateway ``` @@ -117,7 +117,7 @@ overridden at deploy time (e.g. a Render secret file mounted at the same path). | `LITELLM_CONFIG_PATH` | yes (config mode) | — | Path to the config.yaml the gateway loads its `model_list` from. The Docker image defaults this to `/app/config.yaml`. | | `LITELLM_MASTER_KEY` | yes | — | Admin bearer token (checked locally, no proxy call). Unset ⇒ the master-key path is disabled. | | `LITELLM_DATA_PLANE_KEY` | for virtual keys | — | Dedicated secret the gateway sends as `X-LiteLLM-Data-Plane-Key` to authenticate itself to the proxy's verify endpoint. **Must match the proxy's `LITELLM_DATA_PLANE_KEY`.** Not the master key. | -| `LITELLM_AUTH_VERIFY_URL` | for virtual keys | `http://localhost:4000/internal/v1/auth/verify` | The proxy's verify endpoint the gateway delegates virtual-key auth to. | +| `LITELLM_AUTH_VERIFY_URL` | for virtual keys | `http://localhost:4000/v1/rust_control_plane/authentication` | The proxy's verify endpoint the gateway delegates virtual-key auth to. | | `LITELLM_AUTH_CACHE_TTL_SECS` | no | `0` | Verified-key cache TTL. **`0` = off (default)** → every connection re-verifies (budget/rate-limit enforced each time); cheap for realtime since auth is per-connection. Set `> 0` only for high-RPS per-request routes, trading budget/rate-limit freshness for fewer proxy calls. | | `OPENAI_API_KEY` | yes | — | Upstream OpenAI key. Referenced by config.yaml as `os.environ/OPENAI_API_KEY` for the gateway→OpenAI dial. | | `HOST` | no | `127.0.0.1` | **Set to `0.0.0.0` in any container/deploy** or external traffic is refused. | diff --git a/litellm-rust/crates/ai-gateway/render.yaml b/litellm-rust/crates/ai-gateway/render.yaml index 18175c2c36f..5c8f1dea630 100644 --- a/litellm-rust/crates/ai-gateway/render.yaml +++ b/litellm-rust/crates/ai-gateway/render.yaml @@ -35,8 +35,8 @@ services: # must match the proxy's LITELLM_DATA_PLANE_KEY. NOT the master key. - key: LITELLM_DATA_PLANE_KEY sync: false - # The proxy's internal verify endpoint, e.g. - # https:///internal/v1/auth/verify + # The proxy's Rust control-plane verify endpoint, e.g. + # https:///v1/rust_control_plane/authentication - key: LITELLM_AUTH_VERIFY_URL sync: false # Referenced by config.yaml as os.environ/OPENAI_API_KEY for the upstream dial. diff --git a/litellm-rust/crates/ai-gateway/src/main.rs b/litellm-rust/crates/ai-gateway/src/main.rs index 4e3dc852b77..c1c721510b7 100644 --- a/litellm-rust/crates/ai-gateway/src/main.rs +++ b/litellm-rust/crates/ai-gateway/src/main.rs @@ -27,7 +27,7 @@ const DEFAULT_HOST: &str = "127.0.0.1"; const DEFAULT_PORT: u16 = 4001; /// Default endpoint on the Python proxy that verifies virtual keys. -const DEFAULT_AUTH_VERIFY_URL: &str = "http://localhost:4000/internal/v1/auth/verify"; +const DEFAULT_AUTH_VERIFY_URL: &str = "http://localhost:4000/v1/rust_control_plane/authentication"; #[tokio::main] async fn main() { diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e48bb0c8d46..97bd598c406 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -290,9 +290,6 @@ from litellm.proxy.common_request_processing import ( create_response, ) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy -from litellm.proxy.auth.internal_auth_endpoints import ( - router as internal_auth_endpoints_router, -) from litellm.proxy.common_utils.debug_utils import init_verbose_loggers from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router from litellm.proxy.common_utils.encrypt_decrypt_utils import ( @@ -349,6 +346,9 @@ from litellm.proxy.hooks.prompt_injection_detection import ( from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger from litellm.proxy.image_endpoints.endpoints import router as image_router from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request +from litellm.proxy.rust_control_plane.auth_endpoints import ( + router as rust_control_plane_auth_router, +) from litellm.proxy.management_endpoints.budget_management_endpoints import ( router as budget_management_router, ) @@ -16641,7 +16641,7 @@ app.include_router(caching_router) app.include_router(analytics_router) app.include_router(callback_management_endpoints_router) app.include_router(debugging_endpoints_router) -app.include_router(internal_auth_endpoints_router) +app.include_router(rust_control_plane_auth_router) app.include_router(ui_crud_endpoints_router) app.include_router(openai_files_router) app.include_router(team_callback_router) diff --git a/litellm/proxy/rust_control_plane/__init__.py b/litellm/proxy/rust_control_plane/__init__.py new file mode 100644 index 00000000000..10f215aebc1 --- /dev/null +++ b/litellm/proxy/rust_control_plane/__init__.py @@ -0,0 +1 @@ +"""Internal routes exposed by the Python control plane for Rust data planes.""" diff --git a/litellm/proxy/auth/internal_auth_endpoints.py b/litellm/proxy/rust_control_plane/auth_endpoints.py similarity index 94% rename from litellm/proxy/auth/internal_auth_endpoints.py rename to litellm/proxy/rust_control_plane/auth_endpoints.py index 31c13345520..629e67060b1 100644 --- a/litellm/proxy/auth/internal_auth_endpoints.py +++ b/litellm/proxy/rust_control_plane/auth_endpoints.py @@ -1,5 +1,5 @@ """ -Internal data-plane auth endpoints. +Rust control-plane auth endpoints. This module exposes the authentication seam used by the Rust ai-gateway (the data plane). The Rust gateway terminates client connections and needs to @@ -68,7 +68,9 @@ class VerifyKeyRequest(BaseModel): model: Optional[str] = None -def _synthetic_request(route: str, api_key: str, model: Optional[str]) -> Request: +def _synthetic_request( + route: str, authorization_header: str, model: Optional[str] +) -> Request: """ Build a minimal ASGI request standing in for the client's real call, so ``user_api_key_auth`` evaluates the key against the intended data-plane @@ -86,7 +88,7 @@ def _synthetic_request(route: str, api_key: str, model: Optional[str]) -> Reques "path": route, "raw_path": route.encode(), "headers": [ - (b"authorization", f"Bearer {api_key}".encode()), + (b"authorization", authorization_header.encode()), (b"content-type", b"application/json"), ], "query_string": b"", @@ -97,11 +99,11 @@ def _synthetic_request(route: str, api_key: str, model: Optional[str]) -> Reques return Request(scope, receive) -router = APIRouter() +router = APIRouter(prefix="/v1/rust_control_plane", tags=["rust control plane"]) @router.post( - "/internal/v1/auth/verify", + "/authentication", dependencies=[Depends(require_data_plane_key)], # Internal data-plane route: keep it out of the public OpenAPI spec / docs # (and the generated UI schema.d.ts). It's not a client- or UI-facing API. @@ -132,7 +134,7 @@ async def verify_key(body: VerifyKeyRequest) -> dict[str, Any]: body.api_key if body.api_key.startswith("Bearer ") else f"Bearer {body.api_key}" ) synthetic_request = _synthetic_request( - route=body.route, api_key=bearer_key, model=body.model + route=body.route, authorization_header=bearer_key, model=body.model ) try: auth = await user_api_key_auth(request=synthetic_request, api_key=bearer_key) diff --git a/tests/test_litellm/proxy/rust_control_plane/__init__.py b/tests/test_litellm/proxy/rust_control_plane/__init__.py new file mode 100644 index 00000000000..c44841b4a73 --- /dev/null +++ b/tests/test_litellm/proxy/rust_control_plane/__init__.py @@ -0,0 +1 @@ +"""Tests for Rust control-plane endpoints.""" diff --git a/tests/test_litellm/proxy/auth/test_internal_auth_endpoints.py b/tests/test_litellm/proxy/rust_control_plane/test_auth_endpoints.py similarity index 85% rename from tests/test_litellm/proxy/auth/test_internal_auth_endpoints.py rename to tests/test_litellm/proxy/rust_control_plane/test_auth_endpoints.py index 74d2b41cd0a..597ea405306 100644 --- a/tests/test_litellm/proxy/auth/test_internal_auth_endpoints.py +++ b/tests/test_litellm/proxy/rust_control_plane/test_auth_endpoints.py @@ -1,14 +1,15 @@ -"""Unit tests for the internal data-plane auth endpoints.""" +"""Unit tests for the Rust control-plane auth endpoints.""" import pytest from fastapi import HTTPException, Request from litellm.proxy._types import ProxyException, UserAPIKeyAuth -from litellm.proxy.auth.internal_auth_endpoints import ( +from litellm.proxy.rust_control_plane.auth_endpoints import ( DATA_PLANE_KEY_ENV_VAR, DATA_PLANE_KEY_HEADER, VerifyKeyRequest, require_data_plane_key, + router, verify_key, ) @@ -21,7 +22,7 @@ def _make_request(headers: dict) -> Request: scope = { "type": "http", "method": "POST", - "path": "/internal/v1/auth/verify", + "path": "/v1/rust_control_plane/authentication", "headers": raw_headers, } return Request(scope) @@ -77,6 +78,13 @@ def test_require_data_plane_key_passes_when_correct(monkeypatch): assert require_data_plane_key(request) is None +def test_router_mounts_auth_verify_under_rust_control_plane(): + assert any( + getattr(route, "path", None) == "/v1/rust_control_plane/authentication" + for route in router.routes + ) + + @pytest.mark.asyncio async def test_verify_key_returns_model_dump(monkeypatch): expected_auth = UserAPIKeyAuth( @@ -104,6 +112,7 @@ async def test_verify_key_returns_model_dump(monkeypatch): assert captured["api_key"] == "Bearer sk-test-key" # Validation runs against a synthetic request carrying the gateway's route... assert captured["request"].url.path == "/v1/realtime" + assert captured["request"].headers["authorization"] == "Bearer sk-test-key" # ...and the requested model in the body, so model-access checks enforce it. assert (await captured["request"].json())["model"] == "gpt-realtime" assert result == expected_auth.model_dump(exclude_none=True, mode="json") @@ -129,6 +138,27 @@ async def test_verify_key_omits_model_when_absent(monkeypatch): assert (await captured["request"].json()) == {} +@pytest.mark.asyncio +async def test_verify_key_does_not_double_prefix_existing_bearer(monkeypatch): + captured = {} + + async def fake_user_api_key_auth(request, api_key): + captured["api_key"] = api_key + captured["request"] = request + return UserAPIKeyAuth(api_key="hashed-key") + + monkeypatch.setattr( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", + fake_user_api_key_auth, + ) + + body = VerifyKeyRequest(api_key="Bearer sk-test-key", route="/v1/realtime") + await verify_key(body=body) + + assert captured["api_key"] == "Bearer sk-test-key" + assert captured["request"].headers["authorization"] == "Bearer sk-test-key" + + @pytest.mark.asyncio async def test_verify_key_401_on_proxy_exception(monkeypatch): async def fake_user_api_key_auth(request, api_key):