mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
Merge d85f54e726 into 14193435b8
This commit is contained in:
commit
98e25ebd05
31 changed files with 452 additions and 1334 deletions
45
litellm-rust/Cargo.lock
generated
45
litellm-rust/Cargo.lock
generated
|
|
@ -774,6 +774,26 @@ dependencies = [
|
|||
"tower-service",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "axum-extra"
|
||||
version = "0.12.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "be44683b41ccb9ab2d23a5230015c9c3c55be97a25e4428366de8873103f7970"
|
||||
dependencies = [
|
||||
"axum-core",
|
||||
"bytes",
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"headers",
|
||||
"http 1.4.2",
|
||||
"http-body 1.1.0",
|
||||
"http-body-util",
|
||||
"mime",
|
||||
"pin-project-lite",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "axum-login"
|
||||
version = "0.18.0"
|
||||
|
|
@ -2662,6 +2682,30 @@ dependencies = [
|
|||
"hashbrown 0.16.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "headers"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fc2700e3fe555c3310aa7286cac6167449f2c87e0eb58769c9208a1c58a1d106"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"headers-core",
|
||||
"http 1.4.2",
|
||||
"httpdate",
|
||||
"mime",
|
||||
"sha1 0.10.7",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "headers-core"
|
||||
version = "0.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "54b4a22553d4242c49fddb9ba998a99962b5cc6f22cb5a3482bec22522403ce4"
|
||||
dependencies = [
|
||||
"http 1.4.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "heck"
|
||||
version = "0.5.0"
|
||||
|
|
@ -3797,6 +3841,7 @@ name = "litellm-gateway-auth"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"axum",
|
||||
"axum-extra",
|
||||
"axum-login",
|
||||
"futures-util",
|
||||
"litellm-auth-types",
|
||||
|
|
|
|||
|
|
@ -76,6 +76,7 @@ litellm-redis-testing = { path = "crates/redis-testing" }
|
|||
askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] }
|
||||
tracing = "0.1"
|
||||
axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio", "multipart"] }
|
||||
axum-extra = { version = "0.12.6", default-features = false, features = ["typed-header"] }
|
||||
axum-login = "0.18.0"
|
||||
tower-sessions = { version = "0.14.0", features = ["memory-store"] }
|
||||
bytes = "1"
|
||||
|
|
|
|||
|
|
@ -1,27 +1,13 @@
|
|||
# Gateway authentication
|
||||
|
||||
Inbound authentication uses a verifier, identity resolver, and authorizer injected into `Auth`. `Auth::from_config` composes the current master-key implementation. `Auth::new` accepts alternate implementations and an injected clock
|
||||
`Auth` turns a presented secret into a `Caller`. Today the only caller is `Caller::Master`, verified by a constant-time comparison against the configured master key, which may be an `os.environ/` reference resolved through the injected `SecretSource`. New credential kinds (virtual keys, JWT) become new `Caller` variants and new branches in `Auth::authenticate`, not pluggable verifier traits. Ports exist only at real infrastructure boundaries such as secrets and, later, key storage
|
||||
|
||||
`CredentialExtractor` selects a token or transport credential from HTTP request parts. The default `Bearer` extractor is strict. A host can supply another extractor with `Auth::with_extractor` for a separate route profile. A transport credential is only a selection marker, never proof of identity
|
||||
Credential extractors are plain Axum `FromRequestParts` types that convert into `SecretValue`. `Bearer` takes exactly one RFC 6750 Authorization bearer token through typed `headers` decoders, never hand-scanned strings. A route profile picks its extractor with `from_fn_with_state(state, authenticate::<E>)`, and extractor configuration comes from that state through `FromRef`
|
||||
|
||||
`Authenticator` receives that selected credential and HTTP request parts. Transport verifiers must obtain proof from trusted server extensions, not client-supplied identity headers. It returns `VerifiedIdentity` only after verification. Request parts provide headers and trusted transport extensions for integrations. The verifier must validate any client-controlled claims before treating them as identity. There is no fallback chain after verification failure
|
||||
The `authenticate` middleware inserts the `Caller` as a request extension and replaces any preexisting one. Handlers extract `Caller`, which fails closed when the middleware did not run, and call `Caller::authorize` with the typed `AccessRequest` before calling a provider or MCP server. `Caller::owner` scopes MCP sessions and response caches
|
||||
|
||||
`IdentityResolver` maps verified identity into an internal principal and permissions. The caller retains the original authentication evidence and credential restrictions independently of that mapping. Principals include an authority, subject, and kind. External subjects must remain authority-scoped unless the resolver explicitly maps them to a shared internal identity
|
||||
Local UI login uses axum-login and tower-sessions. `UiSession` validates the session and its CSRF bearer and never produces a `Caller`, so a UI session cannot authorize inference
|
||||
|
||||
`AuthenticatedRequest::authorize` checks expiry and intersects resolved permissions with credential restrictions before calling the injected authorizer. The authorizer may impose additional policy but cannot expand either permission set. The returned operation grant binds the exact authenticated caller instance and access request. Consuming it for another caller or target fails
|
||||
`general_settings.custom_auth` is served by the Python proxy, and the gateway refuses to start with it. `CustomAuth` is the reserved seam for running it natively. Its outcome maps Python's modes, where `Caller` and `Reject` are the hook's decision, `Fallthrough` is `auto` mode moving on to built-in auth, and `Rekey` is the enterprise hook returning a replacement key for built-in auth. It runs first in `Auth::authenticate` and receives request parts plus the optional extracted key, so a profile with a hook uses `Option<E>` as its extractor. The runtime behind it (Python extension host, Lua, or other) is undecided
|
||||
|
||||
`Permissions::All`, `None`, and `Only` provide the initial policy vocabulary. `Only` matches exact typed access requests. Model requests contain both the public name and resolved deployment model. MCP requests contain the resolved server ID and upstream operation, including bare tool/prompt names and resource URIs. More expressive policy adapters can be added without changing credential verification
|
||||
|
||||
The Axum `authenticate` middleware currently accepts one Authorization header using the existing Bearer format. Missing, duplicate, empty, or malformed credentials fail. Authentication replaces any preexisting caller extension and checks the method and matched route before dispatch. Handlers extract `AuthenticatedRequest` and authorize their parsed operation before calling a provider. Missing authenticated context fails closed
|
||||
|
||||
The gateway also inserts an MCP `Authorization` extension and a `SessionOwner` derived from the verified principal and credential reference. The MCP host checks the injected policy at initialization and before sending resolved upstream operations. Standalone MCP library hosts retain their existing explicitly trusted registry behavior when no policy is installed. Production mounts must install authentication and the policy extension together
|
||||
|
||||
Local UI login continues using axum-login and tower-sessions. Only after session and CSRF validation does `UiSession` expose an authenticated caller. Its credentials are restricted to session-info and logout operations, so the UI CSRF bearer cannot authorize inference. Future UI management routes must extend that explicit scope
|
||||
|
||||
Authentication evidence and principals contain no raw token, password, request body, or mutable accounting state. Session ownership is scoped by principal authority, subject, verifier, and credential ID, so separate credentials do not silently share an MCP session. Credential rotation may retain ownership when the verifier preserves a stable credential ID. Scope and expiry checks still run for each operation
|
||||
|
||||
Failures distinguish invalid or expired credentials, forbidden operations, unavailable authentication services, and missing server configuration/context. HTTP adapters map those outcomes to status codes. Inference keeps its API-specific error envelopes and MCP translates policy failures to its own protocol
|
||||
|
||||
Virtual-key storage, JWT verification, OAuth2 introspection, trusted-proxy validation, SSO, and custom Python hook adapters are not implemented here yet. They should supply these contracts rather than bypassing the shared authorization boundary. Request-body-dependent custom hooks will need a bounded endpoint adapter after parsing
|
||||
|
||||
Budget reservations, rate limits, usage settlement, retries, and cancellation accounting belong after authorization in the execution lifecycle. This foundation adds no accounting backend or cached authorization decisions. Future reservations must bind the same caller and resolved operation and settle idempotently across streaming completion and cancellation
|
||||
Errors map to HTTP status codes in `Error::status`. Inference keeps its API-specific error envelopes and MCP translates authorization failures to its own protocol
|
||||
|
|
|
|||
|
|
@ -6,7 +6,8 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
axum = { workspace = true, features = ["matched-path"] }
|
||||
axum.workspace = true
|
||||
axum-extra.workspace = true
|
||||
axum-login.workspace = true
|
||||
litellm-auth-types.workspace = true
|
||||
litellm-config.workspace = true
|
||||
|
|
|
|||
|
|
@ -1,189 +1,42 @@
|
|||
use std::{future::Future, pin::Pin, sync::Arc, time::SystemTime};
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::http::{Request, request::Parts};
|
||||
use litellm_auth_types::SecretValue;
|
||||
use litellm_config::Config;
|
||||
use litellm_secrets::source::SecretSource;
|
||||
use sha2::{Digest, Sha256};
|
||||
use subtle::ConstantTimeEq;
|
||||
|
||||
use crate::{
|
||||
AuthenticatedCaller, AuthenticatedRequest, Authentication, AuthenticationMethod, Authorizer,
|
||||
Bearer, CredentialExtractor, Error, NoAdditionalPolicy, Permissions, Principal, PrincipalKind,
|
||||
ResolvedIdentity, VerifiedIdentity,
|
||||
};
|
||||
|
||||
pub type AuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
|
||||
|
||||
pub trait Clock: Send + Sync {
|
||||
fn now(&self) -> SystemTime;
|
||||
}
|
||||
|
||||
pub struct SystemClock;
|
||||
|
||||
impl Clock for SystemClock {
|
||||
fn now(&self) -> SystemTime {
|
||||
SystemTime::now()
|
||||
}
|
||||
}
|
||||
|
||||
pub enum Credential {
|
||||
Token(SecretValue),
|
||||
Transport,
|
||||
}
|
||||
|
||||
pub trait Authenticator: Send + Sync {
|
||||
fn validate(&self) -> AuthFuture<'_, ()> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
|
||||
fn verify<'a>(
|
||||
&'a self,
|
||||
credential: &'a Credential,
|
||||
request: &'a Parts,
|
||||
) -> AuthFuture<'a, VerifiedIdentity>;
|
||||
}
|
||||
|
||||
pub trait IdentityResolver: Send + Sync {
|
||||
fn resolve<'a>(&'a self, identity: &'a VerifiedIdentity) -> AuthFuture<'a, ResolvedIdentity>;
|
||||
}
|
||||
|
||||
pub struct LocalAdministrator;
|
||||
|
||||
impl IdentityResolver for LocalAdministrator {
|
||||
fn resolve<'a>(&'a self, identity: &'a VerifiedIdentity) -> AuthFuture<'a, ResolvedIdentity> {
|
||||
Box::pin(async move {
|
||||
let allowed = match identity.authentication.method {
|
||||
AuthenticationMethod::MasterKey => identity.principal == master_principal(),
|
||||
AuthenticationMethod::Session => {
|
||||
identity.principal.authority() == "litellm:local-ui"
|
||||
&& identity.principal.kind() == PrincipalKind::Human
|
||||
}
|
||||
AuthenticationMethod::External(_) => false,
|
||||
};
|
||||
if !allowed {
|
||||
return Err(Error::Forbidden);
|
||||
}
|
||||
Ok(ResolvedIdentity {
|
||||
principal: identity.principal.clone(),
|
||||
permissions: Permissions::All,
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
use crate::{Caller, Error};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Auth {
|
||||
extractor: Arc<dyn CredentialExtractor>,
|
||||
authenticator: Arc<dyn Authenticator>,
|
||||
identities: Arc<dyn IdentityResolver>,
|
||||
authorizer: Arc<dyn Authorizer>,
|
||||
clock: Arc<dyn Clock>,
|
||||
}
|
||||
|
||||
impl Auth {
|
||||
pub fn new(
|
||||
authenticator: Arc<dyn Authenticator>,
|
||||
identities: Arc<dyn IdentityResolver>,
|
||||
authorizer: Arc<dyn Authorizer>,
|
||||
clock: Arc<dyn Clock>,
|
||||
) -> Self {
|
||||
Self {
|
||||
extractor: Arc::new(Bearer),
|
||||
authenticator,
|
||||
identities,
|
||||
authorizer,
|
||||
clock,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_config(config: &Config, secrets: Arc<dyn SecretSource>) -> Self {
|
||||
let master = Arc::new(MasterKeyAuthenticator {
|
||||
key: config.general_settings.master_key.clone(),
|
||||
secrets,
|
||||
});
|
||||
Self {
|
||||
extractor: Arc::new(Bearer),
|
||||
authenticator: master,
|
||||
identities: Arc::new(LocalAdministrator),
|
||||
authorizer: Arc::new(NoAdditionalPolicy),
|
||||
clock: Arc::new(SystemClock),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_extractor(self, extractor: Arc<dyn CredentialExtractor>) -> Self {
|
||||
Self { extractor, ..self }
|
||||
}
|
||||
|
||||
pub async fn authenticate_parts(&self, request: &Parts) -> Result<AuthenticatedRequest, Error> {
|
||||
let credential = self.extractor.extract(request)?;
|
||||
self.authenticate_request(&credential, request).await
|
||||
}
|
||||
|
||||
pub async fn validate(&self) -> Result<(), Error> {
|
||||
self.authenticator.validate().await
|
||||
}
|
||||
|
||||
pub async fn authenticate(
|
||||
&self,
|
||||
credential: &SecretValue,
|
||||
) -> Result<AuthenticatedRequest, Error> {
|
||||
self.authenticate_request(
|
||||
&Credential::Token(credential.clone()),
|
||||
&Request::new(()).into_parts().0,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn authenticate_request(
|
||||
&self,
|
||||
credential: &Credential,
|
||||
request: &Parts,
|
||||
) -> Result<AuthenticatedRequest, Error> {
|
||||
let verified = self.authenticator.verify(credential, request).await?;
|
||||
resolve(
|
||||
verified,
|
||||
self.identities.as_ref(),
|
||||
self.authorizer.clone(),
|
||||
self.clock.clone(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve(
|
||||
verified: VerifiedIdentity,
|
||||
identities: &dyn IdentityResolver,
|
||||
authorizer: Arc<dyn Authorizer>,
|
||||
clock: Arc<dyn Clock>,
|
||||
) -> Result<AuthenticatedRequest, Error> {
|
||||
if verified
|
||||
.authentication
|
||||
.expires_at
|
||||
.is_some_and(|expiry| expiry <= clock.now())
|
||||
{
|
||||
return Err(Error::Expired);
|
||||
}
|
||||
let resolved = identities.resolve(&verified).await?;
|
||||
Ok(AuthenticatedRequest::new(
|
||||
Arc::new(AuthenticatedCaller::new(verified, resolved)),
|
||||
authorizer,
|
||||
clock,
|
||||
))
|
||||
}
|
||||
|
||||
pub struct MasterKeyAuthenticator {
|
||||
key: Option<SecretValue>,
|
||||
master_key: Option<SecretValue>,
|
||||
secrets: Arc<dyn SecretSource>,
|
||||
}
|
||||
|
||||
impl MasterKeyAuthenticator {
|
||||
pub fn new(key: Option<SecretValue>, secrets: Arc<dyn SecretSource>) -> Self {
|
||||
Self { key, secrets }
|
||||
impl Auth {
|
||||
pub fn from_config(config: &Config, secrets: Arc<dyn SecretSource>) -> Self {
|
||||
Self {
|
||||
master_key: config.general_settings.master_key.clone(),
|
||||
secrets,
|
||||
}
|
||||
}
|
||||
|
||||
async fn key(&self) -> Result<SecretValue, Error> {
|
||||
let configured = self.key.as_ref().ok_or(Error::Unconfigured)?;
|
||||
pub async fn validate(&self) -> Result<(), Error> {
|
||||
self.master_key().await.map(drop)
|
||||
}
|
||||
|
||||
pub async fn authenticate(&self, token: &SecretValue) -> Result<Caller, Error> {
|
||||
let expected = self.master_key().await?;
|
||||
if bool::from(Sha256::digest(token.expose()).ct_eq(&Sha256::digest(expected.expose()))) {
|
||||
Ok(Caller::Master)
|
||||
} else {
|
||||
Err(Error::InvalidToken)
|
||||
}
|
||||
}
|
||||
|
||||
async fn master_key(&self) -> Result<SecretValue, Error> {
|
||||
let configured = self.master_key.as_ref().ok_or(Error::Unconfigured)?;
|
||||
let resolved = match configured.expose().strip_prefix("os.environ/") {
|
||||
Some(name) if !name.is_empty() => self
|
||||
.secrets
|
||||
|
|
@ -199,45 +52,3 @@ impl MasterKeyAuthenticator {
|
|||
Ok(resolved)
|
||||
}
|
||||
}
|
||||
|
||||
impl Authenticator for MasterKeyAuthenticator {
|
||||
fn validate(&self) -> AuthFuture<'_, ()> {
|
||||
Box::pin(async { self.key().await.map(|_| ()) })
|
||||
}
|
||||
|
||||
fn verify<'a>(
|
||||
&'a self,
|
||||
credential: &'a Credential,
|
||||
_: &'a Parts,
|
||||
) -> AuthFuture<'a, VerifiedIdentity> {
|
||||
Box::pin(async move {
|
||||
let expected = self.key().await?;
|
||||
let Credential::Token(credential) = credential else {
|
||||
return Err(Error::InvalidToken);
|
||||
};
|
||||
if !bool::from(
|
||||
Sha256::digest(credential.expose()).ct_eq(&Sha256::digest(expected.expose())),
|
||||
) {
|
||||
return Err(Error::InvalidToken);
|
||||
}
|
||||
Ok(VerifiedIdentity {
|
||||
principal: master_principal(),
|
||||
authentication: Authentication {
|
||||
method: AuthenticationMethod::MasterKey,
|
||||
verifier: "litellm:master-key".into(),
|
||||
credential_id: "master".into(),
|
||||
expires_at: None,
|
||||
},
|
||||
restrictions: Permissions::All,
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn master_principal() -> Principal {
|
||||
Principal::new(
|
||||
"litellm:master-key".into(),
|
||||
"master".into(),
|
||||
PrincipalKind::System,
|
||||
)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,139 +0,0 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use crate::{AuthFuture, AuthenticatedCaller, Clock, Error, SharedCaller};
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub enum McpAction {
|
||||
Connect,
|
||||
ListTools,
|
||||
CallTool(String),
|
||||
ListPrompts,
|
||||
GetPrompt(String),
|
||||
ListResources,
|
||||
ListResourceTemplates,
|
||||
ReadResource(String),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub enum UiAction {
|
||||
SessionInfo,
|
||||
Logout,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub enum AccessRequest {
|
||||
Route { method: String, path: String },
|
||||
Model { name: String, deployment: String },
|
||||
Mcp { server: String, action: McpAction },
|
||||
Ui(UiAction),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub enum Permissions {
|
||||
All,
|
||||
#[default]
|
||||
None,
|
||||
Only(Arc<[AccessRequest]>),
|
||||
}
|
||||
|
||||
impl Permissions {
|
||||
pub fn allows(&self, request: &AccessRequest) -> bool {
|
||||
match self {
|
||||
Self::All => true,
|
||||
Self::None => false,
|
||||
Self::Only(allowed) => allowed.contains(request),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub trait Authorizer: Send + Sync {
|
||||
fn authorize<'a>(
|
||||
&'a self,
|
||||
caller: &'a AuthenticatedCaller,
|
||||
request: &'a AccessRequest,
|
||||
) -> AuthFuture<'a, ()>;
|
||||
}
|
||||
|
||||
pub struct NoAdditionalPolicy;
|
||||
|
||||
impl Authorizer for NoAdditionalPolicy {
|
||||
fn authorize<'a>(
|
||||
&'a self,
|
||||
_: &'a AuthenticatedCaller,
|
||||
_: &'a AccessRequest,
|
||||
) -> AuthFuture<'a, ()> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AuthenticatedRequest {
|
||||
caller: SharedCaller,
|
||||
authorizer: Arc<dyn Authorizer>,
|
||||
clock: Arc<dyn Clock>,
|
||||
}
|
||||
|
||||
impl AuthenticatedRequest {
|
||||
pub(crate) fn new(
|
||||
caller: SharedCaller,
|
||||
authorizer: Arc<dyn Authorizer>,
|
||||
clock: Arc<dyn Clock>,
|
||||
) -> Self {
|
||||
Self {
|
||||
caller,
|
||||
authorizer,
|
||||
clock,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn caller(&self) -> &AuthenticatedCaller {
|
||||
&self.caller
|
||||
}
|
||||
|
||||
pub async fn authorize(&self, request: AccessRequest) -> Result<AuthorizedOperation, Error> {
|
||||
if self
|
||||
.caller
|
||||
.authentication()
|
||||
.expires_at
|
||||
.is_some_and(|expiry| expiry <= self.clock.now())
|
||||
{
|
||||
return Err(Error::Expired);
|
||||
}
|
||||
if !self.caller.restrictions().allows(&request)
|
||||
|| !self.caller.permissions().allows(&request)
|
||||
{
|
||||
return Err(Error::Forbidden);
|
||||
}
|
||||
self.authorizer.authorize(&self.caller, &request).await?;
|
||||
Ok(AuthorizedOperation {
|
||||
caller: self.caller.clone(),
|
||||
request,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub struct AuthorizedOperation {
|
||||
caller: SharedCaller,
|
||||
request: AccessRequest,
|
||||
}
|
||||
|
||||
impl AuthorizedOperation {
|
||||
pub fn caller(&self) -> &AuthenticatedCaller {
|
||||
&self.caller
|
||||
}
|
||||
pub fn request(&self) -> &AccessRequest {
|
||||
&self.request
|
||||
}
|
||||
|
||||
pub fn consume(
|
||||
self,
|
||||
caller: &AuthenticatedCaller,
|
||||
request: &AccessRequest,
|
||||
) -> Result<(), Error> {
|
||||
if std::ptr::eq(self.caller.as_ref(), caller) && &self.request == request {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(Error::Forbidden)
|
||||
}
|
||||
}
|
||||
}
|
||||
52
litellm-rust/crates/gateway-auth/src/caller.rs
Normal file
52
litellm-rust/crates/gateway-auth/src/caller.rs
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
use axum::{extract::FromRequestParts, http::request::Parts};
|
||||
|
||||
use crate::Error;
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub enum McpAction {
|
||||
Connect,
|
||||
ListTools,
|
||||
CallTool(String),
|
||||
ListPrompts,
|
||||
GetPrompt(String),
|
||||
ListResources,
|
||||
ListResourceTemplates,
|
||||
ReadResource(String),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub enum AccessRequest {
|
||||
Model { name: String, deployment: String },
|
||||
Mcp { server: String, action: McpAction },
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub enum Caller {
|
||||
Master,
|
||||
}
|
||||
|
||||
impl Caller {
|
||||
pub fn authorize(&self, _: &AccessRequest) -> Result<(), Error> {
|
||||
match self {
|
||||
Self::Master => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn owner(&self) -> String {
|
||||
match self {
|
||||
Self::Master => "master".into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: Send + Sync> FromRequestParts<S> for Caller {
|
||||
type Rejection = Error;
|
||||
|
||||
async fn from_request_parts(parts: &mut Parts, _: &S) -> Result<Self, Error> {
|
||||
parts
|
||||
.extensions
|
||||
.get::<Self>()
|
||||
.cloned()
|
||||
.ok_or(Error::MissingIdentity)
|
||||
}
|
||||
}
|
||||
21
litellm-rust/crates/gateway-auth/src/custom.rs
Normal file
21
litellm-rust/crates/gateway-auth/src/custom.rs
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
use std::{future::Future, pin::Pin};
|
||||
|
||||
use axum::http::request::Parts;
|
||||
use litellm_auth_types::SecretValue;
|
||||
|
||||
use crate::{Caller, Error};
|
||||
|
||||
pub enum CustomOutcome {
|
||||
Caller(Caller),
|
||||
Rekey(SecretValue),
|
||||
Fallthrough,
|
||||
Reject(Error),
|
||||
}
|
||||
|
||||
pub trait CustomAuth: Send + Sync {
|
||||
fn authenticate<'a>(
|
||||
&'a self,
|
||||
request: &'a Parts,
|
||||
key: Option<&'a SecretValue>,
|
||||
) -> Pin<Box<dyn Future<Output = CustomOutcome> + Send + 'a>>;
|
||||
}
|
||||
|
|
@ -17,12 +17,8 @@ pub enum KeyError {
|
|||
pub enum Error {
|
||||
#[error("operation is not permitted")]
|
||||
Forbidden,
|
||||
#[error("credential has expired")]
|
||||
Expired,
|
||||
#[error("authenticated request context is missing")]
|
||||
MissingIdentity,
|
||||
#[error("authentication service unavailable")]
|
||||
Unavailable,
|
||||
#[error("gateway auth not configured")]
|
||||
Unconfigured,
|
||||
#[error("missing or invalid bearer token")]
|
||||
|
|
@ -34,9 +30,8 @@ pub enum Error {
|
|||
impl Error {
|
||||
pub fn status(&self) -> StatusCode {
|
||||
match self {
|
||||
Self::InvalidToken | Self::Expired => StatusCode::UNAUTHORIZED,
|
||||
Self::InvalidToken => StatusCode::UNAUTHORIZED,
|
||||
Self::Forbidden => StatusCode::FORBIDDEN,
|
||||
Self::Unavailable => StatusCode::SERVICE_UNAVAILABLE,
|
||||
Self::MissingIdentity | Self::Unconfigured | Self::Secret(_) => {
|
||||
StatusCode::INTERNAL_SERVER_ERROR
|
||||
}
|
||||
|
|
|
|||
113
litellm-rust/crates/gateway-auth/src/extract.rs
Normal file
113
litellm-rust/crates/gateway-auth/src/extract.rs
Normal file
|
|
@ -0,0 +1,113 @@
|
|||
use std::iter;
|
||||
|
||||
use axum::{
|
||||
extract::FromRequestParts,
|
||||
http::{HeaderName, HeaderValue, request::Parts},
|
||||
};
|
||||
use axum_extra::{
|
||||
TypedHeader,
|
||||
headers::{
|
||||
self, Header,
|
||||
authorization::{Authorization, Credentials},
|
||||
},
|
||||
};
|
||||
use litellm_auth_types::SecretValue;
|
||||
|
||||
use crate::Error;
|
||||
|
||||
pub struct Bearer(pub SecretValue);
|
||||
|
||||
impl From<Bearer> for SecretValue {
|
||||
fn from(value: Bearer) -> Self {
|
||||
value.0
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: Send + Sync> FromRequestParts<S> for Bearer {
|
||||
type Rejection = Error;
|
||||
|
||||
async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Error> {
|
||||
let TypedHeader(Single(Authorization(BearerToken(token)))) =
|
||||
TypedHeader::from_request_parts(parts, state)
|
||||
.await
|
||||
.map_err(|_| Error::InvalidToken)?;
|
||||
Ok(Self(token))
|
||||
}
|
||||
}
|
||||
|
||||
struct BearerToken(SecretValue);
|
||||
|
||||
impl Credentials for BearerToken {
|
||||
const SCHEME: &'static str = "Bearer";
|
||||
|
||||
fn decode(value: &HeaderValue) -> Option<Self> {
|
||||
let token = value.to_str().ok()?.get(Self::SCHEME.len()..)?.trim();
|
||||
is_token(token).then(|| Self(SecretValue::new(token)))
|
||||
}
|
||||
|
||||
fn encode(&self) -> HeaderValue {
|
||||
HeaderValue::from_str(&format!("{} {}", Self::SCHEME, self.0.expose()))
|
||||
.expect("token was decoded from a valid header value")
|
||||
}
|
||||
}
|
||||
|
||||
struct Single<H>(H);
|
||||
|
||||
impl<H: Header> Header for Single<H> {
|
||||
fn name() -> &'static HeaderName {
|
||||
H::name()
|
||||
}
|
||||
|
||||
fn decode<'i, I: Iterator<Item = &'i HeaderValue>>(
|
||||
values: &mut I,
|
||||
) -> Result<Self, headers::Error> {
|
||||
match (values.next(), values.next()) {
|
||||
(Some(value), None) => H::decode(&mut iter::once(value)).map(Self),
|
||||
_ => Err(headers::Error::invalid()),
|
||||
}
|
||||
}
|
||||
|
||||
fn encode<E: Extend<HeaderValue>>(&self, values: &mut E) {
|
||||
self.0.encode(values);
|
||||
}
|
||||
}
|
||||
|
||||
fn is_token(token: &str) -> bool {
|
||||
!token.is_empty() && !token.contains(char::is_whitespace)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use axum::http::Request;
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::standard(&["Bearer token"], Some("token"))]
|
||||
#[case::scheme_is_case_insensitive(&["bearer token"], Some("token"))]
|
||||
#[case::surrounding_space_is_trimmed(&["Bearer token "], Some("token"))]
|
||||
#[case::empty(&["Bearer "], None)]
|
||||
#[case::inner_whitespace(&["Bearer a b"], None)]
|
||||
#[case::other_scheme(&["Basic token"], None)]
|
||||
#[case::duplicate(&["Bearer token", "Bearer token"], None)]
|
||||
#[tokio::test]
|
||||
async fn bearer_accepts_one_rfc_6750_token(
|
||||
#[case] headers: &[&str],
|
||||
#[case] expected: Option<&str>,
|
||||
) {
|
||||
let request = headers
|
||||
.iter()
|
||||
.fold(Request::builder(), |request, value| {
|
||||
request.header("authorization", *value)
|
||||
})
|
||||
.body(())
|
||||
.unwrap();
|
||||
let (mut parts, ()) = request.into_parts();
|
||||
let token = Bearer::from_request_parts(&mut parts, &())
|
||||
.await
|
||||
.ok()
|
||||
.map(|Bearer(token)| token.expose().to_owned());
|
||||
assert_eq!(token.as_deref(), expected);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,90 +1,19 @@
|
|||
use axum::{
|
||||
extract::{FromRequestParts, MatchedPath, Request, State},
|
||||
http::{header::AUTHORIZATION, request::Parts},
|
||||
extract::{Request, State},
|
||||
middleware::Next,
|
||||
response::Response,
|
||||
};
|
||||
use litellm_auth_types::SecretValue;
|
||||
|
||||
use crate::{AccessRequest, Auth, AuthenticatedRequest, Credential, Error};
|
||||
use crate::{Auth, Error};
|
||||
|
||||
pub trait CredentialExtractor: Send + Sync {
|
||||
fn extract(&self, parts: &Parts) -> Result<Credential, Error>;
|
||||
}
|
||||
|
||||
pub struct Bearer;
|
||||
|
||||
impl CredentialExtractor for Bearer {
|
||||
fn extract(&self, parts: &Parts) -> Result<Credential, Error> {
|
||||
bearer(parts).map(Credential::Token)
|
||||
}
|
||||
}
|
||||
|
||||
fn bearer(parts: &Parts) -> Result<SecretValue, Error> {
|
||||
let mut values = parts.headers.get_all(AUTHORIZATION).iter();
|
||||
let value = values.next().ok_or(Error::InvalidToken)?;
|
||||
if values.next().is_some() {
|
||||
return Err(Error::InvalidToken);
|
||||
}
|
||||
let token = value
|
||||
.to_str()
|
||||
.ok()
|
||||
.and_then(|value| value.strip_prefix("Bearer "))
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty() && !value.bytes().any(|byte| byte.is_ascii_whitespace()))
|
||||
.ok_or(Error::InvalidToken)?;
|
||||
Ok(SecretValue::new(token))
|
||||
}
|
||||
|
||||
pub async fn authenticate(
|
||||
pub async fn authenticate<E: Into<SecretValue>>(
|
||||
State(auth): State<Auth>,
|
||||
request: Request,
|
||||
credential: E,
|
||||
mut request: Request,
|
||||
next: Next,
|
||||
) -> Result<Response, Error> {
|
||||
let (mut parts, body) = request.into_parts();
|
||||
let identity = auth.authenticate_parts(&parts).await?;
|
||||
let route = AccessRequest::Route {
|
||||
method: parts.method.to_string(),
|
||||
path: parts
|
||||
.extensions
|
||||
.get::<MatchedPath>()
|
||||
.map_or(parts.uri.path(), MatchedPath::as_str)
|
||||
.into(),
|
||||
};
|
||||
identity
|
||||
.authorize(route.clone())
|
||||
.await?
|
||||
.consume(identity.caller(), &route)?;
|
||||
parts.extensions.insert(identity);
|
||||
Ok(next.run(Request::from_parts(parts, body)).await)
|
||||
}
|
||||
|
||||
impl<S: Send + Sync> FromRequestParts<S> for AuthenticatedRequest {
|
||||
type Rejection = Error;
|
||||
|
||||
async fn from_request_parts(parts: &mut Parts, _: &S) -> Result<Self, Error> {
|
||||
parts
|
||||
.extensions
|
||||
.get::<Self>()
|
||||
.cloned()
|
||||
.ok_or(Error::MissingIdentity)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct RequireMasterKey;
|
||||
|
||||
impl FromRequestParts<Auth> for RequireMasterKey {
|
||||
type Rejection = Error;
|
||||
|
||||
async fn from_request_parts(parts: &mut Parts, state: &Auth) -> Result<Self, Error> {
|
||||
state.validate().await?;
|
||||
let identity = state
|
||||
.authenticate_request(&Credential::Token(bearer(parts)?), parts)
|
||||
.await?;
|
||||
if identity.caller().authentication().method != crate::AuthenticationMethod::MasterKey {
|
||||
return Err(Error::InvalidToken);
|
||||
}
|
||||
parts.extensions.insert(identity);
|
||||
Ok(Self)
|
||||
}
|
||||
let caller = auth.authenticate(&credential.into()).await?;
|
||||
request.extensions_mut().insert(caller);
|
||||
Ok(next.run(request).await)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,122 +0,0 @@
|
|||
use std::{sync::Arc, time::SystemTime};
|
||||
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::authorization::Permissions;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum PrincipalKind {
|
||||
Human,
|
||||
Service,
|
||||
System,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub struct Principal {
|
||||
authority: String,
|
||||
subject: String,
|
||||
kind: PrincipalKind,
|
||||
}
|
||||
|
||||
impl Principal {
|
||||
pub fn new(authority: String, subject: String, kind: PrincipalKind) -> Self {
|
||||
Self {
|
||||
authority,
|
||||
subject,
|
||||
kind,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn authority(&self) -> &str {
|
||||
&self.authority
|
||||
}
|
||||
pub fn subject(&self) -> &str {
|
||||
&self.subject
|
||||
}
|
||||
pub fn kind(&self) -> PrincipalKind {
|
||||
self.kind
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub enum AuthenticationMethod {
|
||||
MasterKey,
|
||||
Session,
|
||||
External(String),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Authentication {
|
||||
pub method: AuthenticationMethod,
|
||||
pub verifier: String,
|
||||
pub credential_id: String,
|
||||
pub expires_at: Option<SystemTime>,
|
||||
}
|
||||
|
||||
pub struct VerifiedIdentity {
|
||||
pub principal: Principal,
|
||||
pub authentication: Authentication,
|
||||
pub restrictions: Permissions,
|
||||
}
|
||||
|
||||
pub struct ResolvedIdentity {
|
||||
pub principal: Principal,
|
||||
pub permissions: Permissions,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct AuthenticatedCaller {
|
||||
verified_principal: Principal,
|
||||
principal: Principal,
|
||||
authentication: Authentication,
|
||||
permissions: Permissions,
|
||||
restrictions: Permissions,
|
||||
}
|
||||
|
||||
impl AuthenticatedCaller {
|
||||
pub(crate) fn new(verified: VerifiedIdentity, resolved: ResolvedIdentity) -> Self {
|
||||
Self {
|
||||
verified_principal: verified.principal,
|
||||
principal: resolved.principal,
|
||||
authentication: verified.authentication,
|
||||
permissions: resolved.permissions,
|
||||
restrictions: verified.restrictions,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn verified_principal(&self) -> &Principal {
|
||||
&self.verified_principal
|
||||
}
|
||||
|
||||
pub fn principal(&self) -> &Principal {
|
||||
&self.principal
|
||||
}
|
||||
pub fn authentication(&self) -> &Authentication {
|
||||
&self.authentication
|
||||
}
|
||||
pub fn permissions(&self) -> &Permissions {
|
||||
&self.permissions
|
||||
}
|
||||
pub fn restrictions(&self) -> &Permissions {
|
||||
&self.restrictions
|
||||
}
|
||||
|
||||
pub fn session_owner(&self) -> String {
|
||||
let fields = [
|
||||
self.verified_principal.authority(),
|
||||
self.verified_principal.subject(),
|
||||
self.principal.authority(),
|
||||
self.principal.subject(),
|
||||
self.authentication.verifier.as_str(),
|
||||
self.authentication.credential_id.as_str(),
|
||||
];
|
||||
let digest = fields.into_iter().fold(Sha256::new(), |mut digest, field| {
|
||||
digest.update((field.len() as u64).to_be_bytes());
|
||||
digest.update(field.as_bytes());
|
||||
digest
|
||||
});
|
||||
format!("{:x}", digest.finalize())
|
||||
}
|
||||
}
|
||||
|
||||
pub type SharedCaller = Arc<AuthenticatedCaller>;
|
||||
|
|
@ -1,27 +1,20 @@
|
|||
mod authentication;
|
||||
mod authorization;
|
||||
mod caller;
|
||||
mod custom;
|
||||
mod error;
|
||||
pub mod extract;
|
||||
mod http;
|
||||
mod identity;
|
||||
pub mod keys;
|
||||
mod ui;
|
||||
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
pub use authentication::{
|
||||
Auth, AuthFuture, Authenticator, Clock, Credential, IdentityResolver, LocalAdministrator,
|
||||
MasterKeyAuthenticator, SystemClock,
|
||||
};
|
||||
pub use authorization::{
|
||||
AccessRequest, AuthenticatedRequest, AuthorizedOperation, Authorizer, McpAction,
|
||||
NoAdditionalPolicy, Permissions, UiAction,
|
||||
};
|
||||
pub use authentication::Auth;
|
||||
pub use caller::{AccessRequest, Caller, McpAction};
|
||||
pub use custom::{CustomAuth, CustomOutcome};
|
||||
pub use error::{Error, KeyError, UiAuthError};
|
||||
pub use http::{Bearer, CredentialExtractor, RequireMasterKey, authenticate};
|
||||
pub use identity::{
|
||||
AuthenticatedCaller, Authentication, AuthenticationMethod, Principal, PrincipalKind,
|
||||
ResolvedIdentity, SharedCaller, VerifiedIdentity,
|
||||
};
|
||||
pub use extract::Bearer;
|
||||
pub use http::authenticate;
|
||||
pub use ui::{UI_CSRF_KEY, UiAuthSession, UiBackend, UiCredentials, UiSession, UiUser};
|
||||
|
||||
pub fn hash_token(token: &str) -> String {
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
use std::{convert::Infallible, sync::Arc};
|
||||
use std::convert::Infallible;
|
||||
|
||||
use axum::extract::FromRequestParts;
|
||||
use axum::http::{header::AUTHORIZATION, request::Parts};
|
||||
use axum::{extract::FromRequestParts, http::request::Parts};
|
||||
use axum_login::{AuthUser, AuthnBackend, UserId};
|
||||
use litellm_auth_types::SecretValue;
|
||||
use serde::Deserialize;
|
||||
|
|
@ -10,11 +9,7 @@ use subtle::ConstantTimeEq;
|
|||
use tower_sessions::Session;
|
||||
use veil::Redact;
|
||||
|
||||
use crate::{
|
||||
AccessRequest, AuthenticatedRequest, Authentication, AuthenticationMethod, LocalAdministrator,
|
||||
NoAdditionalPolicy, Permissions, Principal, PrincipalKind, SystemClock, UiAction, UiAuthError,
|
||||
VerifiedIdentity,
|
||||
};
|
||||
use crate::{Bearer, UiAuthError};
|
||||
|
||||
pub const UI_CSRF_KEY: &str = "litellm.ui.csrf";
|
||||
pub type UiAuthSession = axum_login::AuthSession<UiBackend>;
|
||||
|
|
@ -85,7 +80,6 @@ impl AuthnBackend for UiBackend {
|
|||
|
||||
pub struct UiSession {
|
||||
pub user: UiUser,
|
||||
pub identity: AuthenticatedRequest,
|
||||
}
|
||||
|
||||
impl<S: Send + Sync> FromRequestParts<S> for UiSession {
|
||||
|
|
@ -103,84 +97,12 @@ impl<S: Send + Sync> FromRequestParts<S> for UiSession {
|
|||
.get::<String>(UI_CSRF_KEY)
|
||||
.await?
|
||||
.ok_or(UiAuthError::Unauthorized)?;
|
||||
let provided = parts
|
||||
.headers
|
||||
.get(AUTHORIZATION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| value.strip_prefix("Bearer "))
|
||||
.ok_or(UiAuthError::Unauthorized)?;
|
||||
if !bool::from(Sha256::digest(expected).ct_eq(&Sha256::digest(provided))) {
|
||||
let Bearer(provided) = Bearer::from_request_parts(parts, state)
|
||||
.await
|
||||
.map_err(|_| UiAuthError::Unauthorized)?;
|
||||
if !bool::from(Sha256::digest(expected).ct_eq(&Sha256::digest(provided.expose()))) {
|
||||
return Err(UiAuthError::Unauthorized);
|
||||
}
|
||||
let identity = session_identity(&user).await?;
|
||||
parts.extensions.insert(identity.clone());
|
||||
Ok(Self { user, identity })
|
||||
}
|
||||
}
|
||||
|
||||
async fn session_identity(user: &UiUser) -> Result<AuthenticatedRequest, UiAuthError> {
|
||||
crate::authentication::resolve(
|
||||
VerifiedIdentity {
|
||||
principal: Principal::new(
|
||||
"litellm:local-ui".into(),
|
||||
user.username.clone(),
|
||||
PrincipalKind::Human,
|
||||
),
|
||||
authentication: Authentication {
|
||||
method: AuthenticationMethod::Session,
|
||||
verifier: "litellm:local-ui".into(),
|
||||
credential_id: user.username.clone(),
|
||||
expires_at: None,
|
||||
},
|
||||
restrictions: Permissions::Only(Arc::from([
|
||||
AccessRequest::Ui(UiAction::SessionInfo),
|
||||
AccessRequest::Ui(UiAction::Logout),
|
||||
])),
|
||||
},
|
||||
&LocalAdministrator,
|
||||
Arc::new(NoAdditionalPolicy),
|
||||
Arc::new(SystemClock),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| UiAuthError::Unauthorized)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use rstest::rstest;
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn ui_identity_is_scoped_to_session_operations() {
|
||||
let backend = UiBackend::new("admin".into(), SecretValue::new("password")).unwrap();
|
||||
let user = backend.get_user(&"admin".into()).await.unwrap().unwrap();
|
||||
let identity = session_identity(&user).await.unwrap();
|
||||
assert_eq!(identity.caller().principal().subject(), user.username);
|
||||
assert_eq!(
|
||||
identity.caller().authentication().method,
|
||||
AuthenticationMethod::Session
|
||||
);
|
||||
assert!(
|
||||
identity
|
||||
.authorize(AccessRequest::Ui(UiAction::SessionInfo))
|
||||
.await
|
||||
.is_ok()
|
||||
);
|
||||
assert!(
|
||||
identity
|
||||
.authorize(AccessRequest::Ui(UiAction::Logout))
|
||||
.await
|
||||
.is_ok()
|
||||
);
|
||||
assert!(matches!(
|
||||
identity
|
||||
.authorize(AccessRequest::Model {
|
||||
name: "model".into(),
|
||||
deployment: "provider/model".into(),
|
||||
})
|
||||
.await,
|
||||
Err(crate::Error::Forbidden)
|
||||
));
|
||||
Ok(Self { user })
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,13 +4,13 @@ use axum::{
|
|||
Router,
|
||||
body::{Body, to_bytes},
|
||||
http::{Request, StatusCode},
|
||||
middleware::from_extractor_with_state,
|
||||
middleware::from_fn_with_state,
|
||||
routing::get,
|
||||
};
|
||||
use futures_util::future::BoxFuture;
|
||||
use litellm_auth_types::SecretValue;
|
||||
use litellm_config::Config;
|
||||
use litellm_gateway_auth::{Auth, RequireMasterKey, hash_token};
|
||||
use litellm_gateway_auth::{Auth, Bearer, Caller, Error, authenticate, hash_token};
|
||||
use litellm_secrets::source::SecretSource;
|
||||
use rstest::{fixture, rstest};
|
||||
use tower::ServiceExt;
|
||||
|
|
@ -53,9 +53,9 @@ fn secrets() -> Arc<dyn SecretSource> {
|
|||
#[case::wrong_scheme("literal-key", Some("Basic literal-key"), 401)]
|
||||
#[case::empty_token("literal-key", Some("Bearer "), 401)]
|
||||
#[case::missing_reference("os.environ/MISSING", Some("Bearer os.environ/MISSING"), 500)]
|
||||
#[case::empty_reference("os.environ/EMPTY", Some("Bearer "), 500)]
|
||||
#[case::empty_key("", Some("Bearer "), 500)]
|
||||
#[case::whitespace_key(" ", Some("Bearer "), 500)]
|
||||
#[case::empty_reference("os.environ/EMPTY", Some("Bearer token"), 500)]
|
||||
#[case::empty_key("", Some("Bearer token"), 500)]
|
||||
#[case::whitespace_key(" ", Some("Bearer token"), 500)]
|
||||
#[case::empty_reference_name("os.environ/", Some("Bearer os.environ/"), 500)]
|
||||
#[case::secret_failure("os.environ/ERROR", Some("Bearer private-backend-detail"), 500)]
|
||||
#[tokio::test]
|
||||
|
|
@ -71,8 +71,9 @@ async fn enforces_configured_keys_without_exposing_secrets(
|
|||
.unwrap();
|
||||
let app = Router::new()
|
||||
.route("/protected", get(|| async { StatusCode::NO_CONTENT }))
|
||||
.layer(from_extractor_with_state::<RequireMasterKey, _>(
|
||||
.layer(from_fn_with_state(
|
||||
Auth::from_config(&config, secrets),
|
||||
authenticate::<Bearer>,
|
||||
));
|
||||
let request = Request::get("/protected");
|
||||
let request = match authorization {
|
||||
|
|
@ -99,368 +100,109 @@ fn hash_token_matches_python_sha256_hexdigest() {
|
|||
);
|
||||
}
|
||||
|
||||
use litellm_gateway_auth::{
|
||||
AccessRequest, AuthFuture, AuthenticatedCaller, AuthenticatedRequest, Authentication,
|
||||
AuthenticationMethod, Authenticator, Authorizer, Clock, Error, IdentityResolver, Permissions,
|
||||
Principal, PrincipalKind, ResolvedIdentity, VerifiedIdentity, authenticate,
|
||||
};
|
||||
use std::{
|
||||
sync::atomic::{AtomicU64, AtomicUsize, Ordering},
|
||||
time::{Duration, SystemTime},
|
||||
};
|
||||
|
||||
struct TestClock(AtomicU64);
|
||||
|
||||
impl Clock for TestClock {
|
||||
fn now(&self) -> SystemTime {
|
||||
SystemTime::UNIX_EPOCH + Duration::from_secs(self.0.load(Ordering::SeqCst))
|
||||
}
|
||||
}
|
||||
|
||||
struct Verifier {
|
||||
authority: &'static str,
|
||||
restrictions: Permissions,
|
||||
failure: Option<u16>,
|
||||
}
|
||||
|
||||
impl Authenticator for Verifier {
|
||||
fn verify<'a>(
|
||||
&'a self,
|
||||
credential: &'a litellm_gateway_auth::Credential,
|
||||
_: &'a axum::http::request::Parts,
|
||||
) -> AuthFuture<'a, VerifiedIdentity> {
|
||||
Box::pin(async move {
|
||||
match self.failure {
|
||||
Some(401) => return Err(Error::InvalidToken),
|
||||
Some(_) => return Err(Error::Unavailable),
|
||||
None => (),
|
||||
}
|
||||
let litellm_gateway_auth::Credential::Token(credential) = credential else {
|
||||
return Err(Error::InvalidToken);
|
||||
};
|
||||
if credential.expose() != "test-token" {
|
||||
return Err(Error::InvalidToken);
|
||||
}
|
||||
Ok(VerifiedIdentity {
|
||||
principal: Principal::new(
|
||||
self.authority.into(),
|
||||
"subject".into(),
|
||||
PrincipalKind::Service,
|
||||
),
|
||||
authentication: Authentication {
|
||||
method: AuthenticationMethod::External("test".into()),
|
||||
verifier: "test-verifier".into(),
|
||||
credential_id: "credential".into(),
|
||||
expires_at: Some(SystemTime::UNIX_EPOCH + Duration::from_secs(100)),
|
||||
},
|
||||
restrictions: self.restrictions.clone(),
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct Identities {
|
||||
mapped_principal: Option<Principal>,
|
||||
calls: AtomicUsize,
|
||||
permissions: Permissions,
|
||||
}
|
||||
|
||||
impl IdentityResolver for Identities {
|
||||
fn resolve<'a>(&'a self, identity: &'a VerifiedIdentity) -> AuthFuture<'a, ResolvedIdentity> {
|
||||
Box::pin(async move {
|
||||
self.calls.fetch_add(1, Ordering::SeqCst);
|
||||
Ok(ResolvedIdentity {
|
||||
principal: self
|
||||
.mapped_principal
|
||||
.clone()
|
||||
.unwrap_or_else(|| identity.principal.clone()),
|
||||
permissions: self.permissions.clone(),
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct Policy(AtomicUsize);
|
||||
|
||||
impl Authorizer for Policy {
|
||||
fn authorize<'a>(
|
||||
&'a self,
|
||||
_: &'a AuthenticatedCaller,
|
||||
_: &'a AccessRequest,
|
||||
) -> AuthFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
self.0.fetch_add(1, Ordering::SeqCst);
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct Pipeline {
|
||||
auth: Auth,
|
||||
identities: Arc<Identities>,
|
||||
policy: Arc<Policy>,
|
||||
clock: Arc<TestClock>,
|
||||
}
|
||||
|
||||
fn pipeline(
|
||||
authority: &'static str,
|
||||
permissions: Permissions,
|
||||
restrictions: Permissions,
|
||||
failure: Option<u16>,
|
||||
) -> Pipeline {
|
||||
let identities = Arc::new(Identities {
|
||||
mapped_principal: None,
|
||||
calls: AtomicUsize::new(0),
|
||||
permissions,
|
||||
});
|
||||
let policy = Arc::new(Policy(AtomicUsize::new(0)));
|
||||
let clock = Arc::new(TestClock(AtomicU64::new(10)));
|
||||
let auth = Auth::new(
|
||||
Arc::new(Verifier {
|
||||
authority,
|
||||
restrictions,
|
||||
failure,
|
||||
}),
|
||||
identities.clone(),
|
||||
policy.clone(),
|
||||
clock.clone(),
|
||||
);
|
||||
Pipeline {
|
||||
auth,
|
||||
identities,
|
||||
policy,
|
||||
clock,
|
||||
}
|
||||
}
|
||||
|
||||
fn model(name: &str) -> AccessRequest {
|
||||
AccessRequest::Model {
|
||||
name: name.into(),
|
||||
deployment: name.into(),
|
||||
}
|
||||
fn master_auth(secrets: Arc<dyn SecretSource>, key: &str) -> Auth {
|
||||
let config = Config::from_yaml(&format!(
|
||||
"model_list: []\ngeneral_settings:\n master_key: {key}\n"
|
||||
))
|
||||
.unwrap();
|
||||
Auth::from_config(&config, secrets)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::invalid_token(401)]
|
||||
#[case::unavailable_verifier(503)]
|
||||
#[case::valid(Some("Bearer master-key"), None, 204)]
|
||||
#[case::duplicate(Some("Bearer master-key"), Some("Bearer master-key"), 401)]
|
||||
#[case::forged_extension_without_credentials(None, None, 401)]
|
||||
#[tokio::test]
|
||||
async fn failed_verification_never_resolves_identity_or_runs_the_handler(#[case] status: u16) {
|
||||
let pipeline = pipeline("issuer-a", Permissions::All, Permissions::All, Some(status));
|
||||
async fn middleware_only_trusts_the_presented_credential(
|
||||
secrets: Arc<dyn SecretSource>,
|
||||
#[case] first: Option<&str>,
|
||||
#[case] second: Option<&str>,
|
||||
#[case] status: u16,
|
||||
) {
|
||||
let app = Router::new()
|
||||
.route("/protected", get(|| async { StatusCode::NO_CONTENT }))
|
||||
.layer(axum::middleware::from_fn_with_state(
|
||||
pipeline.auth,
|
||||
authenticate,
|
||||
.route(
|
||||
"/protected",
|
||||
get(|caller: Caller| async move {
|
||||
assert_eq!(caller, Caller::Master);
|
||||
StatusCode::NO_CONTENT
|
||||
}),
|
||||
)
|
||||
.layer(from_fn_with_state(
|
||||
master_auth(secrets, "master-key"),
|
||||
authenticate::<Bearer>,
|
||||
));
|
||||
let request = [first, second]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.fold(Request::get("/protected"), |request, value| {
|
||||
request.header("authorization", value)
|
||||
})
|
||||
.extension(Caller::Master);
|
||||
let response = app
|
||||
.oneshot(request.body(Body::empty()).unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status().as_u16(), status);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn handlers_fail_closed_without_the_middleware() {
|
||||
let app = Router::new().route(
|
||||
"/protected",
|
||||
get(|_: Caller| async { StatusCode::NO_CONTENT }),
|
||||
);
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::get("/protected")
|
||||
.header("authorization", "Bearer test-token")
|
||||
.header("authorization", "Bearer master-key")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status().as_u16(), status);
|
||||
assert_eq!(pipeline.identities.calls.load(Ordering::SeqCst), 0);
|
||||
assert_eq!(pipeline.policy.0.load(Ordering::SeqCst), 0);
|
||||
assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::membership_restricts(Permissions::Only(Arc::from([model("allowed")])), Permissions::All)]
|
||||
#[case::credential_restricts(Permissions::All, Permissions::Only(Arc::from([model("allowed")])))]
|
||||
#[tokio::test]
|
||||
async fn policy_cannot_expand_membership_or_credential_permissions(
|
||||
#[case] permissions: Permissions,
|
||||
#[case] restrictions: Permissions,
|
||||
) {
|
||||
let pipeline = pipeline("issuer-a", permissions, restrictions, None);
|
||||
let identity = pipeline
|
||||
.auth
|
||||
.authenticate(&SecretValue::new("test-token"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
identity.authorize(model("denied")).await,
|
||||
Err(Error::Forbidden)
|
||||
));
|
||||
assert_eq!(pipeline.policy.0.load(Ordering::SeqCst), 0);
|
||||
let allowed = model("allowed");
|
||||
identity
|
||||
.authorize(allowed.clone())
|
||||
.await
|
||||
.unwrap()
|
||||
.consume(identity.caller(), &allowed)
|
||||
.unwrap();
|
||||
assert_eq!(pipeline.policy.0.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
struct CustomHeader(SecretValue);
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn authorization_grants_bind_the_exact_caller_and_target() {
|
||||
let first = pipeline("issuer-a", Permissions::All, Permissions::All, None);
|
||||
let second = pipeline("issuer-b", Permissions::All, Permissions::All, None);
|
||||
let caller = first
|
||||
.auth
|
||||
.authenticate(&SecretValue::new("test-token"))
|
||||
.await
|
||||
.unwrap();
|
||||
let other = second
|
||||
.auth
|
||||
.authenticate(&SecretValue::new("test-token"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_ne!(caller.caller().principal(), other.caller().principal());
|
||||
assert_ne!(
|
||||
caller.caller().session_owner(),
|
||||
other.caller().session_owner()
|
||||
);
|
||||
let access = model("allowed");
|
||||
assert!(matches!(
|
||||
caller
|
||||
.authorize(access.clone())
|
||||
.await
|
||||
.unwrap()
|
||||
.consume(other.caller(), &access),
|
||||
Err(Error::Forbidden)
|
||||
));
|
||||
assert!(matches!(
|
||||
caller
|
||||
.authorize(access)
|
||||
.await
|
||||
.unwrap()
|
||||
.consume(caller.caller(), &model("different")),
|
||||
Err(Error::Forbidden)
|
||||
));
|
||||
}
|
||||
impl axum::extract::FromRequestParts<Auth> for CustomHeader {
|
||||
type Rejection = Error;
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn expiry_is_checked_again_for_operations_on_existing_connections() {
|
||||
let pipeline = pipeline("issuer-a", Permissions::All, Permissions::All, None);
|
||||
let identity = pipeline
|
||||
.auth
|
||||
.authenticate(&SecretValue::new("test-token"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(identity.authorize(model("allowed")).await.is_ok());
|
||||
pipeline.clock.0.store(100, Ordering::SeqCst);
|
||||
assert!(matches!(
|
||||
identity.authorize(model("allowed")).await,
|
||||
Err(Error::Expired)
|
||||
));
|
||||
assert!(matches!(
|
||||
pipeline
|
||||
.auth
|
||||
.authenticate(&SecretValue::new("test-token"))
|
||||
.await,
|
||||
Err(Error::Expired)
|
||||
));
|
||||
assert_eq!(pipeline.identities.calls.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(pipeline.policy.0.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn bearer_extraction_rejects_duplicates_and_does_not_trust_preexisting_identity() {
|
||||
let pipeline = pipeline("issuer-a", Permissions::All, Permissions::All, None);
|
||||
let identity = pipeline
|
||||
.auth
|
||||
.authenticate(&SecretValue::new("test-token"))
|
||||
.await
|
||||
.unwrap();
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/protected",
|
||||
get(|_: AuthenticatedRequest| async { StatusCode::NO_CONTENT }),
|
||||
)
|
||||
.layer(axum::middleware::from_fn_with_state(
|
||||
pipeline.auth,
|
||||
authenticate,
|
||||
));
|
||||
let duplicate = Request::get("/protected")
|
||||
.header("authorization", "Bearer test-token")
|
||||
.header("authorization", "Bearer test-token")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
assert_eq!(app.clone().oneshot(duplicate).await.unwrap().status(), 401);
|
||||
let forged = Request::get("/protected")
|
||||
.extension(identity)
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
assert_eq!(app.clone().oneshot(forged).await.unwrap().status(), 401);
|
||||
let valid = Request::get("/protected")
|
||||
.header("authorization", "Bearer test-token")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
assert_eq!(app.oneshot(valid).await.unwrap().status(), 204);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn route_permissions_are_checked_before_parsing_or_dispatch() {
|
||||
let allowed = AccessRequest::Route {
|
||||
method: "GET".into(),
|
||||
path: "/allowed/{id}".into(),
|
||||
};
|
||||
let pipeline = pipeline(
|
||||
"issuer-a",
|
||||
Permissions::Only(Arc::from([allowed])),
|
||||
Permissions::All,
|
||||
None,
|
||||
);
|
||||
let app = Router::new()
|
||||
.route("/allowed/{id}", get(|| async { StatusCode::NO_CONTENT }))
|
||||
.route("/denied", get(|| async { StatusCode::NO_CONTENT }))
|
||||
.layer(axum::middleware::from_fn_with_state(
|
||||
pipeline.auth,
|
||||
authenticate,
|
||||
));
|
||||
let allowed = Request::get("/allowed/123")
|
||||
.header("authorization", "Bearer test-token")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
let denied = Request::get("/denied")
|
||||
.header("authorization", "Bearer test-token")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
assert_eq!(app.clone().oneshot(allowed).await.unwrap().status(), 204);
|
||||
assert_eq!(app.oneshot(denied).await.unwrap().status(), 403);
|
||||
}
|
||||
|
||||
struct CustomHeader;
|
||||
|
||||
impl litellm_gateway_auth::CredentialExtractor for CustomHeader {
|
||||
fn extract(
|
||||
&self,
|
||||
parts: &axum::http::request::Parts,
|
||||
) -> Result<litellm_gateway_auth::Credential, Error> {
|
||||
async fn from_request_parts(
|
||||
parts: &mut axum::http::request::Parts,
|
||||
_: &Auth,
|
||||
) -> Result<Self, Error> {
|
||||
let value = parts
|
||||
.headers
|
||||
.get("x-test-key")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.ok_or(Error::InvalidToken)?;
|
||||
Ok(litellm_gateway_auth::Credential::Token(SecretValue::new(
|
||||
value,
|
||||
)))
|
||||
Ok(Self(SecretValue::new(value)))
|
||||
}
|
||||
}
|
||||
|
||||
impl From<CustomHeader> for SecretValue {
|
||||
fn from(value: CustomHeader) -> Self {
|
||||
value.0
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::selected_header("test-token", "wrong-token", 204)]
|
||||
#[case::no_fallback("wrong-token", "test-token", 401)]
|
||||
#[case::selected_header("master-key", "wrong-key", 204)]
|
||||
#[case::no_fallback("wrong-key", "master-key", 401)]
|
||||
#[tokio::test]
|
||||
async fn route_profile_uses_only_its_selected_credentials(
|
||||
secrets: Arc<dyn SecretSource>,
|
||||
#[case] custom: &str,
|
||||
#[case] bearer: &str,
|
||||
#[case] status: u16,
|
||||
) {
|
||||
let pipeline = pipeline("issuer-a", Permissions::All, Permissions::All, None);
|
||||
let app = Router::new()
|
||||
.route("/protected", get(|| async { StatusCode::NO_CONTENT }))
|
||||
.layer(axum::middleware::from_fn_with_state(
|
||||
pipeline.auth.with_extractor(Arc::new(CustomHeader)),
|
||||
authenticate,
|
||||
.layer(from_fn_with_state(
|
||||
master_auth(secrets, "master-key"),
|
||||
authenticate::<CustomHeader>,
|
||||
));
|
||||
let response = app
|
||||
.oneshot(
|
||||
|
|
@ -477,79 +219,17 @@ async fn route_profile_uses_only_its_selected_credentials(
|
|||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn master_rotation_preserves_session_owner_but_rejects_the_old_key(
|
||||
secrets: Arc<dyn SecretSource>,
|
||||
) {
|
||||
let first =
|
||||
Config::from_yaml("model_list: []\ngeneral_settings:\n master_key: first-key\n").unwrap();
|
||||
let second =
|
||||
Config::from_yaml("model_list: []\ngeneral_settings:\n master_key: second-key\n").unwrap();
|
||||
let original = Auth::from_config(&first, secrets.clone())
|
||||
.authenticate(&SecretValue::new("first-key"))
|
||||
.await
|
||||
.unwrap();
|
||||
let rotated = Auth::from_config(&second, secrets);
|
||||
let replacement = rotated
|
||||
.authenticate(&SecretValue::new("second-key"))
|
||||
.await
|
||||
.unwrap();
|
||||
async fn master_rotation_rejects_the_old_key(secrets: Arc<dyn SecretSource>) {
|
||||
let rotated = master_auth(secrets, "second-key");
|
||||
assert_eq!(
|
||||
original.caller().session_owner(),
|
||||
replacement.caller().session_owner()
|
||||
rotated
|
||||
.authenticate(&SecretValue::new("second-key"))
|
||||
.await
|
||||
.unwrap(),
|
||||
Caller::Master
|
||||
);
|
||||
assert!(matches!(
|
||||
rotated.authenticate(&SecretValue::new("first-key")).await,
|
||||
Err(Error::InvalidToken)
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn mapping_two_issuers_to_one_account_does_not_merge_session_ownership() {
|
||||
let resolver = Arc::new(Identities {
|
||||
mapped_principal: Some(Principal::new(
|
||||
"internal".into(),
|
||||
"shared-account".into(),
|
||||
PrincipalKind::Human,
|
||||
)),
|
||||
calls: AtomicUsize::new(0),
|
||||
permissions: Permissions::All,
|
||||
});
|
||||
let first = Auth::new(
|
||||
Arc::new(Verifier {
|
||||
authority: "issuer-a",
|
||||
restrictions: Permissions::All,
|
||||
failure: None,
|
||||
}),
|
||||
resolver.clone(),
|
||||
Arc::new(Policy(AtomicUsize::new(0))),
|
||||
Arc::new(TestClock(AtomicU64::new(10))),
|
||||
);
|
||||
let second = Auth::new(
|
||||
Arc::new(Verifier {
|
||||
authority: "issuer-b",
|
||||
restrictions: Permissions::All,
|
||||
failure: None,
|
||||
}),
|
||||
resolver,
|
||||
Arc::new(Policy(AtomicUsize::new(0))),
|
||||
Arc::new(TestClock(AtomicU64::new(10))),
|
||||
);
|
||||
let original = first
|
||||
.authenticate(&SecretValue::new("test-token"))
|
||||
.await
|
||||
.unwrap();
|
||||
let mapped = second
|
||||
.authenticate(&SecretValue::new("test-token"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(original.caller().principal(), mapped.caller().principal());
|
||||
assert_ne!(
|
||||
original.caller().verified_principal(),
|
||||
mapped.caller().verified_principal()
|
||||
);
|
||||
assert_ne!(
|
||||
original.caller().session_owner(),
|
||||
mapped.caller().session_owner()
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use litellm_gateway_auth::AuthenticatedRequest;
|
||||
use litellm_gateway_auth::Caller;
|
||||
use std::{path::Path, sync::Arc};
|
||||
|
||||
use axum::{Json, extract::State, response::IntoResponse};
|
||||
|
|
@ -13,22 +13,22 @@ use crate::{
|
|||
|
||||
pub(crate) async fn create(
|
||||
State(gateway): State<Arc<Gateway>>,
|
||||
identity: AuthenticatedRequest,
|
||||
caller: Caller,
|
||||
body: InferenceBody,
|
||||
) -> Result<impl IntoResponse, Error> {
|
||||
handle(&gateway, &identity, body).await.map(Json)
|
||||
handle(&gateway, &caller, body).await.map(Json)
|
||||
}
|
||||
|
||||
async fn handle(
|
||||
gateway: &Gateway,
|
||||
identity: &AuthenticatedRequest,
|
||||
caller: &Caller,
|
||||
InferenceBody {
|
||||
fields: body,
|
||||
upload,
|
||||
}: InferenceBody,
|
||||
) -> Result<Value, Error> {
|
||||
let deployment = request::resolve_deployment(gateway, &body)?;
|
||||
request::authorize_model(identity, deployment, &body).await?;
|
||||
request::authorize_model(caller, deployment, &body)?;
|
||||
let audio = match upload {
|
||||
Some(upload) => {
|
||||
let format = upload
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_cache_response::{CacheOptions, CachePolicy, CacheScope};
|
||||
use litellm_gateway_auth::AuthenticatedRequest;
|
||||
use litellm_gateway_auth::Caller;
|
||||
use serde::Deserialize;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
|
|
@ -21,10 +21,7 @@ struct Controls {
|
|||
|
||||
type Prepared = (Map<String, Value>, CacheOptions);
|
||||
|
||||
pub(crate) fn prepare(
|
||||
identity: &AuthenticatedRequest,
|
||||
body: Map<String, Value>,
|
||||
) -> Result<Prepared, Error> {
|
||||
pub(crate) fn prepare(caller: &Caller, body: Map<String, Value>) -> Result<Prepared, Error> {
|
||||
let controls: Controls = match body.get("cache").filter(|value| !value.is_null()) {
|
||||
Some(value) => serde_json::from_value(value.clone())
|
||||
.map_err(|error| Error::InvalidBody(error.to_string()))?,
|
||||
|
|
@ -36,7 +33,6 @@ pub(crate) fn prepare(
|
|||
.map(|value| serde_json::from_value(value.clone()))
|
||||
.transpose()
|
||||
.map_err(|error| Error::InvalidBody(error.to_string()))?;
|
||||
let caller = identity.caller();
|
||||
let options = CacheOptions {
|
||||
policy: CachePolicy {
|
||||
caching,
|
||||
|
|
@ -45,14 +41,7 @@ pub(crate) fn prepare(
|
|||
ttl: controls.ttl.map(duration).transpose()?,
|
||||
max_age: controls.max_age.map(duration).transpose()?,
|
||||
},
|
||||
scope: CacheScope::Isolated(
|
||||
serde_json::json!([
|
||||
caller.principal().authority(),
|
||||
caller.principal().subject(),
|
||||
caller.authentication().credential_id
|
||||
])
|
||||
.to_string(),
|
||||
),
|
||||
scope: CacheScope::Isolated(caller.owner()),
|
||||
};
|
||||
Ok((
|
||||
body.into_iter()
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use litellm_gateway_auth::AuthenticatedRequest;
|
||||
use litellm_gateway_auth::Caller;
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
|
|
@ -13,15 +13,15 @@ use crate::{Error, Gateway, JsonObject, request};
|
|||
|
||||
pub(crate) async fn create(
|
||||
State(gateway): State<Arc<Gateway>>,
|
||||
identity: AuthenticatedRequest,
|
||||
caller: Caller,
|
||||
JsonObject(body): JsonObject,
|
||||
) -> Result<impl IntoResponse, Error> {
|
||||
handle(&gateway, &identity, body).await
|
||||
handle(&gateway, &caller, body).await
|
||||
}
|
||||
|
||||
pub(crate) async fn create_from_model_path(
|
||||
State(gateway): State<Arc<Gateway>>,
|
||||
identity: AuthenticatedRequest,
|
||||
caller: Caller,
|
||||
Path(model): Path<String>,
|
||||
JsonObject(body): JsonObject,
|
||||
) -> Result<impl IntoResponse, Error> {
|
||||
|
|
@ -32,17 +32,17 @@ pub(crate) async fn create_from_model_path(
|
|||
.collect(),
|
||||
Some(_) => body,
|
||||
};
|
||||
handle(&gateway, &identity, body).await
|
||||
handle(&gateway, &caller, body).await
|
||||
}
|
||||
|
||||
async fn handle(
|
||||
gateway: &Gateway,
|
||||
identity: &AuthenticatedRequest,
|
||||
caller: &Caller,
|
||||
body: Map<String, Value>,
|
||||
) -> Result<Response, Error> {
|
||||
let deployment = request::resolve_deployment(gateway, &body)?;
|
||||
request::authorize_model(identity, deployment, &body).await?;
|
||||
let (body, cache_options) = crate::caching::prepare(identity, body)?;
|
||||
request::authorize_model(caller, deployment, &body)?;
|
||||
let (body, cache_options) = crate::caching::prepare(caller, body)?;
|
||||
let route = gateway.chat_completions.clone();
|
||||
let route = match &gateway.cache {
|
||||
Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new(
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
//! `POST /v1/messages`, as the Python proxy's `anthropic_response` serves it.
|
||||
|
||||
use litellm_gateway_auth::AuthenticatedRequest;
|
||||
use litellm_gateway_auth::Caller;
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
|
|
@ -28,13 +28,13 @@ const ANTHROPIC_API_HEADER_PROVIDERS: &str =
|
|||
|
||||
pub async fn create(
|
||||
State(gateway): State<Arc<Gateway>>,
|
||||
identity: AuthenticatedRequest,
|
||||
caller: Caller,
|
||||
RequestId(request_id): RequestId,
|
||||
headers: HeaderMap,
|
||||
body: Result<JsonObject, Error>,
|
||||
) -> impl IntoResponse {
|
||||
let result = match body {
|
||||
Ok(JsonObject(body)) => handle(&gateway, &identity, &headers, body).await,
|
||||
Ok(JsonObject(body)) => handle(&gateway, &caller, &headers, body).await,
|
||||
Err(error) => Err(error),
|
||||
};
|
||||
result.map_err(|error| (error.status(), Json(error.body(request_id.as_deref()))))
|
||||
|
|
@ -42,13 +42,13 @@ pub async fn create(
|
|||
|
||||
async fn handle(
|
||||
gateway: &Gateway,
|
||||
identity: &AuthenticatedRequest,
|
||||
caller: &Caller,
|
||||
headers: &HeaderMap,
|
||||
body: Map<String, Value>,
|
||||
) -> Result<Response, Error> {
|
||||
let deployment = request::resolve_deployment(gateway, &body)?;
|
||||
request::authorize_model(identity, deployment, &body).await?;
|
||||
let (body, cache_options) = crate::caching::prepare(identity, body)?;
|
||||
request::authorize_model(caller, deployment, &body)?;
|
||||
let (body, cache_options) = crate::caching::prepare(caller, body)?;
|
||||
let route = gateway.messages.clone();
|
||||
let route = match &gateway.cache {
|
||||
Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new(
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use litellm_gateway_auth::AuthenticatedRequest;
|
||||
use litellm_gateway_auth::Caller;
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{Json, extract::State, http::HeaderMap, response::IntoResponse};
|
||||
|
|
@ -15,16 +15,16 @@ use crate::{
|
|||
|
||||
pub(crate) async fn create(
|
||||
State(gateway): State<Arc<Gateway>>,
|
||||
identity: AuthenticatedRequest,
|
||||
caller: Caller,
|
||||
headers: HeaderMap,
|
||||
body: InferenceBody,
|
||||
) -> Result<impl IntoResponse, Error> {
|
||||
handle(&gateway, &identity, &headers, body).await.map(Json)
|
||||
handle(&gateway, &caller, &headers, body).await.map(Json)
|
||||
}
|
||||
|
||||
async fn handle(
|
||||
gateway: &Gateway,
|
||||
identity: &AuthenticatedRequest,
|
||||
caller: &Caller,
|
||||
headers: &HeaderMap,
|
||||
InferenceBody {
|
||||
fields: body,
|
||||
|
|
@ -36,7 +36,7 @@ async fn handle(
|
|||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_owned);
|
||||
let deployment = request::resolve_deployment(gateway, &body)?;
|
||||
request::authorize_model(identity, deployment, &body).await?;
|
||||
request::authorize_model(caller, deployment, &body)?;
|
||||
let document = match upload {
|
||||
Some(upload) => OcrDocumentInput::Bytes {
|
||||
bytes: upload.bytes,
|
||||
|
|
|
|||
|
|
@ -90,8 +90,8 @@ fn object(body: &[u8]) -> Result<Map<String, Value>, Error> {
|
|||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn authorize_model(
|
||||
identity: &litellm_gateway_auth::AuthenticatedRequest,
|
||||
pub(crate) fn authorize_model(
|
||||
caller: &litellm_gateway_auth::Caller,
|
||||
deployment: &Deployment,
|
||||
body: &Map<String, Value>,
|
||||
) -> Result<(), Error> {
|
||||
|
|
@ -103,11 +103,7 @@ pub(crate) async fn authorize_model(
|
|||
name: name.into(),
|
||||
deployment: deployment.model.clone(),
|
||||
};
|
||||
identity
|
||||
.authorize(access.clone())
|
||||
.await?
|
||||
.consume(identity.caller(), &access)?;
|
||||
Ok(())
|
||||
Ok(caller.authorize(&access)?)
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_deployment<'a>(
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use axum::{Json, body::Bytes, extract::State, response::Response};
|
||||
use litellm_gateway_auth::AuthenticatedRequest;
|
||||
use litellm_gateway_auth::Caller;
|
||||
use litellm_host_http::Sse;
|
||||
use litellm_inference_responses::{route::Responses, types::ResponsesCall};
|
||||
use serde_json::json;
|
||||
|
|
@ -10,12 +10,12 @@ use crate::{Error, Gateway, JsonObject, request};
|
|||
|
||||
pub(crate) async fn create(
|
||||
State(gateway): State<Arc<Gateway>>,
|
||||
identity: AuthenticatedRequest,
|
||||
caller: Caller,
|
||||
JsonObject(body): JsonObject,
|
||||
) -> Result<Response, Error> {
|
||||
let deployment = request::resolve_deployment(&gateway, &body)?;
|
||||
request::authorize_model(&identity, deployment, &body).await?;
|
||||
let (body, cache_options) = crate::caching::prepare(&identity, body)?;
|
||||
request::authorize_model(&caller, deployment, &body)?;
|
||||
let (body, cache_options) = crate::caching::prepare(&caller, body)?;
|
||||
let route = gateway.responses.clone();
|
||||
let route = match &gateway.cache {
|
||||
Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new(
|
||||
|
|
|
|||
|
|
@ -126,60 +126,3 @@ async fn all_inference_endpoints_share_native_cache(
|
|||
);
|
||||
assert_eq!(to_bytes(restored.into_body(), 4096).await.unwrap(), second);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::different_subject("issuer", "tenant-b")]
|
||||
#[case::different_authority("other-issuer", "tenant-a")]
|
||||
#[tokio::test]
|
||||
async fn authenticated_callers_do_not_share_cached_responses(
|
||||
#[case] authority: &str,
|
||||
#[case] subject: &str,
|
||||
) {
|
||||
use litellm_gateway_auth::{Principal, PrincipalKind};
|
||||
|
||||
let upstream = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"id":"response-1", "model":"test-model", "status":"completed", "output":[]
|
||||
})))
|
||||
.expect(2)
|
||||
.mount(&upstream)
|
||||
.await;
|
||||
let cache: Arc<dyn ResponseCacheService> = Arc::new(ResponseCache::new(Arc::new(
|
||||
InMemoryCache::new(Some(100), Some(Duration::from_secs(60))),
|
||||
)));
|
||||
let first_caller = support::app_with_cache_for_principal(
|
||||
"openai/test-model",
|
||||
&upstream.uri(),
|
||||
cache.clone(),
|
||||
Principal::new("issuer".into(), "tenant-a".into(), PrincipalKind::Service),
|
||||
);
|
||||
let second_caller = support::app_with_cache_for_principal(
|
||||
"openai/test-model",
|
||||
&upstream.uri(),
|
||||
cache,
|
||||
Principal::new(authority.into(), subject.into(), PrincipalKind::Service),
|
||||
);
|
||||
let body = json!({"model":"public/model","input":"same prompt"});
|
||||
let first = support::post(first_caller.clone(), "/v1/responses", body.clone()).await;
|
||||
assert_eq!(first.status(), 200);
|
||||
assert!(!first.headers().contains_key("x-litellm-cache-key"));
|
||||
let first_hit = support::post(first_caller.clone(), "/v1/responses", body.clone()).await;
|
||||
assert_eq!(first_hit.status(), 200);
|
||||
let first_key = first_hit.headers().get("x-litellm-cache-key").unwrap();
|
||||
let second = support::post(second_caller.clone(), "/v1/responses", body.clone()).await;
|
||||
assert_eq!(second.status(), 200);
|
||||
assert!(!second.headers().contains_key("x-litellm-cache-key"));
|
||||
let second_hit = support::post(second_caller, "/v1/responses", body.clone()).await;
|
||||
assert_eq!(second_hit.status(), 200);
|
||||
assert_ne!(
|
||||
second_hit.headers().get("x-litellm-cache-key").unwrap(),
|
||||
first_key
|
||||
);
|
||||
let first_again = support::post(first_caller, "/v1/responses", body).await;
|
||||
assert_eq!(first_again.status(), 200);
|
||||
assert_eq!(
|
||||
first_again.headers().get("x-litellm-cache-key"),
|
||||
Some(first_key)
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -174,25 +174,6 @@ async fn hosted_provider_failure_preserves_status_body_and_request_id() {
|
|||
assert_eq!(body["request_id"], "host-http-request");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::messages("/v1/messages")]
|
||||
#[case::chat("/v1/chat/completions")]
|
||||
#[case::model_path("/engines/path-model/chat/completions")]
|
||||
#[case::ocr("/ocr")]
|
||||
#[case::transcription("/audio/transcriptions")]
|
||||
#[tokio::test]
|
||||
async fn model_permissions_prevent_provider_calls(#[case] path: &str) {
|
||||
let upstream = MockServer::start().await;
|
||||
let app = support::app_with_permissions(
|
||||
"anthropic/test-model",
|
||||
&upstream.uri(),
|
||||
litellm_gateway_auth::Permissions::None,
|
||||
);
|
||||
let response = support::post(app, path, json!({"model": "public/model"})).await;
|
||||
assert_eq!(response.status(), 403);
|
||||
assert!(upstream.received_requests().await.unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::completion(false)]
|
||||
#[case::streaming(true)]
|
||||
|
|
|
|||
|
|
@ -97,22 +97,3 @@ async fn responses_preserve_upstream_errors_without_retry() {
|
|||
.contains("slow down")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn responses_authorize_the_public_model_before_provider_execution() {
|
||||
let upstream = MockServer::start().await;
|
||||
let app = support::app_with_permissions(
|
||||
"openai/test-model",
|
||||
&upstream.uri(),
|
||||
litellm_gateway_auth::Permissions::None,
|
||||
);
|
||||
let response = support::post(
|
||||
app,
|
||||
"/v1/responses",
|
||||
json!({"model": "public/model", "input": "hello"}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), 403);
|
||||
assert!(upstream.received_requests().await.unwrap().is_empty());
|
||||
}
|
||||
|
|
|
|||
|
|
@ -29,15 +29,7 @@ impl SecretSource for NoSecrets {
|
|||
}
|
||||
|
||||
pub fn app(model: &str, api_base: &str) -> Router {
|
||||
app_with_permissions(model, api_base, litellm_gateway_auth::Permissions::All)
|
||||
}
|
||||
|
||||
pub fn app_with_permissions(
|
||||
model: &str,
|
||||
api_base: &str,
|
||||
permissions: litellm_gateway_auth::Permissions,
|
||||
) -> Router {
|
||||
configured_app(model, api_base, permissions, None, None)
|
||||
configured_app(model, api_base, None)
|
||||
}
|
||||
|
||||
pub fn app_with_cache(
|
||||
|
|
@ -45,36 +37,13 @@ pub fn app_with_cache(
|
|||
api_base: &str,
|
||||
cache: Arc<dyn litellm_cache_response::ResponseCacheService>,
|
||||
) -> Router {
|
||||
configured_app(
|
||||
model,
|
||||
api_base,
|
||||
litellm_gateway_auth::Permissions::All,
|
||||
Some(cache),
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn app_with_cache_for_principal(
|
||||
model: &str,
|
||||
api_base: &str,
|
||||
cache: Arc<dyn litellm_cache_response::ResponseCacheService>,
|
||||
principal: litellm_gateway_auth::Principal,
|
||||
) -> Router {
|
||||
configured_app(
|
||||
model,
|
||||
api_base,
|
||||
litellm_gateway_auth::Permissions::All,
|
||||
Some(cache),
|
||||
Some(principal),
|
||||
)
|
||||
configured_app(model, api_base, Some(cache))
|
||||
}
|
||||
|
||||
fn configured_app(
|
||||
model: &str,
|
||||
api_base: &str,
|
||||
permissions: litellm_gateway_auth::Permissions,
|
||||
cache: Option<Arc<dyn litellm_cache_response::ResponseCacheService>>,
|
||||
principal: Option<litellm_gateway_auth::Principal>,
|
||||
) -> Router {
|
||||
let pool = Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver)));
|
||||
let http = Resolution::from(&HttpSettings::default()).config;
|
||||
|
|
@ -102,34 +71,16 @@ fn configured_app(
|
|||
Some(cache) => gateway.with_cache(cache),
|
||||
None => gateway,
|
||||
};
|
||||
router(Arc::new(gateway)).layer(axum::middleware::from_fn_with_state(
|
||||
(permissions, principal),
|
||||
test_identity,
|
||||
))
|
||||
router(Arc::new(gateway)).layer(axum::middleware::from_fn(master_caller))
|
||||
}
|
||||
|
||||
async fn test_identity(
|
||||
axum::extract::State((permissions, principal)): axum::extract::State<(
|
||||
litellm_gateway_auth::Permissions,
|
||||
Option<litellm_gateway_auth::Principal>,
|
||||
)>,
|
||||
async fn master_caller(
|
||||
mut request: axum::extract::Request,
|
||||
next: axum::middleware::Next,
|
||||
) -> Response {
|
||||
let auth = litellm_gateway_auth::Auth::new(
|
||||
Arc::new(litellm_gateway_auth::MasterKeyAuthenticator::new(
|
||||
Some(SecretValue::new("test-inbound-key")),
|
||||
Arc::new(NoSecrets),
|
||||
)),
|
||||
Arc::new(TestPermissions(permissions, principal)),
|
||||
Arc::new(litellm_gateway_auth::NoAdditionalPolicy),
|
||||
Arc::new(litellm_gateway_auth::SystemClock),
|
||||
);
|
||||
let identity = auth
|
||||
.authenticate(&SecretValue::new("test-inbound-key"))
|
||||
.await
|
||||
.unwrap();
|
||||
request.extensions_mut().insert(identity);
|
||||
request
|
||||
.extensions_mut()
|
||||
.insert(litellm_gateway_auth::Caller::Master);
|
||||
next.run(request).await
|
||||
}
|
||||
|
||||
|
|
@ -147,22 +98,3 @@ pub async fn post(app: Router, path: &str, body: Value) -> Response {
|
|||
pub async fn json(response: Response) -> Value {
|
||||
serde_json::from_slice(&to_bytes(response.into_body(), 1024 * 1024).await.unwrap()).unwrap()
|
||||
}
|
||||
|
||||
struct TestPermissions(
|
||||
litellm_gateway_auth::Permissions,
|
||||
Option<litellm_gateway_auth::Principal>,
|
||||
);
|
||||
|
||||
impl litellm_gateway_auth::IdentityResolver for TestPermissions {
|
||||
fn resolve<'a>(
|
||||
&'a self,
|
||||
identity: &'a litellm_gateway_auth::VerifiedIdentity,
|
||||
) -> litellm_gateway_auth::AuthFuture<'a, litellm_gateway_auth::ResolvedIdentity> {
|
||||
Box::pin(async move {
|
||||
Ok(litellm_gateway_auth::ResolvedIdentity {
|
||||
principal: self.1.clone().unwrap_or_else(|| identity.principal.clone()),
|
||||
permissions: self.0.clone(),
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -10,9 +10,7 @@ use tower_cookies::{Cookie, Cookies};
|
|||
use tower_sessions::{Expiry, Session, cookie::SameSite};
|
||||
|
||||
use crate::Error;
|
||||
use litellm_gateway_auth::{
|
||||
AccessRequest, UI_CSRF_KEY, UiAction, UiAuthError, UiAuthSession, UiCredentials, UiSession,
|
||||
};
|
||||
use litellm_gateway_auth::{UI_CSRF_KEY, UiAuthError, UiAuthSession, UiCredentials, UiSession};
|
||||
|
||||
pub struct State {
|
||||
pub signing_key: EncodingKey,
|
||||
|
|
@ -109,11 +107,6 @@ pub async fn login(
|
|||
|
||||
pub async fn info(session: Result<UiSession, UiAuthError>) -> Result<impl IntoResponse, Error> {
|
||||
let session = session?;
|
||||
session
|
||||
.identity
|
||||
.authorize(AccessRequest::Ui(UiAction::SessionInfo))
|
||||
.await
|
||||
.map_err(|_| UiAuthError::Unauthorized)?;
|
||||
Ok(Json(SessionInfo {
|
||||
user_id: session.user.username,
|
||||
user_role: "proxy_admin",
|
||||
|
|
@ -125,12 +118,7 @@ pub async fn logout(
|
|||
mut auth: UiAuthSession,
|
||||
cookies: Cookies,
|
||||
) -> Result<impl IntoResponse, Error> {
|
||||
let validated = validated?;
|
||||
validated
|
||||
.identity
|
||||
.authorize(AccessRequest::Ui(UiAction::Logout))
|
||||
.await
|
||||
.map_err(|_| UiAuthError::Unauthorized)?;
|
||||
validated?;
|
||||
auth.logout().await?;
|
||||
cookies.add(
|
||||
Cookie::build(("token", ""))
|
||||
|
|
|
|||
|
|
@ -1,27 +1,27 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use axum::{extract::Request, middleware::Next, response::Response};
|
||||
use litellm_gateway_auth::{AccessRequest, AuthenticatedRequest, Error, McpAction};
|
||||
use litellm_gateway_auth::{AccessRequest, Caller, Error, McpAction};
|
||||
use litellm_gateway_mcp::{
|
||||
Authorization, GatewayFuture, OperationAuthorizer, ServerInfo, SessionOwner,
|
||||
rmcp::model::ClientRequest,
|
||||
};
|
||||
|
||||
pub(crate) async fn bind_session_owner(
|
||||
identity: AuthenticatedRequest,
|
||||
caller: Caller,
|
||||
mut request: Request,
|
||||
next: Next,
|
||||
) -> Response {
|
||||
request
|
||||
.extensions_mut()
|
||||
.insert(SessionOwner(identity.caller().session_owner()));
|
||||
.insert(SessionOwner(caller.owner()));
|
||||
request
|
||||
.extensions_mut()
|
||||
.insert(Authorization(Arc::new(McpPolicy(identity))));
|
||||
.insert(Authorization(Arc::new(McpPolicy(caller))));
|
||||
next.run(request).await
|
||||
}
|
||||
|
||||
struct McpPolicy(AuthenticatedRequest);
|
||||
struct McpPolicy(Caller);
|
||||
|
||||
impl OperationAuthorizer for McpPolicy {
|
||||
fn authorize<'a>(
|
||||
|
|
@ -49,15 +49,11 @@ impl OperationAuthorizer for McpPolicy {
|
|||
}
|
||||
Some(_) => return Err(litellm_gateway_mcp::Error::Forbidden),
|
||||
};
|
||||
let access = AccessRequest::Mcp {
|
||||
server: server.server_id.clone(),
|
||||
action,
|
||||
};
|
||||
self.0
|
||||
.authorize(access.clone())
|
||||
.await
|
||||
.map_err(mcp_error)?
|
||||
.consume(self.0.caller(), &access)
|
||||
.authorize(&AccessRequest::Mcp {
|
||||
server: server.server_id.clone(),
|
||||
action,
|
||||
})
|
||||
.map_err(mcp_error)
|
||||
})
|
||||
}
|
||||
|
|
@ -65,9 +61,7 @@ impl OperationAuthorizer for McpPolicy {
|
|||
|
||||
fn mcp_error(error: Error) -> litellm_gateway_mcp::Error {
|
||||
match error {
|
||||
Error::InvalidToken | Error::Expired | Error::Forbidden => {
|
||||
litellm_gateway_mcp::Error::Forbidden
|
||||
}
|
||||
Error::InvalidToken | Error::Forbidden => litellm_gateway_mcp::Error::Forbidden,
|
||||
_ => litellm_gateway_mcp::Error::Configuration("MCP authorization unavailable".into()),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,6 +4,8 @@ pub enum Error {
|
|||
Http(#[from] litellm_http::Error),
|
||||
#[error("environment_variables values must be strings, numbers or booleans")]
|
||||
Environment,
|
||||
#[error("general_settings.custom_auth is only supported by the Python proxy")]
|
||||
CustomAuth,
|
||||
#[error("MCP host does not yet support configured setting {0}")]
|
||||
McpSetting(String),
|
||||
#[error(transparent)]
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ use axum::{
|
|||
use http_body_util::BodyExt;
|
||||
|
||||
use litellm_config::Config;
|
||||
use litellm_gateway_auth::Auth;
|
||||
use litellm_gateway_auth::{Auth, Bearer};
|
||||
use litellm_gateway_inference::{Gateway, ModelRouter};
|
||||
use litellm_http::{
|
||||
ClientVariant, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, SslVerify,
|
||||
|
|
@ -29,6 +29,14 @@ use litellm_tracing::ByteChunk;
|
|||
use uuid::Uuid;
|
||||
|
||||
pub fn build_inference(config: &Config) -> Result<Arc<Gateway>, Error> {
|
||||
if config
|
||||
.general_settings
|
||||
.additional_fields
|
||||
.get("custom_auth")
|
||||
.is_some_and(|value| !value.is_null())
|
||||
{
|
||||
return Err(Error::CustomAuth);
|
||||
}
|
||||
let pool = Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver)));
|
||||
let environment = secrets::environment_values(&config.environment_variables)?;
|
||||
let lookup = |name: &str| {
|
||||
|
|
@ -138,7 +146,7 @@ pub fn router(
|
|||
.route_layer(axum::middleware::from_fn(auth::bind_session_owner))
|
||||
.route_layer(axum::middleware::from_fn_with_state(
|
||||
auth,
|
||||
litellm_gateway_auth::authenticate,
|
||||
litellm_gateway_auth::authenticate::<Bearer>,
|
||||
))
|
||||
.layer(axum::middleware::from_fn(log_request));
|
||||
let health = Router::new().route("/health/readiness", axum::routing::get(readiness));
|
||||
|
|
|
|||
|
|
@ -33,6 +33,22 @@ fn inference() -> Arc<Gateway> {
|
|||
litellm_gateway::build_inference(&Config::from_yaml("model_list: []").unwrap()).unwrap()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::configured("custom_auth: auth.user_api_key_auth", false)]
|
||||
#[case::null("custom_auth: null", true)]
|
||||
#[case::absent("master_key: key", true)]
|
||||
fn custom_auth_is_left_to_the_python_proxy(#[case] setting: &str, #[case] builds: bool) {
|
||||
let config =
|
||||
Config::from_yaml(&format!("model_list: []\ngeneral_settings:\n {setting}\n")).unwrap();
|
||||
match litellm_gateway::build_inference(&config) {
|
||||
Ok(_) => assert!(builds),
|
||||
Err(error) => {
|
||||
assert!(!builds);
|
||||
assert!(matches!(error, litellm_gateway::Error::CustomAuth));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::authorized("/v1/messages", Some("Bearer gateway-key"), Some("gateway-key"), 400)]
|
||||
#[case::missing_token("/v1/messages", None, Some("gateway-key"), 401)]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue