From e2916e8e31576d176e68ee83c91d95b3af2c79b6 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 24 Jun 2026 15:35:44 -0700 Subject: [PATCH] fix(ai-gateway): close realtime auth review gaps --- litellm-rust/Cargo.lock | 1 + litellm-rust/Cargo.toml | 1 + litellm-rust/crates/ai-gateway/Cargo.toml | 3 +- litellm-rust/crates/ai-gateway/README.md | 17 +++--- .../crates/ai-gateway/src/auth/mod.rs | 7 +-- .../ai-gateway/src/auth/user_api_key.rs | 54 ++++++++++++++----- .../ai-gateway/src/routes/realtime/mod.rs | 48 ++++++++++++++--- 7 files changed, 101 insertions(+), 30 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 6fe84f1cfbc..371eb5911da 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -571,6 +571,7 @@ dependencies = [ "reqwest", "serde", "serde_json", + "serde_urlencoded", "sha2", "subtle", "tokio", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 25ee2213040..0aef56bf949 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -20,6 +20,7 @@ rand = "0.8" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" +serde_urlencoded = "0.7" sha2 = "0.10" subtle = "2" thiserror = "2.0" diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index 161a3c62d1d..5032259826b 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -22,13 +22,14 @@ futures-util.workspace = true serde_json.workspace = true axum = { workspace = true, features = ["ws"], optional = true } serde = { workspace = true, optional = true } +serde_urlencoded = { workspace = true, optional = true } sha2 = { workspace = true, optional = true } subtle = { workspace = true, optional = true } pyo3 = { workspace = true, features = ["auto-initialize"], optional = true } [features] default = [] -server = ["dep:axum", "dep:subtle", "dep:serde", "dep:sha2"] +server = ["dep:axum", "dep:subtle", "dep:serde", "dep:serde_urlencoded", "dep:sha2"] # Build the gateway's config from the proxy YAML via an embedded Python # interpreter (links libpython; requires `litellm` importable at runtime). python-config = ["dep:pyo3"] diff --git a/litellm-rust/crates/ai-gateway/README.md b/litellm-rust/crates/ai-gateway/README.md index e50fb433178..9e946c6e2ef 100644 --- a/litellm-rust/crates/ai-gateway/README.md +++ b/litellm-rust/crates/ai-gateway/README.md @@ -17,7 +17,7 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame. Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. - **Client endpoint:** `wss:///v1/realtime?model=` (WebSocket) -- **Auth:** `Authorization: Bearer ` — the master key (admin) or a virtual key validated by the LiteLLM proxy (see below). Fails closed. +- **Auth:** `Authorization: Bearer ` — proxy-admin only for realtime serving today. Non-admin virtual keys are verified but rejected until realtime usage is reported back to the LiteLLM proxy. Fails closed. - **Health:** `GET /health/readiness`, `GET /health/liveness`, `GET /health/gil` > **Realtime serving is pure Rust.** Python is used at **load time only** — to @@ -42,7 +42,9 @@ route/model permissions are enforced in exactly one place. A client presents load. A bounded ≤200-entry cache is available for future high-RPS *per-request* routes via `LITELLM_AUTH_CACHE_TTL_SECS > 0`, which trades budget/rate-limit freshness (bounded by the TTL) for fewer control-plane calls — like the proxy's - own ~60s auth cache. + own ~60s auth cache. Non-admin virtual-key identities are rejected on the + realtime route until realtime usage is reported back to the proxy, so callers + cannot consume upstream usage without key/team spend attribution. Data plane → control plane is itself authenticated with a **dedicated data-plane key** (NOT the master key — least privilege): the gateway sends @@ -73,9 +75,12 @@ export LITELLM_AUTH_VERIFY_URL=https:///v1/rust_control_plane/authen ``` Fails closed: a missing/wrong data-plane key, an unreachable proxy, or a rejected key -all yield `401`. Revocation and budget changes take effect within the cache TTL -(if caching is enabled via `LITELLM_AUTH_CACHE_TTL_SECS`; off by default → every connection re-verifies). Keep the proxy on a private network — -the verify endpoint is internal-only and excluded from the public OpenAPI spec. +all yield `401`; a valid non-admin virtual key currently yields `403` on +`/v1/realtime` until realtime spend reporting is available. Revocation and budget +changes take effect within the cache TTL (if caching is enabled via +`LITELLM_AUTH_CACHE_TTL_SECS`; off by default → every connection re-verifies). +Keep the proxy on a private network — the verify endpoint is internal-only and +excluded from the public OpenAPI spec. ## Configuration (config.yaml) @@ -117,7 +122,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/v1/rust_control_plane/authentication` | 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. Non-admin virtual keys are still rejected by `/v1/realtime` until realtime spend reporting is wired. | | `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/src/auth/mod.rs b/litellm-rust/crates/ai-gateway/src/auth/mod.rs index 18938317655..1920ff35b44 100644 --- a/litellm-rust/crates/ai-gateway/src/auth/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/auth/mod.rs @@ -1,9 +1,10 @@ //! Gateway authentication, as an axum **extractor** (the idiomatic pattern — //! keeps handlers clean and auth testable). //! -//! For now this is a single **master key**: any caller presenting it as -//! `Authorization: Bearer ` may invoke the gateway. Per-key auth, budgets, -//! and rate limits are delegated to the Python proxy in a later phase. +//! For now this supports both the local **master key** and virtual keys resolved +//! through the Python control plane. Individual routes can still narrow the +//! accepted identities. For example, realtime rejects non-admin virtual keys +//! until realtime usage is reported back to the proxy for spend attribution. //! //! A handler opts in by adding [`RequireMasterKey`] to its arguments; auth then //! runs during extraction, before the handler body. Routes never re-implement it. diff --git a/litellm-rust/crates/ai-gateway/src/auth/user_api_key.rs b/litellm-rust/crates/ai-gateway/src/auth/user_api_key.rs index 6379a6e18d5..534e72bd2aa 100644 --- a/litellm-rust/crates/ai-gateway/src/auth/user_api_key.rs +++ b/litellm-rust/crates/ai-gateway/src/auth/user_api_key.rs @@ -52,6 +52,10 @@ impl UserApiKeyAuth { ..Self::default() } } + + pub fn is_proxy_admin(&self) -> bool { + self.user_role.as_deref() == Some("proxy_admin") + } } /// SHA-256 over `(route, model, key)`, used as the cache key. Hashing keeps @@ -79,15 +83,25 @@ fn bearer_token(parts: &Parts) -> Option<&str> { .filter(|token| !token.is_empty()) } -/// Pull a query parameter's value out of a raw query string, e.g. `model` from -/// `model=gpt-realtime&foo=bar`. Returns the raw (un-percent-decoded) value; -/// model names use only query-safe characters (alphanumerics, `-`, `_`, `.`, `/`), -/// so no decoding is needed for the values we authorize on. -fn query_param<'a>(query: Option<&'a str>, name: &str) -> Option<&'a str> { - query?.split('&').find_map(|pair| { - let (k, v) = pair.split_once('=')?; - (k == name).then_some(v) - }) +#[derive(serde::Deserialize)] +struct RequestedModelQuery { + model: Option, +} + +/// Decode the requested model with the same query parser Axum's `Query` extractor +/// uses later in the realtime handler. This keeps auth and serving aligned: an +/// encoded key like `mo%64el=restricted` is still treated as `model=restricted` +/// during authorization. +fn requested_model(query: Option<&str>) -> Result, String> { + let Some(query) = query else { + return Ok(None); + }; + let parsed: RequestedModelQuery = + serde_urlencoded::from_str(query).map_err(|err| format!("invalid query string: {err}"))?; + Ok(parsed + .model + .map(|model| model.trim().to_string()) + .filter(|model| !model.is_empty())) } #[axum::async_trait] @@ -115,7 +129,9 @@ impl FromRequestParts for UserApiKeyAuth { // the key's route/model restrictions for what's actually being requested // (not just that the key exists — closes the model-authorization bypass). let route = parts.uri.path(); - let model = query_param(parts.uri.query(), "model"); + let model = requested_model(parts.uri.query()) + .map_err(|message| (StatusCode::BAD_REQUEST, message))?; + let model = model.as_deref(); // Virtual key: serve from cache, else verify via the (swappable) backend // and cache the result keyed by (route, model, key) hash. @@ -303,9 +319,19 @@ mod tests { #[test] fn admin_is_proxy_admin() { - assert_eq!( - UserApiKeyAuth::admin().user_role.as_deref(), - Some("proxy_admin") - ); + assert!(UserApiKeyAuth::admin().is_proxy_admin()); + assert!(!UserApiKeyAuth::default().is_proxy_admin()); + } + + #[test] + fn requested_model_decodes_query_parameter_names_like_axum() { + let model = requested_model(Some("mo%64el=gpt-realtime&foo=bar")).unwrap(); + assert_eq!(model.as_deref(), Some("gpt-realtime")); + } + + #[test] + fn requested_model_trims_empty_values() { + let model = requested_model(Some("model=%20%20%20")).unwrap(); + assert_eq!(model, None); } } diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs index 86d9e6d06d3..91f106b20c5 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs @@ -1,8 +1,7 @@ //! `GET /v1/realtime` (WebSocket). //! //! This file is the **axum surface**: `router()`, the handler, and the small -//! socket↔events adapter. The pure logic (no axum) lives in [`service`]. Auth is -//! the `RequireMasterKey` extractor, so the handler stays thin. +//! socket↔events adapter. The pure logic (no axum) lives in [`service`]. mod service; @@ -33,16 +32,31 @@ struct RealtimeQuery { model: String, } +fn require_realtime_billing_safe_auth(auth: &UserApiKeyAuth) -> Result<(), (StatusCode, String)> { + if auth.is_proxy_admin() { + return Ok(()); + } + Err(( + StatusCode::FORBIDDEN, + "non-admin virtual-key realtime is disabled until realtime usage is reported to the control plane".to_string(), + )) +} + /// Auth runs via the `UserApiKeyAuth` extractor (master key → admin, otherwise -/// cache then the swappable authenticator). We validate the model BEFORE the -/// upgrade so failures are clean HTTP (400/404), not a socket that opens then -/// closes, then hand the socket to `bridge`. +/// cache then the swappable authenticator). Non-admin virtual keys are rejected +/// until realtime usage is reported back to the proxy; otherwise callers could +/// consume upstream realtime spend without that spend being charged to their key +/// or team. We validate the model BEFORE the upgrade so failures are clean HTTP +/// (400/404), not a socket that opens then closes, then hand the socket to +/// `bridge`. async fn handle( - _auth: UserApiKeyAuth, + auth: UserApiKeyAuth, ws: WebSocketUpgrade, State(state): State, Query(query): Query, ) -> Result { + require_realtime_billing_safe_auth(&auth)?; + if query.model.trim().is_empty() { return Err(( StatusCode::BAD_REQUEST, @@ -87,3 +101,25 @@ async fn bridge( futures_util::pin_mut!(client_in, client_out); let _ = service::run(&router, &pool, &model, None, client_in, client_out).await; } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn realtime_allows_proxy_admin_identity() { + assert!(require_realtime_billing_safe_auth(&UserApiKeyAuth::admin()).is_ok()); + } + + #[test] + fn realtime_rejects_non_admin_virtual_keys_until_spend_is_reported() { + let err = require_realtime_billing_safe_auth(&UserApiKeyAuth { + user_role: Some("internal_user".to_string()), + ..UserApiKeyAuth::default() + }) + .unwrap_err(); + + assert_eq!(err.0, StatusCode::FORBIDDEN); + assert!(err.1.contains("usage is reported")); + } +}