scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

refactor(core): type credential resolution lifecycle

Browse source
This commit is contained in:
Yujong Lee 2026-09-09 14:15:45 -07:00
parent 983813b3a7
commit 39beee6529
7 changed files with 587 additions and 242 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm-rust/crates/core/src/providers/auth/azure/mod.rs
View file

@ -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};

250
litellm-rust/crates/core/src/providers/auth/azure/resolve.rs
View file

@ -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<String> + Sync),
) -> Result<Option<ResolvedCredential>, 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<String>,
) -> Result<AzureAuthPlan, AuthError> {
) -> Result<AzureCredentialPlan, AuthError> {
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<String>,
scope: String,
authority: Option<String>,
) -> Result<AzureAuthPlan, AuthError> {
) -> Result<AzureCredentialPlan, AuthError> {
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<SecretValue>) -> Option<String> {
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<String> + Sync),
reference: &CredentialRef,
) -> Result<Option<SecretValue>, 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<SecretValue>) -> Result<Option<CredentialRef>, 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<String, AuthError> {
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")));
}
}

116
litellm-rust/crates/core/src/providers/auth/azure/types.rs
View file

@ -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<SecretValue>,
pub azure_ad_token_provider: Option<TokenCallerHandle>,
pub credential_resolver: Option<CredentialResolverHandle>,
pub tenant_id: ConfigValue<String>,
pub client_id: ConfigValue<String>,
pub client_secret: ConfigValue<SecretValue>,
@ -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<Box<dyn Future<Output = Result<Option<SecretValue>, 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<dyn AzureValueLookup>);
impl AzureValueLookupHandle {
pub(crate) fn new(lookup: Arc<dyn AzureValueLookup>) -> Self {
Self(lookup)
}
pub(crate) async fn resolve(&self, reference: &str) -> Result<Option<SecretValue>, 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<Option<SecretValue>, 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<SecretValue, AuthError> {
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<Vec<std::path::PathBuf>, 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::<Vec<_>>()
})
.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;

177
litellm-rust/crates/core/src/providers/auth/credential.rs Normal file
View file

@ -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<Box<dyn Future<Output = Result<CredentialLookup, AuthError>> + 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<dyn CredentialResolver>);
impl CredentialResolverHandle {
pub fn new(resolver: Arc<dyn CredentialResolver>) -> Self {
Self(resolver)
}
pub async fn resolve(&self, reference: &CredentialRef) -> Result<CredentialLookup, AuthError> {
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<CredentialPlanResolution, AuthError> {
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())
);
}
}

7
litellm-rust/crates/core/src/providers/auth/mod.rs
View file

@ -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};

115
litellm-rust/crates/core/src/providers/auth/policy.rs Normal file
View file

@ -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<Vec<(String, String)>, 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"));
}
}

162
litellm-rust/crates/core/src/providers/azure_ai/auth.rs
View file

@ -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<String> + Sync),
) -> Result<Vec<(String, String)>, 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<String> + Sync),
) -> Result<Vec<(String, String)>, 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<String> + Sync),
policy: &ProviderAuthPolicy,
missing_message: &'static str,
) -> Result<Vec<(String, String)>, 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<String>,
) -> Option<String> {
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<String>,
) -> Option<String> {
explicit
) -> Option<ResolvedCredential> {
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 {