diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 6fb0c2a3c75..10dc3cad13d 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -41,17 +41,6 @@ dependencies = [ "rustversion", ] -[[package]] -name = "async-trait" -version = "0.1.91" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae36dc4177970ef04fde5178d3e2429882def40e57a451f919c098f72baa6cec" -dependencies = [ - "proc-macro2", - "quote", - "syn 3.0.0", -] - [[package]] name = "atomic-waker" version = "1.1.2" @@ -412,14 +401,14 @@ dependencies = [ [[package]] name = "axum" -version = "0.7.9" +version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f" +checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" dependencies = [ - "async-trait", "axum-core", "base64", "bytes", + "form_urlencoded", "futures-util", "http 1.4.2", "http-body 1.1.0", @@ -432,15 +421,14 @@ dependencies = [ "mime", "percent-encoding", "pin-project-lite", - "rustversion", - "serde", + "serde_core", "serde_json", "serde_path_to_error", "serde_urlencoded", "sha1", "sync_wrapper", "tokio", - "tokio-tungstenite", + "tokio-tungstenite 0.29.0", "tower", "tower-layer", "tower-service", @@ -449,19 +437,17 @@ dependencies = [ [[package]] name = "axum-core" -version = "0.4.5" +version = "0.5.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09f2bd6146b97ae3359fa0cc6d6b376d9539582c7b4220f041a33ec24c226199" +checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" dependencies = [ - "async-trait", "bytes", - "futures-util", + "futures-core", "http 1.4.2", "http-body 1.1.0", "http-body-util", "mime", "pin-project-lite", - "rustversion", "sync_wrapper", "tower-layer", "tower-service", @@ -964,6 +950,18 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi 5.3.0", + "wasip2", +] + [[package]] name = "getrandom" version = "0.4.3" @@ -973,7 +971,7 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "r-efi", + "r-efi 6.0.0", "rand_core 0.10.1", "wasm-bindgen", ] @@ -1413,12 +1411,11 @@ dependencies = [ "futures-util", "litellm-config", "litellm-core", + "litellm-gateway-auth", "serde", "serde_json", - "sha2 0.10.9", - "subtle", "tokio", - "tokio-tungstenite", + "tokio-tungstenite 0.24.0", "tower", "tracing", ] @@ -1467,11 +1464,22 @@ dependencies = [ "serde_json", "thiserror 2.0.19", "tokio", - "tokio-tungstenite", + "tokio-tungstenite 0.24.0", "tracing", "tracing-subscriber", ] +[[package]] +name = "litellm-gateway-auth" +version = "0.1.0" +dependencies = [ + "axum", + "reqwest", + "sha2 0.10.9", + "subtle", + "tokio", +] + [[package]] name = "litellm-python-bridge" version = "0.1.0" @@ -1532,9 +1540,9 @@ checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" [[package]] name = "matchit" -version = "0.7.3" +version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94" +checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" [[package]] name = "memchr" @@ -1896,6 +1904,12 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + [[package]] name = "r-efi" version = "6.0.0" @@ -1909,10 +1923,20 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22f6172bdec972074665ed81ed53b71da00bfc44b65a753cfde883ec4c702a1a" dependencies = [ "libc", - "rand_chacha", + "rand_chacha 0.3.1", "rand_core 0.6.4", ] +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + [[package]] name = "rand" version = "0.10.2" @@ -1934,6 +1958,16 @@ dependencies = [ "rand_core 0.6.4", ] +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", +] + [[package]] name = "rand_core" version = "0.6.4" @@ -1943,6 +1977,15 @@ dependencies = [ "getrandom 0.2.17", ] +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + [[package]] name = "rand_core" version = "0.10.1" @@ -2678,7 +2721,19 @@ dependencies = [ "rustls-pki-types", "tokio", "tokio-rustls 0.26.4", - "tungstenite", + "tungstenite 0.24.0", +] + +[[package]] +name = "tokio-tungstenite" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f72a05e828585856dacd553fba484c242c46e391fb0e58917c942ee9202915c" +dependencies = [ + "futures-util", + "log", + "tokio", + "tungstenite 0.29.0", ] [[package]] @@ -2839,6 +2894,22 @@ dependencies = [ "utf-8", ] +[[package]] +name = "tungstenite" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c01152af293afb9c7c2a57e4b559c5620b421f6d133261c60dd2d0cdb38e6b8" +dependencies = [ + "bytes", + "data-encoding", + "http 1.4.2", + "httparse", + "log", + "rand 0.9.5", + "sha1", + "thiserror 2.0.19", +] + [[package]] name = "typenum" version = "1.20.1" @@ -2961,6 +3032,15 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + [[package]] name = "wasm-bindgen" version = "0.2.126" @@ -3186,6 +3266,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + [[package]] name = "writeable" version = "0.6.3" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 7f27946ac22..aad642b4993 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -6,6 +6,7 @@ members = [ "crates/python-interop", "crates/python-bridge", "crates/auth-aws", + "crates/gateway-auth", ] resolver = "2" @@ -23,7 +24,8 @@ litellm-config = { path = "crates/config" } litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false } litellm-python-interop = { path = "crates/python-interop" } litellm-auth-aws = { path = "crates/auth-aws" } -axum = "0.7" +litellm-gateway-auth = { path = "crates/gateway-auth" } +axum = "0.8" pyo3 = "0.29.2" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } pythonize = "0.29.0" diff --git a/litellm-rust/crates/ai-gateway/AGENTS.md b/litellm-rust/crates/ai-gateway/AGENTS.md index 9a970008ae0..dfdb6e7c750 100644 --- a/litellm-rust/crates/ai-gateway/AGENTS.md +++ b/litellm-rust/crates/ai-gateway/AGENTS.md @@ -1,16 +1,16 @@ # ai-gateway -The Axum server fronting the Rust gateway. Owns transport + config + auth only; deployment selection is `core::router`, and the LLM call (transforms, auth headers, provider HTTP) is a `core` route entrypoint. No provider handler lives here. +The Axum server fronting the Rust gateway. Owns transport and config; gateway authentication lives in `litellm-gateway-auth`, deployment selection is `core::router`, and the LLM call (transforms, provider auth headers, provider HTTP) is a `core` route entrypoint. No provider handler lives here. - Routes follow one template: each module exposes `pub fn router() -> Router`; `routes/mod.rs` only merges them. Simple routes are one file, non-trivial routes a folder (`handler`/`service`/`transport`). See `routes/AGENTS.md` -- Auth is an extractor: add `crate::auth::RequireMasterKey` to handler args; never re-implement the check per route +- Auth is an extractor from `litellm-gateway-auth`: add `RequireMasterKey` to handler args; never re-implement the check per route - Handlers are thin: validate and delegate to `service`; no business logic, no provider calls, no transforms - Services call `core`, they never reimplement it: pick the deployment, call the entrypoint. Provider resolution, auth headers, URL, and the HTTP call are `core`'s job - State is shared and cheap to clone: long-lived handles behind `Arc` in `state.rs`; read env/config only in `main.rs` ## Auth (interim) -- Single master key (`LITELLM_MASTER_KEY`) enforced by `auth::RequireMasterKey`; `Authorization: Bearer ` +- Single master key (`LITELLM_MASTER_KEY`) enforced by `litellm_gateway_auth::RequireMasterKey`; `Authorization: Bearer ` - Fails closed (500) when unset; constant-time compare; binds `127.0.0.1` by default (`HOST` to override) - Per-key auth, budgets, rate limits delegated to the Python proxy later; health routes unauthenticated diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index 2d317e65159..9400416f22a 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -17,21 +17,18 @@ required-features = ["server"] tracing.workspace = true litellm-core = { workspace = true, features = ["bedrock-auth"] } litellm-config.workspace = true +litellm-gateway-auth = { workspace = true, optional = true } # `sync` powers the bounded mpsc channel the realtime logger drains. tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time", "sync"] } futures-util.workspace = true serde_json.workspace = true axum = { workspace = true, features = ["ws"], optional = true } serde.workspace = true -subtle = { workspace = true, optional = true } -# sha2 hashes the master key into user_api_key_hash (matches the proxy's -# SHA-256 hash_token) so the plaintext credential never enters a log payload. -sha2 = { workspace = true, optional = true } tower = { version = "0.5.3", features = ["util"], optional = true } [features] default = [] -server = ["dep:axum", "dep:subtle", "dep:sha2"] +server = ["dep:axum", "dep:litellm-gateway-auth"] # Build the gateway's config from the proxy YAML via an embedded Python # interpreter (links libpython; requires `litellm` importable at runtime). python-config = ["litellm-config/python"] diff --git a/litellm-rust/crates/ai-gateway/src/auth/mod.rs b/litellm-rust/crates/ai-gateway/src/auth/mod.rs deleted file mode 100644 index b09d8285c3a..00000000000 --- a/litellm-rust/crates/ai-gateway/src/auth/mod.rs +++ /dev/null @@ -1,93 +0,0 @@ -//! 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. -//! -//! A handler opts in by adding [`RequireMasterKey`] to its arguments; auth then -//! runs during extraction, before the handler body. Routes never re-implement it. - -use axum::extract::FromRequestParts; -use axum::http::StatusCode; -use axum::http::header::AUTHORIZATION; -use axum::http::request::Parts; -use sha2::{Digest, Sha256}; -use subtle::ConstantTimeEq; - -use crate::state::AppState; - -/// SHA-256 hex digest of a token — the exact transform the Python proxy applies -/// (`litellm.proxy.utils.hash_token`). -/// -/// STRICT REQUIREMENT: a raw key (`LITELLM_MASTER_KEY`, a virtual key, …) must -/// **never** leave this gateway in a log payload. Spend logs and every callback -/// integration receive `user_api_key_hash`, so that field must be this hash, not -/// the credential. Hashing here also means the value matches the key's hash in -/// `LiteLLM_SpendLogs.api_key`, so realtime spend joins with the rest of LiteLLM. -pub fn hash_token(token: &str) -> String { - let digest = Sha256::digest(token.as_bytes()); - let mut hex = String::with_capacity(digest.len() * 2); - for byte in digest { - use std::fmt::Write; - let _ = write!(hex, "{byte:02x}"); - } - hex -} - -/// Extractor that requires the configured master key as a bearer token. -/// -/// Rejections: `500` when no master key is configured (permanent -/// misconfiguration, not a transient outage); `401` on a missing/incorrect -/// token. The comparison is constant-time. -pub struct RequireMasterKey; - -#[axum::async_trait] -impl FromRequestParts for RequireMasterKey { - type Rejection = (StatusCode, String); - - async fn from_request_parts( - parts: &mut Parts, - state: &AppState, - ) -> Result { - let Some(expected) = state.master_key.as_deref() else { - return Err(( - StatusCode::INTERNAL_SERVER_ERROR, - "gateway auth not configured (set LITELLM_MASTER_KEY)".to_string(), - )); - }; - let provided = parts - .headers - .get(AUTHORIZATION) - .and_then(|value| value.to_str().ok()) - .and_then(|value| value.strip_prefix("Bearer ")) - .map(str::trim); - match provided { - Some(token) if bool::from(token.as_bytes().ct_eq(expected.as_bytes())) => Ok(Self), - _ => Err(( - StatusCode::UNAUTHORIZED, - "missing or invalid bearer token".to_string(), - )), - } - } -} - -#[cfg(test)] -mod tests { - use super::hash_token; - - #[test] - fn hash_token_matches_python_sha256_hexdigest() { - // Must equal hashlib.sha256("sk-1234".encode()).hexdigest() — the value - // the proxy stores in LiteLLM_SpendLogs.api_key. - assert_eq!( - hash_token("sk-1234"), - "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b" - ); - // 64 lowercase hex chars, and never the raw input. - let h = hash_token("sk-secret"); - assert_eq!(h.len(), 64); - assert!(h.chars().all(|c| c.is_ascii_hexdigit())); - assert_ne!(h, "sk-secret"); - } -} diff --git a/litellm-rust/crates/ai-gateway/src/lib.rs b/litellm-rust/crates/ai-gateway/src/lib.rs index 429b97786c0..c2f53167d23 100644 --- a/litellm-rust/crates/ai-gateway/src/lib.rs +++ b/litellm-rust/crates/ai-gateway/src/lib.rs @@ -4,15 +4,13 @@ //! without pulling in the HTTP server: //! //! - [`io`]: compatibility exports and realtime WebSocket splice helpers. -//! - The server modules ([`auth`], [`routes`], [`state`]) and anything pulling +//! - The server modules ([`routes`], [`state`]) and anything pulling //! `axum` are gated behind the `server` feature, which the `litellm-ai-gateway` //! binary turns on. pub mod audio_transcription; pub mod io; -#[cfg(feature = "server")] -pub mod auth; #[cfg(feature = "server")] pub mod routes; #[cfg(feature = "server")] diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs index 2d326f126b1..18b4c9f7133 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -13,9 +13,9 @@ use litellm_core::Error; use litellm_core::lifecycle::StreamingCall; use serde_json::{Map, Value}; -use crate::auth::RequireMasterKey; use crate::constants::{MESSAGES_HEADERS_NOT_FORWARDED, MESSAGES_ROUTE_PATH}; use crate::state::AppState; +use litellm_gateway_auth::RequireMasterKey; /// This route's contribution to the app router. pub fn router() -> Router { 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 098232874bb..78d3848ce33 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs @@ -22,10 +22,10 @@ use litellm_core::realtime::types::RealtimeEvent; use litellm_core::router::Router as ModelRouter; use serde::Deserialize; -use crate::auth::RequireMasterKey; use crate::state::AppState; use litellm_core::integrations::custom_logger::CustomLogger; use litellm_core::integrations::types::RequestMetadata; +use litellm_gateway_auth::{RequireMasterKey, hash_token}; /// Process-local monotonic counter, mixed into the per-session call id so two /// sessions opened in the same nanosecond still get distinct ids. @@ -93,7 +93,7 @@ async fn bridge( let (ws_sink, ws_stream) = socket.split(); let metadata = RequestMetadata { - user_api_key_hash: master_key.as_deref().map(crate::auth::hash_token), + user_api_key_hash: master_key.as_deref().map(hash_token), ..RequestMetadata::default() }; @@ -105,7 +105,7 @@ async fn bridge( }); let client_out = ws_sink.with(|event: RealtimeEvent| async move { Ok::(Message::Text( - serde_json::to_string(&event).unwrap_or_default(), + serde_json::to_string(&event).unwrap_or_default().into(), )) }); diff --git a/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs index b11fb7b411d..bcd7299ce96 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs @@ -15,10 +15,10 @@ use litellm_core::responses::types::{ResponsesErrorFrame, ResponsesWsEvent, Resp use litellm_core::router::Router as ModelRouter; use serde::Deserialize; -use crate::auth::RequireMasterKey; use crate::state::AppState; use litellm_core::integrations::custom_logger::CustomLogger; use litellm_core::integrations::types::RequestMetadata; +use litellm_gateway_auth::{RequireMasterKey, hash_token}; static CALL_SEQ: AtomicU64 = AtomicU64::new(0); @@ -87,7 +87,7 @@ where S::Error: std::fmt::Display, { if let Ok(payload) = serde_json::to_string(&ResponsesErrorFrame::invalid_request(message)) { - let _ = sink.send(Message::Text(payload)).await; + let _ = sink.send(Message::Text(payload.into())).await; } let _ = sink .send(Message::Close(Some(axum::extract::ws::CloseFrame { @@ -117,7 +117,7 @@ impl Sink for ResponseClientSink { item: ResponsesWsEvent, ) -> Result<(), Self::Error> { let payload = serde_json::to_string(&item).map_err(axum::Error::new)?; - std::pin::Pin::new(&mut self.sink).start_send(Message::Text(payload)) + std::pin::Pin::new(&mut self.sink).start_send(Message::Text(payload.into())) } fn poll_flush( @@ -207,7 +207,7 @@ async fn bridge( let call_id = new_call_id(); let metadata = RequestMetadata { - user_api_key_hash: master_key.as_deref().map(crate::auth::hash_token), + user_api_key_hash: master_key.as_deref().map(hash_token), ..RequestMetadata::default() }; let client_in = Box::pin(stream.filter_map(|message| async move { diff --git a/litellm-rust/crates/ai-gateway/src/state.rs b/litellm-rust/crates/ai-gateway/src/state.rs index 34b4534fc74..259b3a8d7fb 100644 --- a/litellm-rust/crates/ai-gateway/src/state.rs +++ b/litellm-rust/crates/ai-gateway/src/state.rs @@ -4,6 +4,7 @@ use crate::io::realtime_pool::RealtimePool; use litellm_core::router::Router; use litellm_core::integrations::custom_logger::CustomLogger; +use litellm_gateway_auth::MasterKeyState; /// Shared application state handed to every route handler. #[derive(Clone)] @@ -19,3 +20,9 @@ pub struct AppState { /// every realtime connect fresh-dials exactly as before. pub realtime_pool: Arc, } + +impl MasterKeyState for AppState { + fn master_key(&self) -> Option<&str> { + self.master_key.as_deref() + } +} diff --git a/litellm-rust/crates/gateway-auth/Cargo.toml b/litellm-rust/crates/gateway-auth/Cargo.toml new file mode 100644 index 00000000000..a9f7254a8f3 --- /dev/null +++ b/litellm-rust/crates/gateway-auth/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "litellm-gateway-auth" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +axum.workspace = true +sha2.workspace = true +subtle.workspace = true + +[dev-dependencies] +reqwest.workspace = true +tokio = { workspace = true, features = ["macros", "rt"] } diff --git a/litellm-rust/crates/gateway-auth/src/extractor.rs b/litellm-rust/crates/gateway-auth/src/extractor.rs new file mode 100644 index 00000000000..b118fd6c2ec --- /dev/null +++ b/litellm-rust/crates/gateway-auth/src/extractor.rs @@ -0,0 +1,38 @@ +use axum::extract::FromRequestParts; +use axum::http::StatusCode; +use axum::http::header::AUTHORIZATION; +use axum::http::request::Parts; +use subtle::ConstantTimeEq; + +use crate::MasterKeyState; + +pub struct RequireMasterKey; + +impl FromRequestParts for RequireMasterKey +where + S: MasterKeyState + Send + Sync, +{ + type Rejection = (StatusCode, String); + + async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { + let Some(expected) = state.master_key() else { + return Err(( + StatusCode::INTERNAL_SERVER_ERROR, + "gateway auth not configured (set LITELLM_MASTER_KEY)".to_string(), + )); + }; + let provided = parts + .headers + .get(AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.strip_prefix("Bearer ")) + .map(str::trim); + match provided { + Some(token) if bool::from(token.as_bytes().ct_eq(expected.as_bytes())) => Ok(Self), + _ => Err(( + StatusCode::UNAUTHORIZED, + "missing or invalid bearer token".to_string(), + )), + } + } +} diff --git a/litellm-rust/crates/gateway-auth/src/hash.rs b/litellm-rust/crates/gateway-auth/src/hash.rs new file mode 100644 index 00000000000..84e770bb3fc --- /dev/null +++ b/litellm-rust/crates/gateway-auth/src/hash.rs @@ -0,0 +1,11 @@ +use sha2::{Digest, Sha256}; + +pub fn hash_token(token: &str) -> String { + let digest = Sha256::digest(token.as_bytes()); + let mut hex = String::with_capacity(digest.len() * 2); + for byte in digest { + use std::fmt::Write; + let _ = write!(hex, "{byte:02x}"); + } + hex +} diff --git a/litellm-rust/crates/gateway-auth/src/lib.rs b/litellm-rust/crates/gateway-auth/src/lib.rs new file mode 100644 index 00000000000..d2afb023b98 --- /dev/null +++ b/litellm-rust/crates/gateway-auth/src/lib.rs @@ -0,0 +1,7 @@ +mod extractor; +mod hash; +mod state; + +pub use extractor::RequireMasterKey; +pub use hash::hash_token; +pub use state::MasterKeyState; diff --git a/litellm-rust/crates/gateway-auth/src/state.rs b/litellm-rust/crates/gateway-auth/src/state.rs new file mode 100644 index 00000000000..0da3dc05417 --- /dev/null +++ b/litellm-rust/crates/gateway-auth/src/state.rs @@ -0,0 +1,3 @@ +pub trait MasterKeyState { + fn master_key(&self) -> Option<&str>; +} diff --git a/litellm-rust/crates/gateway-auth/tests/auth.rs b/litellm-rust/crates/gateway-auth/tests/auth.rs new file mode 100644 index 00000000000..19bf666065c --- /dev/null +++ b/litellm-rust/crates/gateway-auth/tests/auth.rs @@ -0,0 +1,86 @@ +use std::net::SocketAddr; + +use axum::Router; +use axum::http::StatusCode; +use axum::routing::get; +use litellm_gateway_auth::{MasterKeyState, RequireMasterKey, hash_token}; +use tokio::net::TcpListener; +use tokio::task::JoinHandle; + +#[derive(Clone)] +struct State(Option<&'static str>); + +impl MasterKeyState for State { + fn master_key(&self) -> Option<&str> { + self.0 + } +} + +async fn protected(_: RequireMasterKey) -> StatusCode { + StatusCode::OK +} + +async fn spawn_server(master_key: Option<&'static str>) -> (SocketAddr, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener"); + let address = listener.local_addr().expect("address"); + let app = Router::new() + .route("/protected", get(protected)) + .with_state(State(master_key)); + let server = tokio::spawn(async move { + axum::serve(listener, app).await.expect("server"); + }); + (address, server) +} + +#[test] +fn hash_token_matches_python_sha256_hexdigest() { + assert_eq!( + hash_token("sk-1234"), + "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b" + ); + let hash = hash_token("sk-secret"); + assert_eq!(hash.len(), 64); + assert!(hash.chars().all(|character| character.is_ascii_hexdigit())); + assert_ne!(hash, "sk-secret"); +} + +#[tokio::test] +async fn extractor_enforces_auth_over_http() { + let (configured_address, configured_server) = spawn_server(Some("master-key")).await; + let configured_url = format!("http://{configured_address}/protected"); + let client = reqwest::Client::new(); + + let accepted = client + .get(&configured_url) + .bearer_auth("master-key") + .send() + .await + .expect("accepted response"); + let invalid = client + .get(&configured_url) + .bearer_auth("wrong-key") + .send() + .await + .expect("invalid response"); + let missing = client + .get(&configured_url) + .send() + .await + .expect("missing response"); + + assert_eq!(accepted.status(), StatusCode::OK); + assert_eq!(invalid.status(), StatusCode::UNAUTHORIZED); + assert_eq!(missing.status(), StatusCode::UNAUTHORIZED); + configured_server.abort(); + + let (unconfigured_address, unconfigured_server) = spawn_server(None).await; + let unconfigured = client + .get(format!("http://{unconfigured_address}/protected")) + .bearer_auth("master-key") + .send() + .await + .expect("unconfigured response"); + + assert_eq!(unconfigured.status(), StatusCode::INTERNAL_SERVER_ERROR); + unconfigured_server.abort(); +}