From 39beee6529f67dc10b4c5baa5c82ab167995fe9b Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 9 Sep 2026 14:15:45 -0700 Subject: [PATCH] refactor(core): type credential resolution lifecycle --- .../core/src/providers/auth/azure/mod.rs | 2 +- .../core/src/providers/auth/azure/resolve.rs | 250 +++++++++++++----- .../core/src/providers/auth/azure/types.rs | 116 +------- .../core/src/providers/auth/credential.rs | 177 +++++++++++++ .../crates/core/src/providers/auth/mod.rs | 7 + .../crates/core/src/providers/auth/policy.rs | 115 ++++++++ .../core/src/providers/azure_ai/auth.rs | 162 ++++++++---- 7 files changed, 587 insertions(+), 242 deletions(-) create mode 100644 litellm-rust/crates/core/src/providers/auth/credential.rs create mode 100644 litellm-rust/crates/core/src/providers/auth/policy.rs diff --git a/litellm-rust/crates/core/src/providers/auth/azure/mod.rs b/litellm-rust/crates/core/src/providers/auth/azure/mod.rs index 4f8bc654829..5e6fbc6f965 100644 --- a/litellm-rust/crates/core/src/providers/auth/azure/mod.rs +++ b/litellm-rust/crates/core/src/providers/auth/azure/mod.rs @@ -4,4 +4,4 @@ mod resolve; mod types; pub(crate) use resolve::AzureAuthService; -pub use types::{AzureAuthInputs, AzureCredentialType, ConfigValue}; +pub use types::{AzureAuthInputs, AzureCredentialType, ConfigValue, DEFAULT_AZURE_SCOPE}; diff --git a/litellm-rust/crates/core/src/providers/auth/azure/resolve.rs b/litellm-rust/crates/core/src/providers/auth/azure/resolve.rs index 718101a90e8..96c7c71efeb 100644 --- a/litellm-rust/crates/core/src/providers/auth/azure/resolve.rs +++ b/litellm-rust/crates/core/src/providers/auth/azure/resolve.rs @@ -1,14 +1,11 @@ -use std::sync::Arc; - use crate::error::AuthError; -use crate::providers::auth::secret::SecretValue; -use crate::providers::auth::token::{ResolvedCredential, TokenCallerHandle}; +use crate::providers::auth::{ + CredentialFileRef, CredentialLookup, CredentialRef, ResolvedCredential, SecretValue, + TokenCallerHandle, +}; use super::native::{NativeAzureRequest, NativeAzureTokenAcquirer}; -use super::types::{ - AzureAuthInputs, AzureCredentialType, AzureValueLookupHandle, ConfigValue, DEFAULT_AZURE_SCOPE, - ProcessAzureValueLookup, -}; +use super::types::{AzureAuthInputs, AzureCredentialType, ConfigValue, DEFAULT_AZURE_SCOPE}; const AZURE_AD_TOKEN_ENV: &str = "AZURE_AD_TOKEN"; const AZURE_TENANT_ID_ENV: &str = "AZURE_TENANT_ID"; @@ -20,11 +17,11 @@ const AZURE_CREDENTIAL_ENV: &str = "AZURE_CREDENTIAL"; const AZURE_FEDERATED_TOKEN_FILE_ENV: &str = "AZURE_FEDERATED_TOKEN_FILE"; #[derive(Clone, Debug)] -pub(crate) enum AzureAuthPlan { +pub(crate) enum AzureCredentialPlan { Supplied(ResolvedCredential), Caller(TokenCallerHandle), Oidc { - reference: String, + reference: CredentialRef, tenant_id: String, client_id: String, scope: String, @@ -35,18 +32,9 @@ pub(crate) enum AzureAuthPlan { Missing, } +#[derive(Default)] pub(crate) struct AzureAuthService { native: NativeAzureTokenAcquirer, - lookup: AzureValueLookupHandle, -} - -impl Default for AzureAuthService { - fn default() -> Self { - Self { - native: NativeAzureTokenAcquirer::default(), - lookup: AzureValueLookupHandle::new(Arc::new(ProcessAzureValueLookup)), - } - } } impl AzureAuthService { @@ -56,8 +44,8 @@ impl AzureAuthService { env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result, AuthError> { match select_auth_plan(inputs, env_lookup)? { - AzureAuthPlan::Supplied(credential) => Ok(Some(credential)), - AzureAuthPlan::Caller(caller) => { + AzureCredentialPlan::Supplied(credential) => Ok(Some(credential)), + AzureCredentialPlan::Caller(caller) => { let credential = caller.acquire().await?; if credential.secret().expose().is_empty() { return Err(AuthError::Caller( @@ -66,32 +54,34 @@ impl AzureAuthService { } Ok(Some(credential)) } - AzureAuthPlan::Oidc { + AzureCredentialPlan::Oidc { reference, tenant_id, client_id, scope, authority, } => { - let assertion = self.lookup.resolve(&reference).await?.ok_or_else(|| { - AuthError::Acquisition( - "Azure OIDC reference did not resolve to a value".to_string(), - ) - })?; + let assertion = resolve_reference(inputs, env_lookup, &reference) + .await? + .ok_or_else(|| { + AuthError::Acquisition( + "Azure OIDC reference did not resolve to a value".to_string(), + ) + })?; self.native .acquire(NativeAzureRequest::ClientAssertion { tenant_id, client_id, assertion, - assertion_identity: reference, + assertion_identity: format!("{reference:?}"), scope, authority, }) .await .map(Some) } - AzureAuthPlan::Native(request) => self.native.acquire(request).await.map(Some), - AzureAuthPlan::Chain(requests) => { + AzureCredentialPlan::Native(request) => self.native.acquire(request).await.map(Some), + AzureCredentialPlan::Chain(requests) => { let mut failures = Vec::new(); for request in requests { match self.native.acquire(request).await { @@ -101,7 +91,7 @@ impl AzureAuthService { } Err(AuthError::Acquisition(failures.join("; "))) } - AzureAuthPlan::Missing => Ok(None), + AzureCredentialPlan::Missing => Ok(None), } } } @@ -109,7 +99,7 @@ impl AzureAuthService { pub(crate) fn select_auth_plan( inputs: &AzureAuthInputs, env_lookup: &dyn Fn(&str) -> Option, -) -> Result { +) -> Result { let token = configured_secret(&inputs.azure_ad_token, AZURE_AD_TOKEN_ENV, env_lookup); let tenant_id = configured_string(&inputs.tenant_id, AZURE_TENANT_ID_ENV, env_lookup); let client_id = configured_string(&inputs.client_id, AZURE_CLIENT_ID_ENV, env_lookup); @@ -138,19 +128,23 @@ pub(crate) fn select_auth_plan( && let (Some(tenant_id), Some(client_id), Some(client_secret)) = (tenant_id.clone(), client_id.clone(), client_secret) { - return Ok(AzureAuthPlan::Native(NativeAzureRequest::ClientSecret { - tenant_id, - client_id, - client_secret, - scope, - authority, - })); + return Ok(AzureCredentialPlan::Native( + NativeAzureRequest::ClientSecret { + tenant_id, + client_id, + client_secret, + scope, + authority, + }, + )); } - if let (Some(reference), Some(tenant_id), Some(client_id)) = - (oidc_reference(&token), tenant_id.clone(), client_id.clone()) - { - return Ok(AzureAuthPlan::Oidc { + if let (Some(reference), Some(tenant_id), Some(client_id)) = ( + oidc_reference(&token)?, + tenant_id.clone(), + client_id.clone(), + ) { + return Ok(AzureCredentialPlan::Oidc { reference, tenant_id, client_id, @@ -160,18 +154,20 @@ pub(crate) fn select_auth_plan( } if let Some(caller) = &inputs.azure_ad_token_provider { - return Ok(AzureAuthPlan::Caller(caller.clone())); + return Ok(AzureCredentialPlan::Caller(caller.clone())); } if let Some(token) = token { - return Ok(AzureAuthPlan::Supplied(ResolvedCredential::AccessToken { - token, - expires_on: None, - })); + return Ok(AzureCredentialPlan::Supplied( + ResolvedCredential::AccessToken { + token, + expires_on: None, + }, + )); } if !inputs.enable_azure_ad_token_refresh && selector.is_none() { - return Ok(AzureAuthPlan::Missing); + return Ok(AzureCredentialPlan::Missing); } select_native_plan( @@ -191,7 +187,7 @@ fn select_native_plan( federated_token_file: Option, scope: String, authority: Option, -) -> Result { +) -> Result { let selected = selector.unwrap_or_else(|| { if federated_token_file.is_some() { AzureCredentialType::DefaultAzureCredential @@ -206,15 +202,12 @@ fn select_native_plan( AzureCredentialType::ClientSecretCredential => Err(AuthError::InvalidConfiguration( "ClientSecretCredential requires tenant_id, client_id, and client_secret".to_string(), )), - AzureCredentialType::WorkloadIdentityCredential => Ok(AzureAuthPlan::Native( + AzureCredentialType::WorkloadIdentityCredential => Ok(AzureCredentialPlan::Native( workload_request(tenant_id, client_id, federated_token_file, scope, authority)?, )), - AzureCredentialType::ManagedIdentityCredential => { - Ok(AzureAuthPlan::Native(NativeAzureRequest::ManagedIdentity { - client_id, - scope, - })) - } + AzureCredentialType::ManagedIdentityCredential => Ok(AzureCredentialPlan::Native( + NativeAzureRequest::ManagedIdentity { client_id, scope }, + )), AzureCredentialType::DefaultAzureCredential => { let workload = match (tenant_id, client_id.clone(), federated_token_file) { (Some(tenant_id), Some(client_id), Some(token_file_path)) => { @@ -228,7 +221,7 @@ fn select_native_plan( } _ => None, }; - Ok(AzureAuthPlan::Chain( + Ok(AzureCredentialPlan::Chain( workload .into_iter() .chain(std::iter::once(NativeAzureRequest::ManagedIdentity { @@ -258,7 +251,7 @@ fn select_native_plan( client_id: Some(client_id), scope: scope.clone(), }); - Ok(AzureAuthPlan::Chain( + Ok(AzureCredentialPlan::Chain( workload .into_iter() .chain(user_assigned) @@ -328,20 +321,100 @@ fn configured_secret( }) } -fn oidc_reference(token: &Option) -> Option { - token - .as_ref() - .map(SecretValue::expose) - .filter(|value| value.starts_with("oidc/")) - .map(str::to_string) +async fn resolve_reference( + inputs: &AzureAuthInputs, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + reference: &CredentialRef, +) -> Result, AuthError> { + let lookup = match reference { + CredentialRef::Explicit(secret) => return Ok(Some(secret.clone())), + CredentialRef::Env(name) => env_lookup(name) + .filter(|value| !value.is_empty()) + .map(SecretValue::new) + .map_or(CredentialLookup::Missing, CredentialLookup::Found), + CredentialRef::None => return Ok(None), + CredentialRef::File(_) | CredentialRef::Request(_) | CredentialRef::Host(_) => { + let resolver = inputs.credential_resolver.as_ref().ok_or_else(|| { + AuthError::InvalidConfiguration( + "credential reference requires a host credential resolver".to_string(), + ) + })?; + resolver.resolve(reference).await? + } + }; + Ok(match lookup { + CredentialLookup::Found(secret) => Some(secret), + CredentialLookup::Missing | CredentialLookup::Declined => None, + }) +} + +fn oidc_reference(token: &Option) -> Result, AuthError> { + let Some(value) = token.as_ref().map(SecretValue::expose) else { + return Ok(None); + }; + if let Some(name) = value.strip_prefix("oidc/env/") { + return non_empty_reference(name, "OIDC environment reference") + .map(CredentialRef::Env) + .map(Some); + } + if let Some(name) = value.strip_prefix("oidc/env_path/") { + return non_empty_reference(name, "OIDC environment path reference") + .map(|name| CredentialRef::File(CredentialFileRef::EnvironmentVariable(name))) + .map(Some); + } + if let Some(path) = value.strip_prefix("oidc/file/") { + let path = non_empty_reference(path, "OIDC file reference")?; + return Ok(Some(CredentialRef::File(CredentialFileRef::Path( + path.into(), + )))); + } + if value.starts_with("oidc/") { + return Err(AuthError::InvalidConfiguration(format!( + "unsupported OIDC reference: {value}" + ))); + } + Ok(None) +} + +fn non_empty_reference(value: &str, kind: &str) -> Result { + if value.is_empty() { + return Err(AuthError::InvalidConfiguration(format!( + "{kind} cannot be empty" + ))); + } + Ok(value.to_string()) } #[cfg(test)] mod tests { + use std::sync::Arc; + use serde_json::json; - use super::{AzureAuthPlan, select_auth_plan}; + use super::{AzureCredentialPlan, oidc_reference, resolve_reference, select_auth_plan}; use crate::providers::auth::azure::AzureAuthInputs; + use crate::providers::auth::{ + CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialRef, + CredentialResolver, CredentialResolverHandle, SecretValue, + }; + + #[derive(Debug)] + struct FileResolver; + + impl CredentialResolver for FileResolver { + fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a> { + Box::pin(async move { + Ok(match reference { + CredentialRef::File(CredentialFileRef::Path(path)) + if path == std::path::Path::new("/run/secrets/assertion") => + { + CredentialLookup::Found(SecretValue::new("rotated-assertion")) + } + _ => CredentialLookup::Declined, + }) + }) + } + } #[test] fn null_and_empty_values_fall_back_to_environment() { @@ -355,7 +428,7 @@ mod tests { }) .unwrap(); - assert!(matches!(plan, AzureAuthPlan::Native(_))); + assert!(matches!(plan, AzureCredentialPlan::Native(_))); } #[test] @@ -365,7 +438,7 @@ mod tests { assert!(matches!( select_auth_plan(&inputs, &|_| None).unwrap(), - AzureAuthPlan::Supplied(_) + AzureCredentialPlan::Supplied(_) )); } @@ -380,7 +453,44 @@ mod tests { assert!(matches!( select_auth_plan(&inputs, &|_| None).unwrap(), - AzureAuthPlan::Oidc { .. } + AzureCredentialPlan::Oidc { + reference: CredentialRef::Env(name), + .. + } if name == "ASSERTION" )); } + + #[test] + fn oidc_file_location_is_typed_before_resolution() { + assert_eq!( + oidc_reference(&Some(SecretValue::new("oidc/file//run/secrets/assertion"))).unwrap(), + Some(CredentialRef::File(CredentialFileRef::Path( + "/run/secrets/assertion".into() + ))) + ); + } + + #[test] + fn unsupported_oidc_reference_is_rejected_during_plan_creation() { + let error = oidc_reference(&Some(SecretValue::new("oidc/vault/assertion"))) + .expect_err("unsupported backend must fail validation"); + + assert!(error.to_string().contains("unsupported OIDC reference")); + } + + #[tokio::test] + async fn host_resolver_owns_file_access() { + let inputs = AzureAuthInputs { + credential_resolver: Some(CredentialResolverHandle::new(Arc::new(FileResolver))), + ..AzureAuthInputs::default() + }; + let reference = + CredentialRef::File(CredentialFileRef::Path("/run/secrets/assertion".into())); + + let resolved = resolve_reference(&inputs, &|_| None, &reference) + .await + .unwrap(); + + assert_eq!(resolved, Some(SecretValue::new("rotated-assertion"))); + } } diff --git a/litellm-rust/crates/core/src/providers/auth/azure/types.rs b/litellm-rust/crates/core/src/providers/auth/azure/types.rs index 8079f7425e6..4b52837b8f6 100644 --- a/litellm-rust/crates/core/src/providers/auth/azure/types.rs +++ b/litellm-rust/crates/core/src/providers/auth/azure/types.rs @@ -1,14 +1,8 @@ -use std::future::Future; -use std::pin::Pin; -use std::sync::Arc; - use serde_json::{Map, Value}; use strum::EnumString; -use veil::Redact; use crate::error::AuthError; -use crate::providers::auth::secret::SecretValue; -use crate::providers::auth::token::TokenCallerHandle; +use crate::providers::auth::{CredentialResolverHandle, SecretValue, TokenCallerHandle}; pub const DEFAULT_AZURE_SCOPE: &str = "https://cognitiveservices.azure.com/.default"; @@ -42,6 +36,7 @@ pub enum AzureCredentialType { pub struct AzureAuthInputs { pub azure_ad_token: ConfigValue, pub azure_ad_token_provider: Option, + pub credential_resolver: Option, pub tenant_id: ConfigValue, pub client_id: ConfigValue, pub client_secret: ConfigValue, @@ -57,6 +52,7 @@ impl AzureAuthInputs { Ok(Self { azure_ad_token: secret_config(params, "azure_ad_token")?, azure_ad_token_provider: None, + credential_resolver: None, tenant_id: string_config(params, "tenant_id")?, client_id: string_config(params, "client_id")?, client_secret: secret_config(params, "client_secret")?, @@ -97,112 +93,6 @@ fn secret_config( }) } -pub(crate) type LookupFuture<'a> = - Pin, AuthError>> + Send + 'a>>; - -pub(crate) trait AzureValueLookup: std::fmt::Debug + Send + Sync { - fn resolve<'a>(&'a self, reference: &'a str) -> LookupFuture<'a>; -} - -#[derive(Clone, Redact)] -pub(crate) struct AzureValueLookupHandle(#[redact(with = "[REDACTED]")] Arc); - -impl AzureValueLookupHandle { - pub(crate) fn new(lookup: Arc) -> Self { - Self(lookup) - } - - pub(crate) async fn resolve(&self, reference: &str) -> Result, AuthError> { - self.0.resolve(reference).await - } -} - -#[derive(Clone, Copy, Debug, Default)] -pub(crate) struct ProcessAzureValueLookup; - -impl AzureValueLookup for ProcessAzureValueLookup { - fn resolve<'a>(&'a self, reference: &'a str) -> LookupFuture<'a> { - Box::pin(async move { resolve_process_value(reference) }) - } -} - -fn resolve_process_value(reference: &str) -> Result, AuthError> { - if let Some(name) = reference.strip_prefix("oidc/env/") { - return Ok(std::env::var(name).ok().map(SecretValue::new)); - } - if let Some(name) = reference.strip_prefix("oidc/env_path/") { - let Some(path) = std::env::var(name).ok() else { - return Ok(None); - }; - return read_secret_file(&path, false).map(Some); - } - if let Some(path) = reference.strip_prefix("oidc/file/") { - return read_secret_file(path, true).map(Some); - } - if reference.starts_with("oidc/") { - return Err(AuthError::InvalidConfiguration(format!( - "unsupported OIDC reference: {reference}" - ))); - } - Ok(None) -} - -fn read_secret_file(path: &str, enforce_allowlist: bool) -> Result { - let requested = std::path::Path::new(path); - if !requested.is_absolute() { - return Err(AuthError::InvalidConfiguration( - "oidc/file path must be absolute".to_string(), - )); - } - let resolved = requested.canonicalize().map_err(|error| { - AuthError::Acquisition(format!("failed to resolve Azure assertion file: {error}")) - })?; - if enforce_allowlist { - let allowed = oidc_allowed_directories()?; - if !allowed - .iter() - .any(|directory| resolved.starts_with(directory)) - { - return Err(AuthError::InvalidConfiguration( - "oidc/file path is outside the allowed credential directories".to_string(), - )); - } - } - std::fs::read_to_string(resolved) - .map(SecretValue::new) - .map_err(|error| { - AuthError::Acquisition(format!("failed to read Azure assertion file: {error}")) - }) -} - -fn oidc_allowed_directories() -> Result, AuthError> { - let configured = std::env::var("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS").ok(); - let raw = configured - .as_deref() - .map(|value| { - value - .split(',') - .filter(|part| !part.trim().is_empty()) - .collect::>() - }) - .unwrap_or_else(|| vec!["/var/run/secrets", "/run/secrets"]); - raw.into_iter() - .map(|path| { - let candidate = std::path::Path::new(path.trim()); - if !candidate.is_absolute() { - return Err(AuthError::InvalidConfiguration( - "OIDC allowed credential directories must be absolute".to_string(), - )); - } - candidate.canonicalize().map_err(|error| { - AuthError::InvalidConfiguration(format!( - "failed to resolve OIDC allowed credential directory: {error}" - )) - }) - }) - .collect() -} - #[cfg(test)] mod tests { use serde_json::json; diff --git a/litellm-rust/crates/core/src/providers/auth/credential.rs b/litellm-rust/crates/core/src/providers/auth/credential.rs new file mode 100644 index 00000000000..2f7ac710bda --- /dev/null +++ b/litellm-rust/crates/core/src/providers/auth/credential.rs @@ -0,0 +1,177 @@ +use std::future::Future; +use std::path::PathBuf; +use std::pin::Pin; +use std::sync::Arc; + +use veil::Redact; + +use crate::error::AuthError; + +use super::{ResolvedCredential, SecretValue, TokenCallerHandle}; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum CredentialFileRef { + Path(PathBuf), + EnvironmentVariable(String), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum CredentialRef { + Explicit(SecretValue), + Env(String), + File(CredentialFileRef), + Request(String), + Host(String), + None, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum CredentialLookup { + Found(SecretValue), + Missing, + Declined, +} + +pub type CredentialLookupFuture<'a> = + Pin> + Send + 'a>>; + +pub trait CredentialResolver: std::fmt::Debug + Send + Sync { + fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a>; +} + +#[derive(Clone, Redact)] +pub struct CredentialResolverHandle(#[redact(with = "[REDACTED]")] Arc); + +impl CredentialResolverHandle { + pub fn new(resolver: Arc) -> Self { + Self(resolver) + } + + pub async fn resolve(&self, reference: &CredentialRef) -> Result { + self.0.resolve(reference).await + } +} + +#[derive(Clone, Debug)] +pub enum CredentialPlan { + Static(CredentialRef), + Caller(TokenCallerHandle), + None, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum CredentialPlanResolution { + Resolved(ResolvedCredential), + Unavailable, +} + +impl CredentialPlan { + pub async fn resolve( + &self, + resolver: &CredentialResolverHandle, + ) -> Result { + match self { + Self::Static(CredentialRef::Explicit(secret)) => Ok( + CredentialPlanResolution::Resolved(ResolvedCredential::Static(secret.clone())), + ), + Self::Static(CredentialRef::None) | Self::None => { + Ok(CredentialPlanResolution::Unavailable) + } + Self::Static(reference) => match resolver.resolve(reference).await? { + CredentialLookup::Found(secret) => Ok(CredentialPlanResolution::Resolved( + ResolvedCredential::Static(secret), + )), + CredentialLookup::Missing | CredentialLookup::Declined => { + Ok(CredentialPlanResolution::Unavailable) + } + }, + Self::Caller(caller) => { + let credential = caller.acquire().await?; + if credential.secret().expose().is_empty() { + return Err(AuthError::Caller( + "credential caller returned an empty credential".to_string(), + )); + } + Ok(CredentialPlanResolution::Resolved(credential)) + } + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::{ + CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution, + CredentialRef, CredentialResolver, CredentialResolverHandle, + }; + use crate::AuthError; + use crate::providers::auth::SecretValue; + + #[derive(Debug)] + struct HostResolver; + + impl CredentialResolver for HostResolver { + fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a> { + Box::pin(async move { + Ok(match reference { + CredentialRef::Host(name) if name == "rotating-token" => { + CredentialLookup::Found(SecretValue::new("resolved")) + } + _ => CredentialLookup::Declined, + }) + }) + } + } + + #[tokio::test] + async fn static_host_reference_resolves_at_acquisition_time() { + let resolver = CredentialResolverHandle::new(Arc::new(HostResolver)); + let plan = CredentialPlan::Static(CredentialRef::Host("rotating-token".to_string())); + + let resolved = plan.resolve(&resolver).await.unwrap(); + + assert!(matches!(resolved, CredentialPlanResolution::Resolved(_))); + } + + #[tokio::test] + async fn declined_reference_is_available_for_pre_acquisition_fallback() { + let resolver = CredentialResolverHandle::new(Arc::new(HostResolver)); + let plan = CredentialPlan::Static(CredentialRef::Request("api-key".to_string())); + + assert_eq!( + plan.resolve(&resolver).await.unwrap(), + CredentialPlanResolution::Unavailable + ); + } + + #[derive(Debug)] + struct FailingResolver; + + impl CredentialResolver for FailingResolver { + fn resolve<'a>(&'a self, _reference: &'a CredentialRef) -> CredentialLookupFuture<'a> { + Box::pin(async { + Err(AuthError::Acquisition( + "credential backend failed".to_string(), + )) + }) + } + } + + #[tokio::test] + async fn acquisition_failure_is_terminal() { + let resolver = CredentialResolverHandle::new(Arc::new(FailingResolver)); + let plan = CredentialPlan::Static(CredentialRef::Host("token".to_string())); + + let error = plan + .resolve(&resolver) + .await + .expect_err("acquisition errors cannot become fallback"); + + assert_eq!( + error, + AuthError::Acquisition("credential backend failed".to_string()) + ); + } +} diff --git a/litellm-rust/crates/core/src/providers/auth/mod.rs b/litellm-rust/crates/core/src/providers/auth/mod.rs index 65761ff6a9c..a4909979564 100644 --- a/litellm-rust/crates/core/src/providers/auth/mod.rs +++ b/litellm-rust/crates/core/src/providers/auth/mod.rs @@ -1,8 +1,15 @@ pub mod azure; +mod credential; pub(crate) mod http; +mod policy; mod secret; mod token; +pub use credential::{ + CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan, + CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, +}; pub use http::CredentialPlacement; +pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; pub use secret::SecretValue; pub use token::{ResolvedCredential, TokenCaller, TokenCallerHandle, TokenFuture}; diff --git a/litellm-rust/crates/core/src/providers/auth/policy.rs b/litellm-rust/crates/core/src/providers/auth/policy.rs new file mode 100644 index 00000000000..e8451a4f66f --- /dev/null +++ b/litellm-rust/crates/core/src/providers/auth/policy.rs @@ -0,0 +1,115 @@ +use crate::error::AuthError; + +use super::http::apply_credential; +use super::{CredentialPlacement, ResolvedCredential}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CredentialPlanKind { + Static, + Entra, + Caller, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct CredentialRule { + pub kind: CredentialPlanKind, + pub placement: CredentialPlacement, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ExistingHeaderBehavior { + Preserve, + Reject, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ProviderAuthPolicy { + pub rules: &'static [CredentialRule], + pub accepted_existing_headers: &'static [&'static str], + pub existing_header_behavior: ExistingHeaderBehavior, + pub scope: Option<&'static str>, + pub audience: Option<&'static str>, +} + +impl ProviderAuthPolicy { + pub fn has_existing_credential(&self, headers: &[(String, String)]) -> bool { + headers.iter().any(|(name, _)| { + self.accepted_existing_headers + .iter() + .any(|accepted| name.eq_ignore_ascii_case(accepted)) + }) + } + + pub fn apply( + &self, + headers: Vec<(String, String)>, + kind: CredentialPlanKind, + credential: &ResolvedCredential, + ) -> Result, AuthError> { + if self.has_existing_credential(&headers) { + return match self.existing_header_behavior { + ExistingHeaderBehavior::Preserve => Ok(headers), + ExistingHeaderBehavior::Reject => Err(AuthError::InvalidConfiguration( + "credential header already exists".to_string(), + )), + }; + } + let rule = self + .rules + .iter() + .find(|rule| rule.kind == kind) + .ok_or_else(|| { + AuthError::InvalidConfiguration( + "credential plan is not allowed by the provider auth policy".to_string(), + ) + })?; + apply_credential(headers, credential.secret().expose(), rule.placement) + } +} + +#[cfg(test)] +mod tests { + use super::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; + use crate::providers::auth::{CredentialPlacement, ResolvedCredential, SecretValue}; + + const RULES: &[CredentialRule] = &[CredentialRule { + kind: CredentialPlanKind::Static, + placement: CredentialPlacement::Header("x-api-key"), + }]; + const POLICY: ProviderAuthPolicy = ProviderAuthPolicy { + rules: RULES, + accepted_existing_headers: &["x-api-key"], + existing_header_behavior: ExistingHeaderBehavior::Preserve, + scope: None, + audience: None, + }; + + #[test] + fn rules_define_allowed_plans_and_credential_placement() { + let headers = POLICY + .apply( + Vec::new(), + CredentialPlanKind::Static, + &ResolvedCredential::Static(SecretValue::new("secret")), + ) + .unwrap(); + + assert_eq!( + headers, + vec![("x-api-key".to_string(), "secret".to_string())] + ); + } + + #[test] + fn unsupported_plan_is_rejected() { + let error = POLICY + .apply( + Vec::new(), + CredentialPlanKind::Entra, + &ResolvedCredential::Static(SecretValue::new("secret")), + ) + .unwrap_err(); + + assert!(error.to_string().contains("not allowed")); + } +} diff --git a/litellm-rust/crates/core/src/providers/azure_ai/auth.rs b/litellm-rust/crates/core/src/providers/azure_ai/auth.rs index cb8239058f8..1b7f875c5c5 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/auth.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/auth.rs @@ -1,15 +1,53 @@ use std::sync::OnceLock; use crate::error::{AuthError, Error}; -use crate::providers::auth::CredentialPlacement; use crate::providers::auth::azure::{AzureAuthInputs, AzureAuthService}; use crate::providers::auth::http::apply_credential; +use crate::providers::auth::{ + CredentialPlacement, CredentialPlanKind, CredentialRef, CredentialRule, ExistingHeaderBehavior, + ProviderAuthPolicy, ResolvedCredential, SecretValue, +}; pub const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY"; pub const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE"; pub const AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY"; pub const AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT"; +const AZURE_AI_AUTH_RULES: &[CredentialRule] = &[ + CredentialRule { + kind: CredentialPlanKind::Static, + placement: CredentialPlacement::Bearer, + }, + CredentialRule { + kind: CredentialPlanKind::Entra, + placement: CredentialPlacement::Bearer, + }, +]; +const AZURE_AI_AUTH_POLICY: ProviderAuthPolicy = ProviderAuthPolicy { + rules: AZURE_AI_AUTH_RULES, + accepted_existing_headers: &["Authorization"], + existing_header_behavior: ExistingHeaderBehavior::Preserve, + scope: Some(crate::providers::auth::azure::DEFAULT_AZURE_SCOPE), + audience: None, +}; +const DOCUMENT_INTELLIGENCE_AUTH_RULES: &[CredentialRule] = &[ + CredentialRule { + kind: CredentialPlanKind::Static, + placement: CredentialPlacement::Header("Ocp-Apim-Subscription-Key"), + }, + CredentialRule { + kind: CredentialPlanKind::Entra, + placement: CredentialPlacement::Bearer, + }, +]; +const DOCUMENT_INTELLIGENCE_AUTH_POLICY: ProviderAuthPolicy = ProviderAuthPolicy { + rules: DOCUMENT_INTELLIGENCE_AUTH_RULES, + accepted_existing_headers: &["Authorization", "Ocp-Apim-Subscription-Key"], + existing_header_behavior: ExistingHeaderBehavior::Preserve, + scope: Some(crate::providers::auth::azure::DEFAULT_AZURE_SCOPE), + audience: None, +}; + fn resolve_value( explicit: Option<&str>, env_name: &str, @@ -103,27 +141,16 @@ pub async fn authenticate( auth_inputs: Option<&AzureAuthInputs>, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result, Error> { - if crate::http_utils::has_header(&headers, "Authorization") { - return Ok(headers); - } - if let Some(api_key) = resolve_optional_api_key(api_key, env_lookup) { - return apply_credential(headers, &api_key, CredentialPlacement::Bearer) - .map_err(Error::from); - } - let credential = resolve_entra_credential(auth_inputs, env_lookup) - .await? - .ok_or_else(|| { - AuthError::InvalidConfiguration( - "Missing Azure AI credentials - set AZURE_AI_API_KEY or configure Entra ID" - .to_string(), - ) - })?; - apply_credential( + authenticate_with_policy( headers, - credential.secret().expose(), - CredentialPlacement::Bearer, + api_key, + AZURE_AI_API_KEY_ENV, + auth_inputs, + env_lookup, + &AZURE_AI_AUTH_POLICY, + "Missing Azure AI credentials - set AZURE_AI_API_KEY or configure Entra ID", ) - .map_err(Error::from) + .await } pub async fn authenticate_document_intelligence( @@ -132,35 +159,50 @@ pub async fn authenticate_document_intelligence( auth_inputs: Option<&AzureAuthInputs>, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result, Error> { - if crate::http_utils::has_header(&headers, "Authorization") - || crate::http_utils::has_header(&headers, "Ocp-Apim-Subscription-Key") - { + authenticate_with_policy( + headers, + api_key, + AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV, + auth_inputs, + env_lookup, + &DOCUMENT_INTELLIGENCE_AUTH_POLICY, + "Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or configure Entra ID", + ) + .await +} + +async fn authenticate_with_policy( + headers: Vec<(String, String)>, + api_key: Option<&str>, + api_key_env: &str, + auth_inputs: Option<&AzureAuthInputs>, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + policy: &ProviderAuthPolicy, + missing_message: &'static str, +) -> Result, Error> { + if policy.has_existing_credential(&headers) { return Ok(headers); } - if let Some(api_key) = - resolve_optional_value(api_key, AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV, env_lookup) - { - return apply_credential( - headers, - &api_key, - CredentialPlacement::Header("Ocp-Apim-Subscription-Key"), - ) - .map_err(Error::from); + for rule in policy.rules { + let credential = match rule.kind { + CredentialPlanKind::Static => { + resolve_static_credential(api_key, api_key_env, env_lookup) + } + CredentialPlanKind::Entra => resolve_entra_credential(auth_inputs, env_lookup).await?, + CredentialPlanKind::Caller => { + return Err(AuthError::InvalidConfiguration( + "caller credential plan requires provider-specific inputs".to_string(), + ) + .into()); + } + }; + if let Some(credential) = credential { + return policy + .apply(headers, rule.kind, &credential) + .map_err(Error::from); + } } - let credential = resolve_entra_credential(auth_inputs, env_lookup) - .await? - .ok_or_else(|| { - AuthError::InvalidConfiguration( - "Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or configure Entra ID" - .to_string(), - ) - })?; - apply_credential( - headers, - credential.secret().expose(), - CredentialPlacement::Bearer, - ) - .map_err(Error::from) + Err(AuthError::InvalidConfiguration(missing_message.to_string()).into()) } async fn resolve_entra_credential( @@ -173,22 +215,26 @@ async fn resolve_entra_credential( } } -fn resolve_optional_api_key( - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Option { - resolve_optional_value(api_key, AZURE_AI_API_KEY_ENV, env_lookup) -} - -fn resolve_optional_value( +fn resolve_static_credential( explicit: Option<&str>, env_name: &str, env_lookup: &dyn Fn(&str) -> Option, -) -> Option { - explicit +) -> Option { + let reference = explicit .filter(|value| !value.is_empty()) - .map(str::to_string) - .or_else(|| env_lookup(env_name).filter(|value| !value.is_empty())) + .map(|value| CredentialRef::Explicit(SecretValue::new(value))) + .unwrap_or_else(|| CredentialRef::Env(env_name.to_string())); + match reference { + CredentialRef::Explicit(secret) => Some(ResolvedCredential::Static(secret)), + CredentialRef::Env(name) => env_lookup(&name) + .filter(|value| !value.is_empty()) + .map(SecretValue::new) + .map(ResolvedCredential::Static), + CredentialRef::File(_) + | CredentialRef::Request(_) + | CredentialRef::Host(_) + | CredentialRef::None => None, + } } fn auth_service() -> &'static AzureAuthService {