diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index f91deaaa401..9a2b2e79364 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3389,15 +3389,18 @@ version = "0.1.0" dependencies = [ "aws-config", "aws-credential-types", + "aws-runtime", "aws-sdk-sts", "aws-sigv4", "aws-smithy-runtime-api", + "aws-smithy-types", "aws-types", "bytes", "futures-util", "litellm-auth-types", "litellm-error", "litellm-http", + "litellm-secret-value", "moka", "reqwest", "rstest", @@ -3406,6 +3409,7 @@ dependencies = [ "sha2 0.10.9", "thiserror 2.0.19", "tokio", + "url", "veil", ] @@ -3431,6 +3435,7 @@ dependencies = [ name = "litellm-auth-gcp" version = "0.1.0" dependencies = [ + "axum", "base64 0.22.1", "gcp_auth", "google-cloud-auth", @@ -3440,10 +3445,12 @@ dependencies = [ "litellm-secret-value", "moka", "rstest", + "rustls 0.23.42", "serde", "serde_json", "sha2 0.10.9", "tempfile", + "time", "tokio", "veil", ] diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml index e4578777a08..9f03cf5aa94 100644 --- a/litellm-rust/crates/auth-aws/Cargo.toml +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -11,14 +11,16 @@ litellm-error.workspace = true litellm-http.workspace = true reqwest.workspace = true -moka = { workspace = true, features = ["sync"] } +moka.workspace = true +litellm-secret-value.workspace = true serde.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } sha2.workspace = true thiserror.workspace = true +url.workspace = true veil.workspace = true -aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"] } +aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio", "credentials-process"] } aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"] } aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"] } aws-sigv4 = "1.5.1" @@ -26,6 +28,8 @@ aws-types = "1.4.0" aws-smithy-runtime-api = "1.13.0" [dev-dependencies] +aws-smithy-types = "1.6.4" +aws-runtime = { version = "1.9.4", default-features = false } rstest.workspace = true litellm-http = { workspace = true, features = ["test-support"] } bytes.workspace = true diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index 81aa4b068bc..113d6d3fbab 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -1,13 +1,16 @@ use std::time::Duration; use std::time::{SystemTime, UNIX_EPOCH}; -use moka::sync::Cache; -use sha2::{Digest, Sha256}; +use aws_credential_types::provider::{SharedCredentialsProvider, error::CredentialsError}; +use litellm_secret_value::SecretValue; +use moka::{Expiry, future::Cache}; +use std::{future::Future, pin::Pin, sync::Arc}; use aws_credential_types::Credentials; use aws_credential_types::provider::ProvideCredentials; use aws_sigv4::http_request::{ - SignableBody, SignableRequest, SigningParams, SigningSettings, sign, + SignableBody, SignableRequest, SignatureLocation, SigningInstructions, SigningParams, + SigningSettings, sign, }; use aws_sigv4::sign::v4; use aws_smithy_runtime_api::client::identity::Identity; @@ -23,72 +26,106 @@ use super::constants::{ }; use crate::error::{CredentialFlow, Error, SigningStage}; -const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60); -const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600); +const CREDENTIAL_CACHE_TTL: Duration = Duration::from_secs(600); +const EXPIRY_BUFFER: Duration = Duration::from_secs(60); #[derive(Clone)] pub struct AwsAuthService { - static_credentials: Cache, - ambient_credentials: Cache, + providers: Cache, + credentials: Cache, + loader: Arc, } impl Default for AwsAuthService { fn default() -> Self { + Self::new(Arc::new(NativeProviderLoader)) + } +} + +impl AwsAuthService { + fn new(loader: Arc) -> Self { Self { - static_credentials: Cache::builder() + providers: Cache::builder().max_capacity(200).build(), + credentials: Cache::builder() .max_capacity(200) - .time_to_live(STATIC_CREDENTIALS_TTL) - .build(), - ambient_credentials: Cache::builder() - .max_capacity(200) - .time_to_live(AMBIENT_CREDENTIALS_TTL) + .expire_after(CredentialExpiry) .build(), + loader, } } } -fn credential_cache_ttl(flow: &AwsAuthFlow) -> Option { - match flow { - AwsAuthFlow::StaticKeys { .. } => Some(STATIC_CREDENTIALS_TTL), - AwsAuthFlow::DefaultChain => Some(AMBIENT_CREDENTIALS_TTL), - AwsAuthFlow::WebIdentity { .. } - | AwsAuthFlow::AssumeRole { .. } - | AwsAuthFlow::Profile { .. } - | AwsAuthFlow::SessionToken { .. } => None, +struct CredentialExpiry; + +impl Expiry for CredentialExpiry { + fn expire_after_create( + &self, + _: &AwsAuthConfig, + credentials: &Credentials, + _: std::time::Instant, + ) -> Option { + Some(credentials.expiry().map_or(CREDENTIAL_CACHE_TTL, |expiry| { + expiry + .duration_since(SystemTime::now()) + .unwrap_or_default() + .saturating_sub(EXPIRY_BUFFER) + .min(CREDENTIAL_CACHE_TTL) + })) } } -#[derive(Clone, Debug, Default, PartialEq, Eq)] +type ProviderFuture<'a> = + Pin> + Send + 'a>>; + +trait ProviderLoader: Send + Sync { + fn load(&self, config: AwsAuthConfig, flow: AwsAuthFlow) -> ProviderFuture<'_>; +} + +struct NativeProviderLoader; + +impl ProviderLoader for NativeProviderLoader { + fn load(&self, config: AwsAuthConfig, flow: AwsAuthFlow) -> ProviderFuture<'_> { + Box::pin(build_provider(config, flow)) + } +} + +#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)] pub struct AwsAuthConfig { - pub access_key_id: Option, - pub secret_access_key: Option, - pub session_token: Option, + pub access_key_id: Option, + pub secret_access_key: Option, + pub session_token: Option, pub region_name: Option, pub session_name: Option, pub profile_name: Option, pub role_name: Option, - pub web_identity_token: Option, + pub web_identity_token: Option, pub sts_endpoint: Option, - pub external_id: Option, + pub external_id: Option, } impl AwsAuthConfig { fn with_environment(self, env_lookup: &(dyn Fn(&str) -> Option + Sync)) -> Self { Self { - access_key_id: self.access_key_id.or_else(|| env_lookup(AWS_ACCESS_KEY_ID)), + access_key_id: self + .access_key_id + .or_else(|| env_lookup(AWS_ACCESS_KEY_ID).map(SecretValue::new)), secret_access_key: self .secret_access_key - .or_else(|| env_lookup(AWS_SECRET_ACCESS_KEY)), - session_token: self.session_token.or_else(|| env_lookup(AWS_SESSION_TOKEN)), + .or_else(|| env_lookup(AWS_SECRET_ACCESS_KEY).map(SecretValue::new)), + session_token: self + .session_token + .or_else(|| env_lookup(AWS_SESSION_TOKEN).map(SecretValue::new)), region_name: self.region_name.or_else(|| env_lookup(AWS_REGION_NAME)), session_name: self.session_name.or_else(|| env_lookup(AWS_SESSION_NAME)), profile_name: self.profile_name.or_else(|| env_lookup(AWS_PROFILE_NAME)), role_name: self.role_name.or_else(|| env_lookup(AWS_ROLE_NAME)), web_identity_token: self .web_identity_token - .or_else(|| env_lookup(AWS_WEB_IDENTITY_TOKEN)), + .or_else(|| env_lookup(AWS_WEB_IDENTITY_TOKEN).map(SecretValue::new)), sts_endpoint: self.sts_endpoint.or_else(|| env_lookup(AWS_STS_ENDPOINT)), - external_id: self.external_id.or_else(|| env_lookup(AWS_EXTERNAL_ID)), + external_id: self + .external_id + .or_else(|| env_lookup(AWS_EXTERNAL_ID).map(SecretValue::new)), } } } @@ -96,52 +133,31 @@ impl AwsAuthConfig { #[derive(Clone, Debug, PartialEq, Eq)] pub enum AwsAuthFlow { WebIdentity { - token: String, + token: SecretValue, role: String, session_name: String, }, AssumeRole { role: String, session_name: Option, + web_identity_role: Option, }, Profile { name: String, }, SessionToken { - access_key_id: String, - secret_access_key: String, - session_token: String, + access_key_id: SecretValue, + secret_access_key: SecretValue, + session_token: SecretValue, }, StaticKeys { - access_key_id: String, - secret_access_key: String, + access_key_id: SecretValue, + secret_access_key: SecretValue, region_name: String, }, DefaultChain, } -fn cache_key(config: &AwsAuthConfig, flow: &AwsAuthFlow) -> String { - let mut hasher = Sha256::new(); - hasher.update(format!("{config:?}:{flow:?}")); - format!("{:x}", hasher.finalize()) -} - -impl AwsAuthService { - fn get_cached_credentials(&self, key: &str) -> Option { - self.static_credentials - .get(key) - .or_else(|| self.ambient_credentials.get(key)) - } - - fn set_cached_credentials(&self, key: String, credentials: Credentials, ttl: Duration) { - if ttl == STATIC_CREDENTIALS_TTL { - self.static_credentials.insert(key, credentials); - } else { - self.ambient_credentials.insert(key, credentials); - } - } -} - fn role_identity(arn: &str) -> Option<(&str, &str, &str)> { let mut parts = arn.splitn(6, ':'); let ("arn", partition, _, _, account, resource) = ( @@ -170,6 +186,9 @@ pub fn classify_auth( config: AwsAuthConfig, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> AwsAuthFlow { + let web_identity_role = env_lookup(AWS_WEB_IDENTITY_TOKEN_FILE) + .filter(|token_file| !token_file.is_empty()) + .and_then(|_| env_lookup(AWS_ROLE_ARN)); let config = config.with_environment(env_lookup); if let (Some(token), Some(role), Some(session_name)) = ( config.web_identity_token.clone(), @@ -186,6 +205,7 @@ pub fn classify_auth( return AwsAuthFlow::AssumeRole { role, session_name: config.session_name.clone(), + web_identity_role, }; } if let Some(name) = config.profile_name { @@ -216,6 +236,30 @@ pub fn classify_auth( AwsAuthFlow::DefaultChain } +async fn profile_provider( + name: &str, + region: Option, + config: aws_config::provider_config::ProviderConfig, + builder: aws_config::profile::credentials::Builder, +) -> Result { + let region = match region { + Some(region) => Some(aws_types::region::Region::new(region)), + None => { + aws_config::default_provider::region::DefaultRegionChain::builder() + .profile_name(name) + .build() + .region() + .await + } + }; + Ok(SharedCredentialsProvider::new( + builder + .configure(&config.with_region(region)) + .profile_name(name) + .build(), + )) +} + impl AwsAuthService { pub async fn resolve_credentials( &self, @@ -224,183 +268,221 @@ impl AwsAuthService { ) -> Result { let resolved = config.clone().with_environment(env_lookup); let flow = classify_auth(config, env_lookup); - match flow { - AwsAuthFlow::SessionToken { - access_key_id, - secret_access_key, - session_token, - } => Ok(Credentials::new( - access_key_id, - secret_access_key, - Some(session_token), - None, - "litellm-static-session", - )), - AwsAuthFlow::StaticKeys { - access_key_id, - secret_access_key, - region_name, - } => { - let flow = AwsAuthFlow::StaticKeys { - access_key_id: access_key_id.clone(), - secret_access_key: secret_access_key.clone(), - region_name, - }; - let key = cache_key(&resolved, &flow); - if let Some(credentials) = self.get_cached_credentials(&key) { - return Ok(credentials); - } - let credentials = Credentials::new( - access_key_id, - secret_access_key, - None, - None, - "litellm-static", - ); - self.set_cached_credentials( - key, - credentials.clone(), - credential_cache_ttl(&flow).unwrap_or(STATIC_CREDENTIALS_TTL), - ); - Ok(credentials) - } - AwsAuthFlow::Profile { name } => { - let provider = aws_config::profile::ProfileFileCredentialsProvider::builder() - .profile_name(name) - .build(); + let cached = self.cached_credentials(&resolved, &flow).await?; + if fresh(&cached) { + return Ok(cached); + } + self.credentials.invalidate(&resolved).await; + let refetched = self.cached_credentials(&resolved, &flow).await?; + if fresh(&refetched) { + return Ok(refetched); + } + self.credentials.invalidate(&resolved).await; + Err(Error::ExpiredCredentials) + } + + async fn cached_credentials( + &self, + resolved: &AwsAuthConfig, + flow: &AwsAuthFlow, + ) -> Result { + self.credentials + .try_get_with(resolved.clone(), async { + let provider = self + .providers + .try_get_with( + resolved.clone(), + self.loader.load(resolved.clone(), flow.clone()), + ) + .await + .map_err(|error| (*error).clone())?; provider .provide_credentials() .await - .map_err(Error::credentials(CredentialFlow::Profile)) - } - AwsAuthFlow::AssumeRole { role, session_name } => { - if is_already_running_as_role(&role, &resolved).await? { - let ambient_flow = AwsAuthFlow::DefaultChain; - let key = cache_key(&resolved, &ambient_flow); - if let Some(credentials) = self.get_cached_credentials(&key) { - return Ok(credentials); - } - let provider = - aws_config::default_provider::credentials::DefaultCredentialsChain::builder() - .build() - .await; - let credentials = provider - .provide_credentials() - .await - .map_err(Error::credentials(CredentialFlow::DefaultChain))?; - self.set_cached_credentials( - key, - credentials.clone(), - credential_cache_ttl(&ambient_flow).unwrap_or(AMBIENT_CREDENTIALS_TTL), - ); - return Ok(credentials); - } - let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); - if let Some(region) = resolved.region_name.clone() { - loader = loader.region(aws_types::region::Region::new(region)); - } - if let Some(endpoint) = resolved.sts_endpoint.clone() { - loader = loader.endpoint_url(endpoint); - } - if let (Some(access_key_id), Some(secret_access_key)) = - (resolved.access_key_id, resolved.secret_access_key) - { - loader = loader.credentials_provider(Credentials::new( - access_key_id, - secret_access_key, - resolved.session_token, - None, - "litellm-role-source", - )); - } - let sdk_config = loader.load().await; - let builder = aws_config::sts::AssumeRoleProvider::builder(role); - let builder = match session_name { - Some(name) => builder.session_name(name), - None => builder.session_name(default_session_name()), - }; - let builder = match resolved.external_id { - Some(id) => builder.external_id(id), - None => builder, - }; - let provider = builder.configure(&sdk_config).build().await; - provider - .provide_credentials() - .await - .map_err(Error::credentials(CredentialFlow::AssumeRole)) - } - AwsAuthFlow::WebIdentity { - token, - role, - session_name, - } => { - let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); - if let Some(region) = resolved.region_name { - loader = loader.region(aws_types::region::Region::new(region)); - } - if let Some(endpoint) = resolved.sts_endpoint { - loader = loader.endpoint_url(endpoint); - } - let sdk_config = loader.load().await; - let client = aws_sdk_sts::Client::new(&sdk_config); - let response = client - .assume_role_with_web_identity() - .role_arn(role) - .role_session_name(session_name) - .web_identity_token(token) - .send() - .await - .map_err(Error::credentials(CredentialFlow::WebIdentity))?; - let credentials = response - .credentials() - .ok_or(Error::MissingWebIdentityCredentials)?; - let expiration = - SystemTime::try_from(*credentials.expiration()).map_err(|source| { - Error::WebIdentityExpiration(litellm_error::Source::new(source)) - })?; - Ok(Credentials::new( - credentials.access_key_id(), - credentials.secret_access_key(), - Some(credentials.session_token().to_string()), - Some(expiration), - "litellm-web-identity", - )) - } - AwsAuthFlow::DefaultChain => { - let key = cache_key(&resolved, &AwsAuthFlow::DefaultChain); - if let Some(credentials) = self.get_cached_credentials(&key) { - return Ok(credentials); - } - let provider = - aws_config::default_provider::credentials::DefaultCredentialsChain::builder() - .build() - .await; - let credentials = provider - .provide_credentials() - .await - .map_err(Error::credentials(CredentialFlow::DefaultChain))?; - self.set_cached_credentials( - key, - credentials.clone(), - credential_cache_ttl(&AwsAuthFlow::DefaultChain) - .unwrap_or(AMBIENT_CREDENTIALS_TTL), - ); - Ok(credentials) + .map_err(Error::credentials(flow.credential_flow())) + }) + .await + .map_err(|error| (*error).clone()) + } +} + +fn fresh(credentials: &Credentials) -> bool { + credentials + .expiry() + .is_none_or(|expiry| expiry > SystemTime::now() + EXPIRY_BUFFER) +} + +impl AwsAuthFlow { + fn credential_flow(&self) -> CredentialFlow { + match self { + Self::Profile { .. } => CredentialFlow::Profile, + Self::AssumeRole { .. } => CredentialFlow::AssumeRole, + Self::WebIdentity { .. } => CredentialFlow::WebIdentity, + Self::DefaultChain | Self::StaticKeys { .. } | Self::SessionToken { .. } => { + CredentialFlow::DefaultChain } } } } -async fn is_already_running_as_role(role: &str, config: &AwsAuthConfig) -> Result { +async fn build_provider( + resolved: AwsAuthConfig, + flow: AwsAuthFlow, +) -> Result { + match flow { + AwsAuthFlow::SessionToken { + access_key_id, + secret_access_key, + session_token, + } => Ok(SharedCredentialsProvider::new(Credentials::new( + access_key_id.expose(), + secret_access_key.expose(), + Some(session_token.expose().to_owned()), + None, + "litellm-static-session", + ))), + AwsAuthFlow::StaticKeys { + access_key_id, + secret_access_key, + .. + } => Ok(SharedCredentialsProvider::new(Credentials::new( + access_key_id.expose(), + secret_access_key.expose(), + None, + None, + "litellm-static", + ))), + AwsAuthFlow::Profile { name } => { + profile_provider( + &name, + resolved.region_name, + aws_config::provider_config::ProviderConfig::without_region(), + aws_config::profile::ProfileFileCredentialsProvider::builder(), + ) + .await + } + AwsAuthFlow::DefaultChain => { + let builder = + aws_config::default_provider::credentials::DefaultCredentialsChain::builder(); + let builder = match resolved.region_name { + Some(region) => builder.region(aws_types::region::Region::new(region)), + None => builder, + }; + Ok(SharedCredentialsProvider::new(builder.build().await)) + } + AwsAuthFlow::AssumeRole { + role, + session_name, + web_identity_role, + } => { + if is_already_running_as_role(&role, web_identity_role.as_deref(), &resolved).await? { + return Ok(SharedCredentialsProvider::new( + aws_config::default_provider::credentials::DefaultCredentialsChain::builder() + .build() + .await, + )); + } + let sdk_config = sdk_config(resolved.clone()).await; + let builder = aws_config::sts::AssumeRoleProvider::builder(role) + .session_name(session_name.unwrap_or_else(default_session_name)); + let builder = match resolved.external_id { + Some(id) => builder.external_id(id.expose()), + None => builder, + }; + Ok(SharedCredentialsProvider::new( + builder.configure(&sdk_config).build().await, + )) + } + AwsAuthFlow::WebIdentity { + token, + role, + session_name, + } => Ok(SharedCredentialsProvider::new(WebIdentityProvider { + client: aws_sdk_sts::Client::new(&sdk_config(resolved).await), + token, + role, + session_name, + })), + } +} + +async fn sdk_config(config: AwsAuthConfig) -> aws_types::SdkConfig { + let loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); + let loader = match config.region_name { + Some(region) => loader.region(aws_types::region::Region::new(region)), + None => loader, + }; + let loader = match config.sts_endpoint { + Some(endpoint) => loader.endpoint_url(endpoint), + None => loader, + }; + let loader = match (config.access_key_id, config.secret_access_key) { + (Some(key), Some(secret)) => loader.credentials_provider(Credentials::new( + key.expose(), + secret.expose(), + config.session_token.map(|token| token.expose().to_owned()), + None, + "litellm-role-source", + )), + _ => loader, + }; + loader.load().await +} + +#[derive(Debug)] +struct WebIdentityProvider { + client: aws_sdk_sts::Client, + token: SecretValue, + role: String, + session_name: String, +} + +impl ProvideCredentials for WebIdentityProvider { + fn provide_credentials<'a>( + &'a self, + ) -> aws_credential_types::provider::future::ProvideCredentials<'a> + where + Self: 'a, + { + aws_credential_types::provider::future::ProvideCredentials::new(async { + let response = self + .client + .assume_role_with_web_identity() + .role_arn(&self.role) + .role_session_name(&self.session_name) + .web_identity_token(self.token.expose()) + .send() + .await + .map_err(CredentialsError::provider_error)?; + let credentials = response.credentials().ok_or_else(|| { + CredentialsError::provider_error(Error::MissingWebIdentityCredentials) + })?; + let expiration = SystemTime::try_from(*credentials.expiration()) + .map_err(CredentialsError::provider_error)?; + Ok(Credentials::new( + credentials.access_key_id(), + credentials.secret_access_key(), + Some(credentials.session_token().to_owned()), + Some(expiration), + "litellm-web-identity", + )) + }) + } +} + +async fn is_already_running_as_role( + role: &str, + web_identity_role: Option<&str>, + config: &AwsAuthConfig, +) -> Result { + if config.access_key_id.is_some() || config.secret_access_key.is_some() { + return Ok(false); + } if role_identity(role).is_none() { return Ok(false); } - if let (Ok(current_role), Ok(token_file)) = ( - std::env::var(AWS_ROLE_ARN), - std::env::var(AWS_WEB_IDENTITY_TOKEN_FILE), - ) && !token_file.is_empty() - { - return Ok(same_role_arns(role, ¤t_role)); + if let Some(current_role) = web_identity_role { + return Ok(same_role_arns(role, current_role)); } let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); @@ -473,16 +555,6 @@ pub fn sign_headers( headers: &HeaderMap, scope: SigningScope<'_>, ) -> Result { - let identity: Identity = scope.credentials.clone().into(); - let params = v4::SigningParams::builder() - .identity(&identity) - .region(scope.region) - .name(scope.service) - .time(scope.time) - .settings(SigningSettings::default()) - .build() - .map(SigningParams::from) - .map_err(Error::signing(SigningStage::Parameters))?; let header_refs = headers .iter() .map(|(name, value)| { @@ -494,16 +566,14 @@ pub fn sign_headers( }) }) .collect::, _>>()?; - let request = SignableRequest::new( + let instructions = sign_request( method, url, header_refs.into_iter(), - SignableBody::Bytes(body), - ) - .map_err(Error::signing(SigningStage::Request))?; - let (instructions, _) = sign(request, ¶ms) - .map_err(Error::signing(SigningStage::Signature))? - .into_parts(); + body, + SigningSettings::default(), + scope, + )?; instructions .headers() .map(|(name, value)| { @@ -516,6 +586,60 @@ pub fn sign_headers( .collect() } +pub fn presign_url( + url: &str, + expires_in: Duration, + scope: SigningScope<'_>, +) -> Result { + let mut signed = url::Url::parse(url).map_err(Error::signing(SigningStage::Request))?; + let mut settings = SigningSettings::default(); + settings.signature_location = SignatureLocation::QueryParams; + settings.expires_in = Some(expires_in); + let instructions = sign_request( + "GET", + signed.as_str(), + std::iter::empty(), + &[], + settings, + scope, + )?; + let (signature, auth): (Vec<_>, Vec<_>) = instructions + .params() + .iter() + .map(|(name, value)| (*name, value.as_ref())) + .partition(|(name, _)| *name == "X-Amz-Signature"); + signed + .query_pairs_mut() + .extend_pairs(auth.into_iter().chain(signature)); + Ok(signed.into()) +} + +fn sign_request<'a>( + method: &'a str, + url: &'a str, + headers: impl Iterator, + body: &'a [u8], + settings: SigningSettings, + scope: SigningScope<'_>, +) -> Result { + let identity: Identity = scope.credentials.clone().into(); + let params = v4::SigningParams::builder() + .identity(&identity) + .region(scope.region) + .name(scope.service) + .time(scope.time) + .settings(settings) + .build() + .map(SigningParams::from) + .map_err(Error::signing(SigningStage::Parameters))?; + let request = SignableRequest::new(method, url, headers, SignableBody::Bytes(body)) + .map_err(Error::signing(SigningStage::Request))?; + let (instructions, _) = sign(request, ¶ms) + .map_err(Error::signing(SigningStage::Signature))? + .into_parts(); + Ok(instructions) +} + /// Model-id and region parsing shared by every Bedrock route. pub fn bedrock_model_id_and_region(model: &str) -> (String, Option) { let mut stripped = model; @@ -586,16 +710,26 @@ pub fn aws_auth_config( env_lookup: &dyn Fn(&str) -> Option, ) -> AwsAuthConfig { AwsAuthConfig { - access_key_id: params.resolve(&AwsParams::ACCESS_KEY_ID, env_lookup), - secret_access_key: params.resolve(&AwsParams::SECRET_ACCESS_KEY, env_lookup), - session_token: params.resolve(&AwsParams::SESSION_TOKEN, env_lookup), + access_key_id: params + .resolve(&AwsParams::ACCESS_KEY_ID, env_lookup) + .map(SecretValue::new), + secret_access_key: params + .resolve(&AwsParams::SECRET_ACCESS_KEY, env_lookup) + .map(SecretValue::new), + session_token: params + .resolve(&AwsParams::SESSION_TOKEN, env_lookup) + .map(SecretValue::new), region_name: params.resolve(&AwsParams::REGION, env_lookup), session_name: params.resolve(&AwsParams::SESSION_NAME, env_lookup), profile_name: params.resolve(&AwsParams::PROFILE_NAME, env_lookup), role_name: params.resolve(&AwsParams::ROLE_NAME, env_lookup), - web_identity_token: params.resolve(&AwsParams::WEB_IDENTITY_TOKEN, env_lookup), + web_identity_token: params + .resolve(&AwsParams::WEB_IDENTITY_TOKEN, env_lookup) + .map(SecretValue::new), sts_endpoint: params.resolve(&AwsParams::STS_ENDPOINT, env_lookup), - external_id: params.resolve(&AwsParams::EXTERNAL_ID, env_lookup), + external_id: params + .resolve(&AwsParams::EXTERNAL_ID, env_lookup) + .map(SecretValue::new), } } @@ -660,10 +794,133 @@ mod tests { use crate::constants::{AWS_REGION, BEDROCK_SERVICE}; use reqwest::header::CONTENT_TYPE; + #[cfg(unix)] + #[rstest::fixture] + fn process_profile() -> aws_runtime::env_config::file::EnvConfigFiles { + let credentials = serde_json::json!({ + "Version": 1, + "AccessKeyId": "process-access", + "SecretAccessKey": "process-secret", + "SessionToken": "process-session", + }); + aws_runtime::env_config::file::EnvConfigFiles::builder() + .include_default_config_file(false) + .include_default_credentials_file(false) + .with_contents( + aws_runtime::env_config::file::EnvConfigFileKind::Config, + format!( + "[profile smoke]\ncredential_process = /usr/bin/printf '%s' '{credentials}'\n" + ), + ) + .build() + } + + #[cfg(unix)] + #[rstest::rstest] + #[tokio::test] + async fn profile_process_credentials_can_be_loaded( + process_profile: aws_runtime::env_config::file::EnvConfigFiles, + ) { + let provider = profile_provider( + "smoke", + Some("us-west-2".into()), + aws_config::provider_config::ProviderConfig::without_region(), + aws_config::profile::ProfileFileCredentialsProvider::builder() + .profile_files(process_profile), + ) + .await + .unwrap(); + let credentials = provider.provide_credentials().await.unwrap(); + + assert_eq!(credentials.access_key_id(), "process-access"); + assert_eq!(credentials.secret_access_key(), "process-secret"); + assert_eq!(credentials.session_token(), Some("process-session")); + } + + #[derive(Debug)] + struct StsConnector; + + impl aws_smithy_runtime_api::client::http::HttpConnector for StsConnector { + fn call( + &self, + request: aws_smithy_runtime_api::client::orchestrator::HttpRequest, + ) -> aws_smithy_runtime_api::client::http::HttpConnectorFuture { + assert!(request.uri().contains("sts.us-west-2.")); + aws_smithy_runtime_api::client::http::HttpConnectorFuture::new(async { + Ok( + aws_smithy_runtime_api::client::orchestrator::HttpResponse::new( + 200_u16.try_into().unwrap(), + aws_smithy_types::body::SdkBody::from( + r#" + + role-accessrole-secret + role-session2100-01-01T00:00:00Z + arn:aws:sts::123456789012:assumed-role/smoke/session + smoke:session + smoke"#, + ), + ), + ) + }) + } + } + + #[rstest::rstest] + #[tokio::test] + async fn role_profile_uses_region_when_acquiring_sts_credentials() { + use aws_smithy_runtime_api::client::http::{SharedHttpConnector, http_client_fn}; + let files = aws_runtime::env_config::file::EnvConfigFiles::builder() + .include_default_config_file(false) + .include_default_credentials_file(false) + .with_contents( + aws_runtime::env_config::file::EnvConfigFileKind::Config, + "[profile smoke]\nrole_arn = arn:aws:iam::123456789012:role/smoke\nsource_profile = source\n\n[profile source]\naws_access_key_id = source-access\naws_secret_access_key = source-secret\n", + ).build(); + let config = aws_config::provider_config::ProviderConfig::without_region() + .with_http_client(http_client_fn(|_, _| { + SharedHttpConnector::new(StsConnector) + })); + let credentials = profile_provider( + "smoke", + Some("us-west-2".into()), + config, + aws_config::profile::ProfileFileCredentialsProvider::builder().profile_files(files), + ) + .await + .unwrap(); + let credentials = credentials.provide_credentials().await.unwrap(); + assert_eq!(credentials.access_key_id(), "role-access"); + assert_eq!(credentials.secret_access_key(), "role-secret"); + assert_eq!(credentials.session_token(), Some("role-session")); + assert!(credentials.expiry().is_some()); + } + fn no_env(_: &str) -> Option { None } + #[rstest::rstest] + #[case::access_key(Some("source-access"), None)] + #[case::secret_key(None, Some("source-secret"))] + #[case::source_pair(Some("source-access"), Some("source-secret"))] + #[tokio::test] + async fn explicit_source_keys_cannot_reuse_an_ambient_role( + #[case] access_key: Option<&str>, + #[case] secret_key: Option<&str>, + ) { + let role = "arn:aws:iam::123456789012:role/target"; + let config = AwsAuthConfig { + access_key_id: access_key.map(SecretValue::new), + secret_access_key: secret_key.map(SecretValue::new), + ..Default::default() + }; + assert!( + !is_already_running_as_role(role, Some(role), &config) + .await + .unwrap() + ); + } + #[test] fn secret_names_cover_environment_reads() { let seen = std::sync::Arc::new(std::sync::Mutex::new( @@ -706,7 +963,10 @@ mod tests { }; assert_eq!( - aws_auth_config(¶ms, &env).secret_access_key.as_deref(), + aws_auth_config(¶ms, &env) + .secret_access_key + .as_ref() + .map(SecretValue::expose), expected ); assert_eq!( @@ -831,6 +1091,31 @@ mod tests { ); } + #[rstest::rstest] + #[case::running_with_web_identity("/var/run/token", Some("arn:aws:iam::1:role/current"))] + #[case::empty_token_file("", None)] + fn assume_role_reads_the_ambient_role_from_the_injected_environment( + #[case] token_file: &str, + #[case] expected: Option<&str>, + ) { + let env = |key: &str| match key { + AWS_WEB_IDENTITY_TOKEN_FILE => Some(token_file.to_string()), + AWS_ROLE_ARN => Some("arn:aws:iam::1:role/current".to_string()), + _ => None, + }; + let config = AwsAuthConfig { + role_name: Some("arn:aws:iam::1:role/target".into()), + ..Default::default() + }; + let AwsAuthFlow::AssumeRole { + web_identity_role, .. + } = classify_auth(config, &env) + else { + panic!("a role name selects AssumeRole"); + }; + assert_eq!(web_identity_role.as_deref(), expected); + } + #[tokio::test] async fn static_credentials_do_not_use_network() { let credentials = AwsAuthService::default() @@ -849,110 +1134,180 @@ mod tests { assert_eq!(credentials.session_token(), None); } - #[test] - fn cache_policy_matches_python_flows() { - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::StaticKeys { - access_key_id: "ak".into(), - secret_access_key: "sk".into(), - region_name: "us-east-1".into(), - }), - Some(STATIC_CREDENTIALS_TTL) - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::DefaultChain), - Some(AMBIENT_CREDENTIALS_TTL) - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::SessionToken { - access_key_id: "ak".into(), - secret_access_key: "sk".into(), - session_token: "token".into(), - }), - None - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::Profile { - name: "profile".into() - }), - None - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::AssumeRole { - role: "arn:aws:iam::123456789012:role/demo".into(), - session_name: None, - }), - None - ); - assert_eq!( - credential_cache_ttl(&AwsAuthFlow::WebIdentity { - token: "token".into(), - role: "arn:aws:iam::123456789012:role/demo".into(), - session_name: "session".into(), - }), - None - ); + struct CountingLoader { + loads: Arc, + calls: Arc, + short_lifetime: Duration, + short_calls: usize, } - #[rstest::rstest] - fn cache_round_trip_preserves_credentials() { - let auth = AwsAuthService::default(); - let key = format!("cache-test-{}", std::process::id()); - let credentials = Credentials::new("cache-ak", "cache-sk", None, None, "test"); - auth.set_cached_credentials(key.clone(), credentials.clone(), STATIC_CREDENTIALS_TTL); - assert_eq!( - auth.get_cached_credentials(&key) - .map(|value| value.access_key_id().to_string()), - Some("cache-ak".to_string()) - ); + impl ProviderLoader for CountingLoader { + fn load(&self, config: AwsAuthConfig, _: AwsAuthFlow) -> ProviderFuture<'_> { + let loads = self.loads.clone(); + let provider = CountingProvider { + calls: self.calls.clone(), + short_lifetime: self.short_lifetime, + short_calls: self.short_calls, + key: config.access_key_id.unwrap(), + }; + Box::pin(async move { + loads.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(SharedCredentialsProvider::new(provider)) + }) + } + } + + #[derive(Debug)] + struct CountingProvider { + calls: Arc, + short_lifetime: Duration, + short_calls: usize, + key: SecretValue, + } + + impl ProvideCredentials for CountingProvider { + fn provide_credentials<'a>( + &'a self, + ) -> aws_credential_types::provider::future::ProvideCredentials<'a> + where + Self: 'a, + { + aws_credential_types::provider::future::ProvideCredentials::new(async move { + tokio::task::yield_now().await; + let call = self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let lifetime = if call < self.short_calls { + self.short_lifetime + } else { + Duration::from_secs(3600) + }; + Ok(Credentials::new( + self.key.expose(), + "credential-value", + Some(format!("session-{call}")), + Some(SystemTime::now() + lifetime), + "counting", + )) + }) + } } #[rstest::rstest] #[tokio::test] - async fn cloned_services_reuse_credentials_but_independent_services_do_not() { - let auth = AwsAuthService::default(); + async fn near_expiry_credentials_are_refreshed_without_rebuilding_the_provider() { + use std::sync::atomic::{AtomicUsize, Ordering}; + let loads = Arc::new(AtomicUsize::new(0)); + let calls = Arc::new(AtomicUsize::new(0)); + let auth = AwsAuthService::new(Arc::new(CountingLoader { + loads: loads.clone(), + calls: calls.clone(), + short_lifetime: Duration::from_secs(30), + short_calls: 1, + })); let config = AwsAuthConfig { - access_key_id: Some("configured-key".into()), - secret_access_key: Some("configured-secret".into()), - region_name: Some("us-east-1".into()), - ..AwsAuthConfig::default() + access_key_id: Some("identity".into()), + ..Default::default() }; - let flow = classify_auth(config.clone(), &no_env); - let cached = Credentials::new("cached-key", "cached-secret", None, None, "test"); - auth.set_cached_credentials( - cache_key(&config, &flow), - cached.clone(), - STATIC_CREDENTIALS_TTL, - ); - - let reused = auth + let first = auth + .resolve_credentials(config.clone(), &no_env) + .await + .unwrap(); + let second = auth .clone() .resolve_credentials(config.clone(), &no_env) .await .unwrap(); - let independent = AwsAuthService::default() - .resolve_credentials(config.clone(), &no_env) - .await - .unwrap(); - let different = AwsAuthConfig { - access_key_id: Some("different-key".into()), - ..config.clone() - }; - let other_identity = auth - .resolve_credentials(different.clone(), &no_env) - .await - .unwrap(); + assert_eq!(first.session_token(), Some("session-1")); + assert_eq!(first.session_token(), second.session_token()); + assert_eq!(loads.load(Ordering::SeqCst), 1); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } - assert_eq!(reused.access_key_id(), cached.access_key_id()); - assert_eq!(reused.secret_access_key(), cached.secret_access_key()); + #[rstest::rstest] + #[case::recovered_by_one_refetch(1, Some("session-1"))] + #[case::still_expired_after_refetch(2, None)] + #[tokio::test] + async fn expired_credentials_are_refetched_once( + #[case] short_calls: usize, + #[case] session: Option<&str>, + ) { + use std::sync::atomic::{AtomicUsize, Ordering}; + let calls = Arc::new(AtomicUsize::new(0)); + let auth = AwsAuthService::new(Arc::new(CountingLoader { + loads: Arc::new(AtomicUsize::new(0)), + calls: calls.clone(), + short_lifetime: Duration::ZERO, + short_calls, + })); + let config = AwsAuthConfig { + access_key_id: Some("identity".into()), + ..Default::default() + }; + let result = auth.resolve_credentials(config, &no_env).await; assert_eq!( - Some(independent.access_key_id()), - config.access_key_id.as_deref() + result.as_ref().ok().and_then(Credentials::session_token), + session ); - assert_eq!( - Some(other_identity.access_key_id()), - different.access_key_id.as_deref() + assert!(session.is_some() || matches!(result, Err(Error::ExpiredCredentials))); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[rstest::rstest] + #[tokio::test] + async fn concurrent_calls_share_acquisition_and_distinct_identities_stay_isolated() { + use std::sync::atomic::{AtomicUsize, Ordering}; + let loads = Arc::new(AtomicUsize::new(0)); + let calls = Arc::new(AtomicUsize::new(0)); + let auth = AwsAuthService::new(Arc::new(CountingLoader { + loads: loads.clone(), + calls: calls.clone(), + short_lifetime: Duration::from_secs(3600), + short_calls: 0, + })); + let first = AwsAuthConfig { + access_key_id: Some("first".into()), + ..Default::default() + }; + let second = AwsAuthConfig { + access_key_id: Some("second".into()), + ..Default::default() + }; + let (a, b, c) = tokio::join!( + auth.resolve_credentials(first.clone(), &no_env), + auth.resolve_credentials(first, &no_env), + auth.resolve_credentials(second, &no_env) ); + let a = a.unwrap(); + let b = b.unwrap(); + let c = c.unwrap(); + assert_eq!(a.session_token(), b.session_token()); + assert_ne!(a.access_key_id(), c.access_key_id()); + assert_eq!(loads.load(Ordering::SeqCst), 2); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[rstest::rstest] + fn auth_config_debug_never_discloses_credentials() { + let config = AwsAuthConfig { + access_key_id: Some("credential-value".into()), + secret_access_key: Some("credential-value".into()), + session_token: Some("credential-value".into()), + web_identity_token: Some("credential-value".into()), + external_id: Some("credential-value".into()), + ..Default::default() + }; + let output = format!("{config:?}"); + assert!(!output.contains("credential-value")); + assert!(output.contains("[REDACTED]")); + } + + #[rstest::rstest] + #[case::session(AwsAuthFlow::SessionToken { access_key_id: "credential-value".into(), secret_access_key: "credential-value".into(), session_token: "credential-value".into() })] + #[case::static_keys(AwsAuthFlow::StaticKeys { access_key_id: "credential-value".into(), secret_access_key: "credential-value".into(), region_name: "region".into() })] + #[case::web_identity(AwsAuthFlow::WebIdentity { token: "credential-value".into(), role: "role".into(), session_name: "session".into() })] + fn auth_flow_debug_never_discloses_credentials(#[case] flow: AwsAuthFlow) { + let output = format!("{flow:?}"); + assert!(!output.contains("credential-value")); + assert!(output.contains("[REDACTED]")); } #[test] @@ -1142,8 +1497,8 @@ mod tests { let credentials = AwsAuthService::default() .resolve_credentials( AwsAuthConfig { - access_key_id: Some(access_key_id), - secret_access_key: Some(secret_access_key), + access_key_id: Some(access_key_id.into()), + secret_access_key: Some(secret_access_key.into()), region_name: Some("us-west-2".to_string()), ..Default::default() }, @@ -1191,4 +1546,44 @@ mod tests { failures.join("; ") ); } + + #[rstest::rstest] + #[case::provisioned( + None, + "https://cache-1/?Action=connect&User=iam-user", + "cache-1/?Action=connect&User=iam-user&X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIDEXAMPLE%2F20261029%2Fus-east-1%2Felasticache%2Faws4_request&X-Amz-Date=20261029T120000Z&X-Amz-Expires=900&X-Amz-SignedHeaders=host&X-Amz-Signature=dfa58492a6ae9ba065088724170a52437329d03f3e9bc9a4c90a80cc8a18e2f6" + )] + #[case::serverless_with_a_session_token( + Some("session-token"), + "https://cache-1/?Action=connect&User=iam-user&ResourceType=ServerlessCache", + "cache-1/?Action=connect&User=iam-user&ResourceType=ServerlessCache&X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIDEXAMPLE%2F20261029%2Fus-east-1%2Felasticache%2Faws4_request&X-Amz-Date=20261029T120000Z&X-Amz-Expires=900&X-Amz-SignedHeaders=host&X-Amz-Security-Token=session-token&X-Amz-Signature=e48bdbb3bf9d5e8f1d325bf33d5260e333df1f6b8868e8edd1e4c3cc01b4c3f9" + )] + fn presigned_urls_match_botocore_sigv4_query_auth( + #[case] session_token: Option<&str>, + #[case] url: &str, + #[case] expected: &str, + ) { + // Expected values from botocore 1.43.117 SigV4QueryAuth with get_current_datetime patched to + // 2026-10-29T12:00:00Z (generated 2026-10-29), the signer Python's ElastiCache IAM provider uses + let credentials = Credentials::new( + "AKIDEXAMPLE", + "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY", + session_token.map(str::to_owned), + None, + "test", + ); + let signed = presign_url( + url, + Duration::from_secs(900), + SigningScope { + region: "us-east-1", + service: "elasticache", + credentials: &credentials, + time: UNIX_EPOCH + Duration::from_secs(1_793_275_200), + }, + ) + .unwrap(); + + assert_eq!(signed.strip_prefix("https://"), Some(expected)); + } } diff --git a/litellm-rust/crates/auth-aws/src/error.rs b/litellm-rust/crates/auth-aws/src/error.rs index 5f88fc8a171..5445de1e6c3 100644 --- a/litellm-rust/crates/auth-aws/src/error.rs +++ b/litellm-rust/crates/auth-aws/src/error.rs @@ -2,6 +2,8 @@ use litellm_error::Source; #[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] pub enum Error { + #[error("AWS credential provider returned expired credentials")] + ExpiredCredentials, #[error("could not load AWS {flow} credentials")] Credentials { flow: CredentialFlow, diff --git a/litellm-rust/crates/auth-aws/src/signer.rs b/litellm-rust/crates/auth-aws/src/signer.rs index dff917d7c13..2d9ad9ba123 100644 --- a/litellm-rust/crates/auth-aws/src/signer.rs +++ b/litellm-rust/crates/auth-aws/src/signer.rs @@ -75,7 +75,8 @@ impl RequestSigner for SigV4Signer { ) .map_err(|error| match error { Error::NonAsciiHeader { .. } => litellm_http::Error::unsignable_header(error), - Error::Credentials { .. } + Error::ExpiredCredentials + | Error::Credentials { .. } | Error::MissingWebIdentityCredentials | Error::WebIdentityExpiration(_) | Error::Signing { .. } => litellm_http::Error::signature(error), diff --git a/litellm-rust/crates/auth-azure/Cargo.toml b/litellm-rust/crates/auth-azure/Cargo.toml index 7dcb0e513a7..540e0d443f6 100644 --- a/litellm-rust/crates/auth-azure/Cargo.toml +++ b/litellm-rust/crates/auth-azure/Cargo.toml @@ -8,17 +8,18 @@ repository.workspace = true [dependencies] litellm-secret-value.workspace = true litellm-auth-types.workspace = true +litellm-error.workspace = true moka.workspace = true serde_json.workspace = true sha2.workspace = true strum.workspace = true +tokio.workspace = true url.workspace = true azure_core = "1.0.0" azure_identity = { version = "1.0.0", features = ["tokio"] } [dev-dependencies] -litellm-error.workspace = true rstest.workspace = true -tokio.workspace = true +tokio = { workspace = true, features = ["test-util"] } diff --git a/litellm-rust/crates/auth-azure/src/native.rs b/litellm-rust/crates/auth-azure/src/native.rs index f40745ab7a2..03021428ea7 100644 --- a/litellm-rust/crates/auth-azure/src/native.rs +++ b/litellm-rust/crates/auth-azure/src/native.rs @@ -59,6 +59,7 @@ pub(crate) enum NativeAzureRequest { pub(crate) struct ValidatedAzureRequest { request: NativeAzureRequest, credential_source: InputSource, + probe_imds: bool, } impl ValidatedAzureRequest { @@ -68,9 +69,18 @@ impl ValidatedAzureRequest { Ok(Self { request, credential_source, + probe_imds: false, }) } + pub(crate) fn with_imds_probe(self, enabled: bool) -> Self { + Self { + probe_imds: enabled + && matches!(self.request, NativeAzureRequest::ManagedIdentity { .. }), + ..self + } + } + pub(crate) fn credential_source(&self) -> InputSource { self.credential_source } @@ -87,9 +97,34 @@ impl ValidatedAzureRequest { } } +pub(crate) enum CredentialAttempt { + Resolved(ResolvedCredential), + Unavailable(Error), + Failed(Error), +} + +impl CredentialAttempt { + pub(crate) fn into_result(self) -> Result { + match self { + Self::Resolved(credential) => Ok(credential), + Self::Unavailable(error) | Self::Failed(error) => Err(error), + } + } +} + pub(crate) struct NativeAzureTokenAcquirer { cache: AzureCredentialProviderCache, transport: Option, + imds_available: tokio::sync::Mutex, +} + +enum ImdsAvailability { + Unknown, + Available, + Unavailable { + retry_at: tokio::time::Instant, + error: Error, + }, } impl Default for NativeAzureTokenAcquirer { @@ -103,6 +138,7 @@ impl NativeAzureTokenAcquirer { Self { cache: AzureCredentialProviderCache::new(cache_capacity), transport: None, + imds_available: tokio::sync::Mutex::new(ImdsAvailability::Unknown), } } @@ -114,10 +150,51 @@ impl NativeAzureTokenAcquirer { Self { cache: AzureCredentialProviderCache::new(cache_capacity), transport: Some(transport), + imds_available: tokio::sync::Mutex::new(ImdsAvailability::Unknown), } } - pub(crate) async fn acquire( + pub(crate) async fn acquire(&self, request: ValidatedAzureRequest) -> CredentialAttempt { + let discover_imds = request.probe_imds; + if discover_imds && let Err(error) = self.discover_imds().await { + return CredentialAttempt::Unavailable(error); + } + match self.acquire_token(request).await { + Ok(credential) => CredentialAttempt::Resolved(credential), + Err(error) => { + let missing_identity = discover_imds && unassigned_identity(&error); + if missing_identity { + CredentialAttempt::Unavailable(error) + } else { + CredentialAttempt::Failed(error) + } + } + } + } + + async fn discover_imds(&self) -> Result<(), Error> { + let mut availability = self.imds_available.lock().await; + match &*availability { + ImdsAvailability::Available => return Ok(()), + ImdsAvailability::Unavailable { retry_at, error } + if tokio::time::Instant::now() < *retry_at => + { + return Err(error.clone()); + } + ImdsAvailability::Unknown | ImdsAvailability::Unavailable { .. } => {} + } + let result = probe_imds(self.transport.clone().unwrap_or_default()).await; + *availability = match &result { + Ok(()) => ImdsAvailability::Available, + Err(error) => ImdsAvailability::Unavailable { + retry_at: tokio::time::Instant::now() + Duration::from_secs(30), + error: error.clone(), + }, + }; + result + } + + async fn acquire_token( &self, request: ValidatedAzureRequest, ) -> Result { @@ -148,6 +225,54 @@ impl NativeAzureTokenAcquirer { } } +fn unassigned_identity(error: &(dyn std::error::Error + 'static)) -> bool { + std::iter::successors(Some(error), |cause| cause.source()).any(|cause| { + cause + .downcast_ref::() + .or_else(|| { + cause + .downcast_ref::()? + .downcast_ref::() + }) + .is_some_and(|error| { + error.http_status() == Some(azure_core::http::StatusCode::BadRequest) + }) + }) +} + +async fn probe_imds(transport: azure_core::http::Transport) -> Result<(), Error> { + use azure_core::http::{Context, Method, Request, Url}; + + let endpoint = + Url::parse("http://169.254.169.254/metadata/identity/oauth2/token?api-version=2019-08-01") + .map_err(|error| { + Error::InvalidConfiguration(ErrorDetail::failed("Azure IMDS probe endpoint", error)) + })?; + let mut request = Request::new(endpoint, Method::Get); + let response = tokio::time::timeout( + Duration::from_secs(1), + transport.send(&Context::new(), &mut request), + ) + .await + .map_err(|error| { + Error::CredentialAcquisition(ErrorDetail::failed("Azure IMDS availability probe", error)) + })? + .map_err(|error| { + Error::CredentialAcquisition(ErrorDetail::failed("Azure IMDS availability probe", error)) + })?; + // Headerless IMDS discovery expects 400, checked 2026-11-01 + // https://learn.microsoft.com/en-us/azure/virtual-machines/instance-metadata-service#errors-and-debugging + if response.status() != azure_core::http::StatusCode::BadRequest { + return Err(Error::CredentialAcquisition( + ErrorDetail::UnexpectedStatus { + operation: "Azure IMDS availability probe", + status: u16::from(response.status()), + }, + )); + } + Ok(()) +} + impl NativeAzureRequest { fn scope(&self) -> &str { match self { @@ -493,6 +618,7 @@ impl ClientAssertion for StaticAssertion { #[cfg(test)] mod tests { use std::sync::{Arc, Mutex}; + use std::time::Duration; use azure_core::http::headers::Headers; use azure_core::http::{AsyncRawResponse, HttpClient, Request, StatusCode, Transport}; @@ -569,6 +695,296 @@ mod tests { } } + #[derive(Debug)] + struct AvailableImdsClient { + probe_status: StatusCode, + probe_delay: Duration, + recover_after: Option, + response_status: Option, + probes: std::sync::atomic::AtomicUsize, + } + + impl Default for AvailableImdsClient { + fn default() -> Self { + Self { + probe_status: StatusCode::BadRequest, + probe_delay: Duration::ZERO, + recover_after: None, + response_status: None, + probes: std::sync::atomic::AtomicUsize::new(0), + } + } + } + + impl HttpClient for AvailableImdsClient { + fn execute_request<'life0, 'life1, 'async_trait>( + &'life0 self, + request: &'life1 Request, + ) -> std::pin::Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + if request + .headers() + .get_optional_str(&azure_core::http::headers::HeaderName::from_static( + "metadata", + )) + .is_none() + { + let probe = self + .probes + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + if self.recover_after.is_some_and(|failures| probe >= failures) { + return Ok(AsyncRawResponse::from_bytes( + StatusCode::BadRequest, + Headers::new(), + "metadata header required", + )); + } + tokio::time::sleep(self.probe_delay).await; + return Ok(AsyncRawResponse::from_bytes( + self.probe_status, + Headers::new(), + "metadata header required", + )); + } + if let Some(status) = self.response_status { + return Ok(AsyncRawResponse::from_bytes( + status, + Headers::new(), + "identity endpoint failure", + )); + } + tokio::time::sleep(std::time::Duration::from_secs(2)).await; + let expires = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs() + + 3600; + Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + Headers::new(), + serde_json::json!({ + "access_token": "managed-token", "expires_on": expires.to_string(), + "token_type": "Bearer", "resource": "https://service.test", + }) + .to_string(), + )) + }) + } + } + + #[rstest::rstest] + #[tokio::test(start_paused = true)] + async fn available_imds_retains_full_sdk_token_acquisition() { + let acquirer = NativeAzureTokenAcquirer::with_transport( + 4, + Transport::new(Arc::new(AvailableImdsClient { + probe_status: StatusCode::BadRequest, + response_status: None, + ..Default::default() + })), + ); + let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(true); + let credential = acquirer.acquire(request).await.into_result().unwrap(); + assert_eq!(credential.secret().expose(), "managed-token"); + } + + #[rstest::rstest] + #[case::another_metadata_service(StatusCode::NotFound)] + #[case::token_required(StatusCode::Unauthorized)] + #[case::proxy_rejection(StatusCode::Forbidden)] + #[case::unavailable(StatusCode::ServiceUnavailable)] + #[tokio::test(start_paused = true)] + async fn unrelated_metadata_responses_do_not_start_managed_identity_acquisition( + #[case] status: StatusCode, + ) { + let client = Arc::new(AvailableImdsClient { + probe_status: status, + ..Default::default() + }); + let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone())); + let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(true); + assert!(matches!( + acquirer.acquire(request.clone()).await, + super::CredentialAttempt::Unavailable(_) + )); + assert!(matches!( + acquirer.acquire(request).await, + super::CredentialAttempt::Unavailable(_) + )); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1); + } + + #[rstest::rstest] + #[tokio::test(start_paused = true)] + async fn an_imds_timeout_is_not_repeated_on_the_next_request() { + let client = Arc::new(AvailableImdsClient { + probe_delay: Duration::from_secs(10), + ..Default::default() + }); + let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone())); + let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(true); + assert!(matches!( + acquirer.acquire(request.clone()).await, + super::CredentialAttempt::Unavailable(_) + )); + let second = tokio::time::timeout(Duration::from_millis(10), acquirer.acquire(request)) + .await + .unwrap(); + assert!(matches!(second, super::CredentialAttempt::Unavailable(_))); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1); + } + + #[rstest::fixture] + fn imds_request() -> ValidatedAzureRequest { + ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(true) + } + + #[rstest::rstest] + #[case::timeout(StatusCode::BadRequest, Duration::from_secs(10))] + #[case::service_unavailable(StatusCode::ServiceUnavailable, Duration::ZERO)] + #[tokio::test(start_paused = true)] + async fn failed_imds_discovery_recovers_without_restarting_the_acquirer( + imds_request: ValidatedAzureRequest, + #[case] status: StatusCode, + #[case] delay: Duration, + ) { + let client = Arc::new(AvailableImdsClient { + probe_status: status, + probe_delay: delay, + recover_after: Some(1), + ..Default::default() + }); + let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone())); + let (first, second) = tokio::join!( + acquirer.acquire(imds_request.clone()), + acquirer.acquire(imds_request.clone()), + ); + assert!(matches!(first, super::CredentialAttempt::Unavailable(_))); + assert!(matches!(second, super::CredentialAttempt::Unavailable(_))); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1); + + tokio::time::advance(Duration::from_secs(30)).await; + let (first, second) = tokio::join!( + acquirer.acquire(imds_request.clone()), + acquirer.acquire(imds_request.clone()), + ); + assert_eq!( + first.into_result().unwrap().secret().expose(), + "managed-token" + ); + assert_eq!( + second.into_result().unwrap().secret().expose(), + "managed-token" + ); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2); + + tokio::time::advance(Duration::from_secs(60)).await; + assert_eq!( + acquirer + .acquire(imds_request) + .await + .into_result() + .unwrap() + .secret() + .expose(), + "managed-token", + ); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2); + } + + #[rstest::rstest] + #[tokio::test(start_paused = true)] + async fn a_persistent_imds_outage_reuses_each_failed_probe( + imds_request: ValidatedAzureRequest, + ) { + let client = Arc::new(AvailableImdsClient { + probe_delay: Duration::from_secs(10), + ..Default::default() + }); + let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone())); + assert!(matches!( + acquirer.acquire(imds_request.clone()).await, + super::CredentialAttempt::Unavailable(_), + )); + tokio::time::advance(Duration::from_secs(30)).await; + let (first, second) = tokio::join!( + acquirer.acquire(imds_request.clone()), + acquirer.acquire(imds_request), + ); + assert!(matches!(first, super::CredentialAttempt::Unavailable(_))); + assert!(matches!(second, super::CredentialAttempt::Unavailable(_))); + assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2); + } + + #[rstest::rstest] + #[case::default_identity_not_assigned(true, StatusCode::BadRequest, true)] + #[case::explicit_identity_not_assigned(false, StatusCode::BadRequest, false)] + #[case::default_identity_forbidden(true, StatusCode::Forbidden, false)] + #[case::default_identity_service_failure(true, StatusCode::InternalServerError, false)] + #[tokio::test(start_paused = true)] + async fn only_unavailable_default_identities_allow_fallback( + #[case] discovery: bool, + #[case] status: StatusCode, + #[case] unavailable: bool, + ) { + let acquirer = NativeAzureTokenAcquirer::with_transport( + 4, + Transport::new(Arc::new(AvailableImdsClient { + probe_status: StatusCode::BadRequest, + response_status: Some(status), + ..Default::default() + })), + ); + let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope: deployment("https://service.test/.default".into()), + selection_source: InputSource::Deployment, + }) + .unwrap() + .with_imds_probe(discovery); + let outcome = acquirer.acquire(request).await; + let is_unavailable = matches!(outcome, super::CredentialAttempt::Unavailable(_)); + let error = outcome.into_result().unwrap_err(); + assert_eq!( + is_unavailable, + unavailable, + "{}", + litellm_error::Report::new(&error) + ); + } + + #[rstest::rstest] #[tokio::test] async fn client_secret_uses_sdk_protocol_and_reuses_cached_credential() { let transport = Arc::new(RecordingTokenClient::default()); @@ -582,8 +998,12 @@ mod tests { "https://login.test", ); - let first = acquirer.acquire(request.clone()).await.unwrap(); - let second = acquirer.acquire(request).await.unwrap(); + let first = acquirer + .acquire(request.clone()) + .await + .into_result() + .unwrap(); + let second = acquirer.acquire(request).await.into_result().unwrap(); assert_eq!(first.secret().expose(), "native-token"); assert_eq!(second.secret().expose(), "native-token"); @@ -599,6 +1019,7 @@ mod tests { ); } + #[rstest::rstest] #[tokio::test] async fn credential_provider_cache_isolates_every_client_secret_identity_field() { let transport = Arc::new(RecordingTokenClient::default()); @@ -645,10 +1066,10 @@ mod tests { ), ]; - acquirer.acquire(base.clone()).await.unwrap(); - acquirer.acquire(base).await.unwrap(); + acquirer.acquire(base.clone()).await.into_result().unwrap(); + acquirer.acquire(base).await.into_result().unwrap(); for request in variants.into_iter().skip(1) { - acquirer.acquire(request).await.unwrap(); + acquirer.acquire(request).await.into_result().unwrap(); } assert_eq!(transport.requests.lock().unwrap().len(), 6); diff --git a/litellm-rust/crates/auth-azure/src/resolve.rs b/litellm-rust/crates/auth-azure/src/resolve.rs index a92630af26a..8e52bed5941 100644 --- a/litellm-rust/crates/auth-azure/src/resolve.rs +++ b/litellm-rust/crates/auth-azure/src/resolve.rs @@ -1,3 +1,4 @@ +use crate::native::CredentialAttempt; use litellm_auth_types::{ CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, Sourced, TokenProviderHandle, @@ -57,14 +58,14 @@ trait AzureTokenAcquirer: Send + Sync { fn acquire( &self, request: ValidatedAzureRequest, - ) -> Pin> + Send + '_>>; + ) -> Pin + Send + '_>>; } impl AzureTokenAcquirer for NativeAzureTokenAcquirer { fn acquire( &self, request: ValidatedAzureRequest, - ) -> Pin> + Send + '_>> { + ) -> Pin + Send + '_>> { Box::pin(NativeAzureTokenAcquirer::acquire(self, request)) } } @@ -123,6 +124,7 @@ impl AzureAuthService { self.native .acquire(request) .await + .into_result() .map(|credential| Sourced::new(credential, source)) .map(Some) } @@ -131,6 +133,7 @@ impl AzureAuthService { self.native .acquire(request) .await + .into_result() .map(|credential| Some(Sourced::new(credential, source))) } AzureCredentialPlan::Chain(requests) => { @@ -138,8 +141,16 @@ impl AzureAuthService { for request in requests { let source = request.credential_source(); match self.native.acquire(request).await { - Ok(credential) => return Ok(Some(Sourced::new(credential, source))), - Err(error) => failures.push(error), + CredentialAttempt::Resolved(credential) => { + return Ok(Some(Sourced::new(credential, source))); + } + CredentialAttempt::Unavailable(error) => failures.push(error), + CredentialAttempt::Failed(error) => { + return Err(Error::CredentialChain { + last: Box::new(error), + earlier: failures, + }); + } } } match failures.pop() { @@ -237,7 +248,12 @@ pub(crate) fn select_auth_plan( return Ok(AzureCredentialPlan::Missing); } - select_native_plan( + let probe_imds = selector + .as_ref() + .is_none_or(|selector| *selector.value() == AzureCredentialType::DefaultAzureCredential) + && env_lookup("IDENTITY_ENDPOINT").is_none() + && env_lookup("MSI_ENDPOINT").is_none(); + match select_native_plan( selector, tenant_id, client_id, @@ -245,7 +261,15 @@ pub(crate) fn select_auth_plan( scope, authority, inputs.enable_azure_ad_token_refresh.source(), - ) + )? { + AzureCredentialPlan::Chain(requests) => Ok(AzureCredentialPlan::Chain( + requests + .into_iter() + .map(|request| request.with_imds_probe(probe_imds)) + .collect(), + )), + plan => Ok(plan), + } } fn select_native_plan( @@ -509,7 +533,10 @@ mod tests { #[derive(Debug)] struct FileResolver; + use crate::native::CredentialAttempt; + struct ChainAcquirer { + terminal_on: Option<&'static str>, requests: Mutex>, succeed_on: Option<&'static str>, } @@ -518,20 +545,25 @@ mod tests { fn acquire( &self, request: ValidatedAzureRequest, - ) -> std::pin::Pin> + Send + '_>> - { + ) -> std::pin::Pin + Send + '_>> { let kind = request.kind(); self.requests.lock().unwrap().push(kind); Box::pin(async move { if self.succeed_on == Some(kind) { - Ok(ResolvedCredential::AccessToken { + CredentialAttempt::Resolved(ResolvedCredential::AccessToken { token: SecretValue::new("chain-token"), expires_on: None, }) } else { - Err(Error::CredentialAcquisition( - litellm_auth_types::ErrorDetail::Unresolved { subject: kind }, - )) + let error = + Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::Unresolved { + subject: kind, + }); + if self.terminal_on == Some(kind) { + CredentialAttempt::Failed(error) + } else { + CredentialAttempt::Unavailable(error) + } } }) } @@ -690,9 +722,97 @@ mod tests { assert_eq!(resolved, Some(SecretValue::new("rotated-assertion"))); } + #[derive(Debug)] + struct UnavailableImdsClient; + + impl azure_core::http::HttpClient for UnavailableImdsClient { + fn execute_request<'life0, 'life1, 'async_trait>( + &'life0 self, + _request: &'life1 azure_core::http::Request, + ) -> std::pin::Pin< + Box< + dyn Future> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(std::future::pending()) + } + } + + struct OfflineChainAcquirer(super::NativeAzureTokenAcquirer); + + impl AzureTokenAcquirer for OfflineChainAcquirer { + fn acquire( + &self, + request: ValidatedAzureRequest, + ) -> std::pin::Pin + Send + '_>> { + Box::pin(async move { + if request.kind() == "developer-tools" { + return CredentialAttempt::Resolved(ResolvedCredential::AccessToken { + token: SecretValue::new("cli-token"), + expires_on: None, + }); + } + self.0.acquire(request).await + }) + } + } + + #[rstest::rstest] + #[case::default_off_azure("DefaultAzureCredential", None, true)] + #[case::explicit_managed_identity("ManagedIdentityCredential", None, false)] + #[case::app_service("DefaultAzureCredential", Some("http://app-service.test"), false)] + #[tokio::test(start_paused = true)] + async fn only_default_imds_discovery_is_bounded( + #[case] selector: &str, + #[case] identity_endpoint: Option<&str>, + #[case] expect_cli: bool, + ) { + let native = super::NativeAzureTokenAcquirer::with_transport( + 4, + azure_core::http::Transport::new(Arc::new(UnavailableImdsClient)), + ); + let service = AzureAuthService::with_acquirer(Arc::new(OfflineChainAcquirer(native))); + let inputs = AzureAuthInputs { + azure_credential: crate::types::ConfigValue::Value(Sourced::new( + selector.into(), + InputSource::Deployment, + )), + enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment), + ..Default::default() + }; + let environment = |name: &str| { + (name == "IDENTITY_ENDPOINT") + .then_some(identity_endpoint) + .flatten() + .map(str::to_string) + }; + let outcome = tokio::time::timeout( + std::time::Duration::from_secs(2), + service.get_azure_ad_token(&inputs, &environment), + ) + .await; + + assert_eq!(outcome.is_ok(), expect_cli); + if expect_cli { + assert_eq!( + outcome.unwrap().unwrap().unwrap().value().secret().expose(), + "cli-token" + ); + } + } + + #[rstest::rstest] #[tokio::test] async fn default_chain_uses_declared_order_and_stops_after_success() { let acquirer = Arc::new(ChainAcquirer { + terminal_on: None, requests: Mutex::new(Vec::new()), succeed_on: Some("developer-tools"), }); @@ -715,9 +835,13 @@ mod tests { ); } + #[rstest::rstest] + #[case::all_unavailable(false)] + #[case::terminal_failure_after_unavailable(true)] #[tokio::test] - async fn chain_reports_each_acquisition_failure() { + async fn chain_reports_each_acquisition_failure(#[case] terminal: bool) { let acquirer = Arc::new(ChainAcquirer { + terminal_on: terminal.then_some("developer-tools"), requests: Mutex::new(Vec::new()), succeed_on: None, }); @@ -744,6 +868,27 @@ mod tests { ); } + #[rstest::rstest] + #[tokio::test] + async fn a_configured_identity_failure_does_not_fall_back_to_a_developer_identity() { + let acquirer = Arc::new(ChainAcquirer { + terminal_on: Some("managed-identity"), + requests: Mutex::new(Vec::new()), + succeed_on: Some("developer-tools"), + }); + let service = AzureAuthService::with_acquirer(acquirer.clone()); + let inputs = AzureAuthInputs { + enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment), + ..Default::default() + }; + let error = service + .get_azure_ad_token(&inputs, &|_| None) + .await + .unwrap_err(); + assert_eq!(*acquirer.requests.lock().unwrap(), ["managed-identity"]); + assert!(matches!(error, Error::CredentialChain { earlier, .. } if earlier.is_empty())); + } + #[derive(Debug)] struct CallerToken(&'static str); diff --git a/litellm-rust/crates/auth-gcp/Cargo.toml b/litellm-rust/crates/auth-gcp/Cargo.toml index ad60f0db817..90aab8d281c 100644 --- a/litellm-rust/crates/auth-gcp/Cargo.toml +++ b/litellm-rust/crates/auth-gcp/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [features] -google-sdk = ["dep:google-cloud-auth", "dep:http"] +google-sdk = ["dep:http"] [dependencies] litellm-secret-value.workspace = true @@ -21,10 +21,14 @@ tokio.workspace = true veil.workspace = true gcp_auth = "0.12.7" -google-cloud-auth = { workspace = true, optional = true } +google-cloud-auth.workspace = true http = { workspace = true, optional = true } [dev-dependencies] +time = { workspace = true, features = ["formatting"] } +axum = { workspace = true, features = ["json"] } litellm-error.workspace = true rstest.workspace = true +rustls.workspace = true tempfile.workspace = true +tokio = { workspace = true, features = ["test-util"] } diff --git a/litellm-rust/crates/auth-gcp/src/auth.rs b/litellm-rust/crates/auth-gcp/src/auth.rs index 650d5831796..5976ef82344 100644 --- a/litellm-rust/crates/auth-gcp/src/auth.rs +++ b/litellm-rust/crates/auth-gcp/src/auth.rs @@ -1,20 +1,11 @@ -use std::{future::Future, pin::Pin, sync::Arc}; +use std::sync::Arc; -use gcp_auth::{CustomServiceAccount, TokenProvider}; -use litellm_auth_types::{CredentialPlacement, Error, ErrorDetail, http::apply_credential}; -use moka::future::Cache; -use serde_json::Value; -use sha2::{Digest, Sha256}; +use litellm_auth_types::{CredentialPlacement, Error, http::apply_credential}; use crate::{ - config::{ - CredentialSource, VertexConfig, credential_source, get_vertex_ai_project, - is_base64_encoded_json, names_a_credentials_file, non_empty_env, - }, - constants::{ - CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT, VERTEX_AI_API_KEY_ENV, - VERTEXAI_API_KEY_ENV, - }, + config::{VertexConfig, credential_source, get_vertex_ai_project, non_empty_env}, + constants::{VERTEX_AI_API_KEY_ENV, VERTEXAI_API_KEY_ENV}, + identity::{GcpAuth, GcpProviderLoader, GcpTokenSource}, }; pub struct VertexEnvironment { @@ -22,28 +13,15 @@ pub struct VertexEnvironment { pub project_id: String, } -struct VertexAccessToken { - token: String, - project_id: String, -} - -#[derive(Clone)] +#[derive(Clone, Default)] pub struct VertexAuth { - providers: Cache>, - loader: Arc, -} - -impl Default for VertexAuth { - fn default() -> Self { - Self::new(Arc::new(GcpProviderLoader)) - } + identity: GcpAuth, } impl VertexAuth { - pub fn new(loader: Arc) -> Self { + pub fn new(loader: Arc) -> Self { Self { - providers: Cache::builder().max_capacity(64).build(), - loader, + identity: GcpAuth::new(loader), } } @@ -52,7 +30,10 @@ impl VertexAuth { config: &VertexConfig, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result { - self.load_provider(config, env_lookup).await?.token().await + self.identity + .access_token(credential_source(config, env_lookup)) + .await + .map(|token| token.expose().to_owned()) } pub async fn validate_environment( @@ -77,10 +58,18 @@ impl VertexAuth { (true, _) => None, (false, Some(token)) => Some(token), (false, None) => { - let access = self.get_access_token(config, env_lookup).await?; + let provider = self.load_provider(config, env_lookup).await?; + let (token, project_id) = match project_id { + Some(project_id) => (provider.token().await?, project_id), + None => tokio::try_join!(provider.token(), provider.project_id())?, + }; return Ok(VertexEnvironment { - headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?, - project_id: project_id.unwrap_or(access.project_id), + headers: apply_credential( + headers, + token.expose(), + CredentialPlacement::Bearer, + )?, + project_id, }); } }; @@ -104,248 +93,171 @@ impl VertexAuth { }) } - async fn get_access_token( - &self, - config: &VertexConfig, - env_lookup: &(dyn Fn(&str) -> Option + Sync), - ) -> Result { - let provider = self.load_provider(config, env_lookup).await?; - let (token, project_id) = tokio::try_join!(provider.token(), provider.project_id())?; - Ok(VertexAccessToken { token, project_id }) - } - async fn load_provider( &self, config: &VertexConfig, env_lookup: &(dyn Fn(&str) -> Option + Sync), - ) -> Result, Error> { - let source = credential_source(config, env_lookup); - let key = source.cache_key(); - self.providers - .try_get_with(key, self.loader.load(source)) + ) -> Result, Error> { + self.identity + .load_provider(credential_source(config, env_lookup)) .await - .map_err(|error| (*error).clone()) } } -pub trait VertexTokenSource: Send + Sync { - fn project_id(&self) -> VertexAuthFuture<'_, String>; - fn token(&self) -> VertexAuthFuture<'_, String>; -} - -pub trait VertexProviderLoader: Send + Sync { - fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc>; -} - -pub type VertexAuthFuture<'a, T> = Pin> + Send + 'a>>; - -struct GcpTokenSource(Arc); - -impl VertexTokenSource for GcpTokenSource { - fn project_id(&self) -> VertexAuthFuture<'_, String> { - Box::pin(async move { - self.0 - .project_id() - .await - .map(|project| project.to_string()) - .map_err(auth_acquisition_error) - }) - } - - fn token(&self) -> VertexAuthFuture<'_, String> { - Box::pin(async move { - self.0 - .token(&[CLOUD_PLATFORM_SCOPE]) - .await - .map(|token| token.as_str().to_string()) - .map_err(auth_acquisition_error) - }) - } -} - -struct GcpProviderLoader; - -impl VertexProviderLoader for GcpProviderLoader { - fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc> { - Box::pin(async move { - let provider: Arc = match source { - CredentialSource::Inline(configured) => Arc::new( - CustomServiceAccount::from_json(validate_request_credentials( - configured.expose(), - )?) - .map_err(auth_acquisition_error)?, - ), - CredentialSource::Trusted(configured) => { - Arc::new(load_trusted_service_account(configured.expose())?) - } - CredentialSource::ApplicationCredentials(path) => { - Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?) - } - CredentialSource::Adc => { - gcp_auth::provider().await.map_err(auth_acquisition_error)? - } - }; - Ok(Arc::new(GcpTokenSource(provider)) as Arc) - }) - } -} - -fn validate_request_credentials(configured: &str) -> Result<&str, Error> { - let token_uri = serde_json::from_str::(configured) - .ok() - .and_then(|credentials| { - credentials - .get("token_uri") - .and_then(Value::as_str) - .map(str::to_string) - }); - if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) { - return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType { - field: "request-controlled vertex_credentials token_uri".into(), - expected: "the canonical Google OAuth token endpoint", - })); - } - Ok(configured) -} - -impl CredentialSource { - fn cache_key(&self) -> CredentialCacheKey { - match self { - Self::Inline(configured) => { - CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into()) - } - Self::Trusted(configured) => { - CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into()) - } - Self::ApplicationCredentials(path) => { - CredentialCacheKey::ApplicationCredentials(path.clone()) - } - Self::Adc => CredentialCacheKey::Adc, - } - } -} - -#[derive(Clone, Debug, Hash, PartialEq, Eq)] -enum CredentialCacheKey { - Inline([u8; 32]), - Trusted([u8; 32]), - ApplicationCredentials(String), - Adc, -} - -fn load_trusted_service_account(configured: &str) -> Result { - if is_base64_encoded_json(configured) { - return Err(Error::InvalidConfiguration( - ErrorDetail::Base64EncodedJson { - field: "vertex_credentials", - }, - )); - } - if names_a_credentials_file(configured) { - CustomServiceAccount::from_file(configured) - } else { - CustomServiceAccount::from_json(configured) - } - .map_err(auth_acquisition_error) -} - -fn auth_acquisition_error(error: gcp_auth::Error) -> Error { - Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed( - "Vertex AI credentials", - error, - )) -} - #[cfg(test)] mod tests { + use std::{ + collections::{BTreeMap, BTreeSet}, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + }; + + use crate::{ + CredentialSource, GcpAuthFuture, GcpProviderLoader, GcpTokenSource, VertexAuth, + VertexConfig, get_vertex_ai_location, secret_names, + }; use litellm_secret_value::SecretValue; + use rstest::rstest; + use serde_json::{Value, json}; - use super::*; - - #[rstest::rstest] - #[case::canonical_endpoint(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#, true)] - #[case::noncanonical_endpoint(r#"{"token_uri":"http://127.0.0.1/token"}"#, false)] - #[case::missing_endpoint("{}", false)] - fn request_credentials_require_canonical_token_endpoint( - #[case] credentials: &str, - #[case] accepted: bool, - ) { - assert_eq!(validate_request_credentials(credentials).is_ok(), accepted); + struct FakeProvider { + project: Option<&'static str>, + calls: Arc, } - #[rstest::rstest] - #[case::missing_file( - "/nonexistent/vertex-key.json", - "failed to open application credentials file" - )] - #[case::malformed_json( - " {\"type\": \"service_account\",", - "failed to deserialize ApplicationCredentials" - )] - #[case::quoted_json( - "'{\"type\": \"service_account\"}'", - "failed to deserialize ApplicationCredentials" - )] - #[tokio::test] - async fn trusted_credentials_are_read_as_a_file_unless_they_are_json( - #[case] configured: &str, - #[case] expected: &str, - ) { - let Err(error) = GcpProviderLoader - .load(CredentialSource::Trusted(SecretValue::new(configured))) - .await - else { - panic!("{configured:?} should not load"); - }; - let report = litellm_error::Report::new(&error).to_string(); - assert!(report.contains(expected), "{report}"); - } + impl GcpTokenSource for FakeProvider { + fn project_id(&self) -> GcpAuthFuture<'_, String> { + self.calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async move { + self.project.map(str::to_owned).ok_or( + litellm_auth_types::Error::CredentialAcquisition( + litellm_auth_types::ErrorDetail::Unresolved { + subject: "credential project", + }, + ), + ) + }) + } - #[rstest::rstest] - #[case::padded("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9")] - #[case::quoted_prefixed_unpadded("\"base64:eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9\"")] - #[case::url_safe_wrapped("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsICJr\nZXkiOiAiPz8_In0=")] - #[tokio::test] - async fn base64_encoded_credentials_are_named_as_base64(#[case] configured: &str) { - assert!( - trusted_load_error(configured) - .await - .contains("is base64-encoded JSON, not the JSON itself") - ); - } - - async fn trusted_load_error(configured: &str) -> String { - match GcpProviderLoader - .load(CredentialSource::Trusted(SecretValue::new(configured))) - .await - { - Ok(_) => panic!("{configured:?} should not load"), - Err(error) => litellm_error::Report::new(&error).to_string(), + fn token(&self) -> GcpAuthFuture<'_, SecretValue> { + self.calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async { Ok(SecretValue::new("adc-token")) }) } } - #[rstest::rstest] - #[tokio::test] - async fn a_relative_key_file_named_like_json_is_read_as_a_file() { - let key = tempfile::Builder::new() - .prefix("{vertex") - .suffix(".json") - .tempfile_in(".") - .unwrap(); - std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap(); - let relative_name = key.path().file_name().unwrap().to_str().unwrap(); - let absolute_path = std::fs::canonicalize(key.path()).unwrap(); - assert_eq!( - trusted_load_error(relative_name).await, - trusted_load_error(absolute_path.to_str().unwrap()).await - ); + struct FakeLoader { + loads: Arc, + provider: Arc, } - #[rstest::rstest] - fn inline_and_trusted_credentials_never_share_a_cache_entry() { - assert_ne!( - CredentialSource::Inline(SecretValue::new("same-value")).cache_key(), - CredentialSource::Trusted(SecretValue::new("same-value")).cache_key() - ); + impl GcpProviderLoader for FakeLoader { + fn load(&self, _source: CredentialSource) -> GcpAuthFuture<'_, Arc> { + let loads = self.loads.clone(); + let provider = self.provider.clone(); + Box::pin(async move { + loads.fetch_add(1, Ordering::SeqCst); + Ok(provider) + }) + } + } + + fn config(value: Value) -> VertexConfig { + VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new()) + .unwrap() + } + + fn auth(calls: Arc, loads: Arc) -> VertexAuth { + let provider: Arc = Arc::new(FakeProvider { + calls, + project: Some("adc-project"), + }); + VertexAuth::new(Arc::new(FakeLoader { loads, provider })) + } + + #[rstest] + #[tokio::test] + async fn secret_names_cover_environment_reads() { + let seen = Arc::new(std::sync::Mutex::new(BTreeSet::::new())); + let recorded = seen.clone(); + let env = |name: &str| { + recorded.lock().unwrap().insert(name.to_string()); + None + }; + let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0))); + auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env) + .await + .unwrap(); + get_vertex_ai_location(&VertexConfig::default(), &env); + let allowed = secret_names() + .iter() + .map(|name| (*name).to_owned()) + .collect::>(); + assert!(seen.lock().unwrap().is_subset(&allowed)); + } + + #[rstest] + #[tokio::test] + async fn explicit_token_and_header_do_not_acquire_adc() { + let loads = Arc::new(AtomicUsize::new(0)); + let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone()); + let configured = config(json!({"vertex_project":"project-1"})); + let explicit = auth + .validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None) + .await + .unwrap(); + assert_eq!(explicit.headers[0].1, "Bearer access-token"); + let existing = auth + .validate_environment( + vec![("authorization".into(), "Bearer existing".into())], + None, + &configured, + &|_| None, + ) + .await + .unwrap(); + assert_eq!(existing.headers[0].1, "Bearer existing"); + assert_eq!(loads.load(Ordering::SeqCst), 0); + } + + #[rstest] + #[tokio::test] + async fn provider_is_reused_across_authentication_calls() { + let calls = Arc::new(AtomicUsize::new(0)); + let loads = Arc::new(AtomicUsize::new(0)); + let auth = auth(calls.clone(), loads.clone()); + let config = VertexConfig::default(); + let environment = |_: &str| None; + let (first, second) = tokio::try_join!( + auth.validate_environment(Vec::new(), None, &config, &environment), + auth.validate_environment(Vec::new(), None, &config, &environment), + ) + .unwrap(); + assert_eq!(first.project_id, "adc-project"); + assert_eq!(first.headers[0].1, "Bearer adc-token"); + assert_eq!(first.project_id, second.project_id); + assert_eq!(first.headers, second.headers); + assert_eq!(loads.load(Ordering::SeqCst), 1); + assert_eq!(calls.load(Ordering::SeqCst), 4); + } + + #[rstest] + #[tokio::test] + async fn an_explicit_project_does_not_require_project_metadata_in_the_credentials() { + let auth = VertexAuth::new(Arc::new(FakeLoader { + loads: Arc::new(AtomicUsize::new(0)), + provider: Arc::new(FakeProvider { + calls: Arc::new(AtomicUsize::new(0)), + project: None, + }), + })); + let configured = VertexConfig::new(None, Some("configured-project".into()), None); + let environment = auth + .validate_environment(Vec::new(), None, &configured, &|_| None) + .await + .unwrap(); + assert_eq!(environment.project_id, "configured-project"); + assert_eq!(environment.headers[0].1, "Bearer adc-token"); } } diff --git a/litellm-rust/crates/auth-gcp/src/config.rs b/litellm-rust/crates/auth-gcp/src/config.rs index 133e1c26903..048e40b0a2f 100644 --- a/litellm-rust/crates/auth-gcp/src/config.rs +++ b/litellm-rust/crates/auth-gcp/src/config.rs @@ -146,6 +146,20 @@ pub enum CredentialSource { Adc, } +impl CredentialSource { + pub fn from_vertex_environment(env_lookup: &dyn Fn(&str) -> Option) -> Self { + VertexParams::CREDENTIALS + .resolve(&|_| None, env_lookup) + .map(|configured| Self::Trusted(SecretValue::new(configured))) + .unwrap_or_else(|| Self::application_default(env_lookup)) + } + + pub fn application_default(env_lookup: &dyn Fn(&str) -> Option) -> Self { + non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV) + .map_or(Self::Adc, Self::ApplicationCredentials) + } +} + pub(crate) fn credential_source( config: &VertexConfig, env_lookup: &dyn Fn(&str) -> Option, @@ -158,12 +172,7 @@ pub(crate) fn credential_source( } }; } - if let Some(configured) = VertexParams::CREDENTIALS.resolve(&|_| None, env_lookup) { - return CredentialSource::Trusted(SecretValue::new(configured)); - } - non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV) - .map(CredentialSource::ApplicationCredentials) - .unwrap_or(CredentialSource::Adc) + CredentialSource::from_vertex_environment(env_lookup) } fn optional_credentials( @@ -319,6 +328,24 @@ mod tests { )); } + #[rstest::rstest] + fn application_default_ignores_vertex_credentials() { + let env = |name: &str| match name { + VERTEXAI_CREDENTIALS_ENV => Some("environment-json".to_string()), + GOOGLE_APPLICATION_CREDENTIALS_ENV => Some("adc.json".to_string()), + _ => None, + }; + assert!( + matches!(CredentialSource::application_default(&env), CredentialSource::ApplicationCredentials(path) if path == "adc.json") + ); + assert!(matches!( + CredentialSource::application_default( + &|name| (name == VERTEXAI_CREDENTIALS_ENV).then(|| "environment-json".into()) + ), + CredentialSource::Adc + )); + } + #[rstest::rstest] fn params_become_a_deployment_sourced_config() { let params = VertexParams { diff --git a/litellm-rust/crates/auth-gcp/src/identity.rs b/litellm-rust/crates/auth-gcp/src/identity.rs new file mode 100644 index 00000000000..5144221b218 --- /dev/null +++ b/litellm-rust/crates/auth-gcp/src/identity.rs @@ -0,0 +1,445 @@ +use crate::{ + config::{CredentialSource, is_base64_encoded_json, names_a_credentials_file}, + constants::{CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT}, +}; +use gcp_auth::{CustomServiceAccount, TokenProvider}; +use litellm_auth_types::{Error, ErrorDetail}; +use litellm_secret_value::SecretValue; +use moka::future::Cache; +use serde_json::Value; +use sha2::{Digest, Sha256}; +use std::{future::Future, pin::Pin, sync::Arc}; + +#[derive(Clone)] +pub struct GcpAuth { + providers: Cache>, + loader: Arc, +} + +impl Default for GcpAuth { + fn default() -> Self { + Self::new(Arc::new(NativeProviderLoader)) + } +} + +impl GcpAuth { + pub fn new(loader: Arc) -> Self { + Self { + providers: Cache::builder().max_capacity(64).build(), + loader, + } + } + + pub async fn access_token(&self, source: CredentialSource) -> Result { + self.load_provider(source).await?.token().await + } + + pub(crate) async fn load_provider( + &self, + source: CredentialSource, + ) -> Result, Error> { + self.providers + .try_get_with(source.cache_key(), self.loader.load(source)) + .await + .map_err(|error| (*error).clone()) + } +} + +pub trait GcpTokenSource: Send + Sync { + fn project_id(&self) -> GcpAuthFuture<'_, String>; + fn token(&self) -> GcpAuthFuture<'_, SecretValue>; +} + +pub trait GcpProviderLoader: Send + Sync { + fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc>; +} + +pub type GcpAuthFuture<'a, T> = Pin> + Send + 'a>>; + +struct ServiceAccountTokenSource(Arc); + +impl GcpTokenSource for ServiceAccountTokenSource { + fn project_id(&self) -> GcpAuthFuture<'_, String> { + Box::pin(async move { + self.0 + .project_id() + .await + .map(|project| project.to_string()) + .map_err(auth_acquisition_error) + }) + } + + fn token(&self) -> GcpAuthFuture<'_, SecretValue> { + Box::pin(async move { + self.0 + .token(&[CLOUD_PLATFORM_SCOPE]) + .await + .map(|token| SecretValue::new(token.as_str())) + .map_err(auth_acquisition_error) + }) + } +} + +struct SdkTokenSource { + credentials: google_cloud_auth::credentials::AccessTokenCredentials, + project_id: Option, +} + +impl GcpTokenSource for SdkTokenSource { + fn project_id(&self) -> GcpAuthFuture<'_, String> { + Box::pin(async move { + self.project_id + .clone() + .ok_or(Error::CredentialAcquisition(ErrorDetail::Unresolved { + subject: "GCP resource project", + })) + }) + } + + fn token(&self) -> GcpAuthFuture<'_, SecretValue> { + Box::pin(async move { + self.credentials + .access_token() + .await + .map(|token| SecretValue::new(token.token)) + .map_err(auth_acquisition_error) + }) + } +} + +async fn load_application_credentials(path: &str) -> Result, Error> { + let contents = tokio::fs::read_to_string(path) + .await + .map_err(auth_acquisition_error)?; + load_credentials_json(&contents) +} + +fn load_credentials_json(contents: &str) -> Result, Error> { + let configured: Value = serde_json::from_str(contents).map_err(auth_acquisition_error)?; + let kind = configured.get("type").and_then(Value::as_str); + let project_id = configured + .get("project_id") + .and_then(Value::as_str) + .map(str::to_string); + let credentials = match kind { + Some("authorized_user") => { + google_cloud_auth::credentials::user_account::Builder::new(configured) + .with_scopes([CLOUD_PLATFORM_SCOPE]) + .build_access_token_credentials() + } + Some("impersonated_service_account") => { + google_cloud_auth::credentials::impersonated::Builder::new(configured) + .with_scopes([CLOUD_PLATFORM_SCOPE]) + .build_access_token_credentials() + } + Some("external_account") => { + google_cloud_auth::credentials::external_account::Builder::new(configured) + .with_scopes([CLOUD_PLATFORM_SCOPE]) + .build_access_token_credentials() + } + Some("service_account") | None => { + let provider = + CustomServiceAccount::from_json(contents).map_err(auth_acquisition_error)?; + return Ok(Arc::new(ServiceAccountTokenSource(Arc::new(provider)))); + } + Some(_) => { + return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType { + field: "Google application credentials type".into(), + expected: "a supported Google credential type", + })); + } + } + .map_err(auth_acquisition_error)?; + Ok(Arc::new(SdkTokenSource { + credentials, + project_id, + })) +} + +struct NativeProviderLoader; + +impl GcpProviderLoader for NativeProviderLoader { + fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc> { + Box::pin(async move { + let provider: Arc = match source { + CredentialSource::Inline(configured) => Arc::new( + CustomServiceAccount::from_json(validate_request_credentials( + configured.expose(), + )?) + .map_err(auth_acquisition_error)?, + ), + CredentialSource::Trusted(configured) => { + return load_trusted_credentials(configured.expose()).await; + } + CredentialSource::ApplicationCredentials(path) => { + return load_application_credentials(&path).await; + } + CredentialSource::Adc => { + gcp_auth::provider().await.map_err(auth_acquisition_error)? + } + }; + Ok(Arc::new(ServiceAccountTokenSource(provider)) as Arc) + }) + } +} + +fn validate_request_credentials(configured: &str) -> Result<&str, Error> { + let token_uri = serde_json::from_str::(configured) + .ok() + .and_then(|credentials| { + credentials + .get("token_uri") + .and_then(Value::as_str) + .map(str::to_string) + }); + if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) { + return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType { + field: "request-controlled vertex_credentials token_uri".into(), + expected: "the canonical Google OAuth token endpoint", + })); + } + Ok(configured) +} + +impl CredentialSource { + fn cache_key(&self) -> CredentialCacheKey { + match self { + Self::Inline(configured) => { + CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into()) + } + Self::Trusted(configured) => { + CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into()) + } + Self::ApplicationCredentials(path) => { + CredentialCacheKey::ApplicationCredentials(path.clone()) + } + Self::Adc => CredentialCacheKey::Adc, + } + } +} + +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +enum CredentialCacheKey { + Inline([u8; 32]), + Trusted([u8; 32]), + ApplicationCredentials(String), + Adc, +} + +async fn load_trusted_credentials(configured: &str) -> Result, Error> { + if is_base64_encoded_json(configured) { + return Err(Error::InvalidConfiguration( + ErrorDetail::Base64EncodedJson { + field: "vertex_credentials", + }, + )); + } + if names_a_credentials_file(configured) { + load_application_credentials(configured).await + } else { + load_credentials_json(configured) + } +} + +fn auth_acquisition_error(error: impl std::error::Error + Send + Sync + 'static) -> Error { + Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed( + "Google credentials", + error, + )) +} + +#[cfg(test)] +mod tests { + use litellm_secret_value::SecretValue; + + use super::*; + + #[rstest::rstest] + #[case::authorized_user("authorized_user", 2)] + #[case::external_account("external_account", 2)] + #[case::impersonated_service_account("impersonated_service_account", 4)] + #[tokio::test] + async fn explicit_application_credentials_refresh_expired_tokens( + #[case] kind: &str, + #[case] expected_requests: usize, + #[values(false, true)] inline: bool, + ) { + use axum::{Json, Router, routing::post}; + use std::sync::atomic::{AtomicUsize, Ordering}; + + let _ = rustls::crypto::ring::default_provider().install_default(); + let requests = Arc::new(AtomicUsize::new(0)); + let recorded = requests.clone(); + let impersonations = requests.clone(); + let app = Router::new().route("/token", post(move || { + let recorded = recorded.clone(); + async move { + Json(serde_json::json!({ + "access_token": format!("token-{}", recorded.fetch_add(1, Ordering::SeqCst)), + "expires_in": 3600, + "token_type": "Bearer", + "issued_token_type": "urn:ietf:params:oauth:token-type:access_token", + })) + } + })); + let app = app.route("/v1/projects/-/serviceAccounts/test-principal:generateAccessToken", post(move || { + let recorded = impersonations.clone(); + async move { + let expiration = (time::OffsetDateTime::now_utc() + time::Duration::hours(1)).format(&time::format_description::well_known::Rfc3339).unwrap(); + Json(serde_json::json!({ + "accessToken": format!("impersonated-{}", recorded.fetch_add(1, Ordering::SeqCst)), + "expireTime": expiration, + })) + } + })); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let file = tempfile::NamedTempFile::new().unwrap(); + let subject_token = tempfile::NamedTempFile::new().unwrap(); + std::fs::write(subject_token.path(), "subject-token").unwrap(); + let configured = if kind == "external_account" { + serde_json::json!({ + "type": "external_account", "audience": "//iam.googleapis.com/projects/test/locations/global/workloadIdentityPools/test/providers/test", + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", "token_url": format!("http://{address}/token"), + "credential_source": { "file": subject_token.path().to_str().unwrap() }, "quota_project_id": "quota-project" + }) + } else { + let source = serde_json::json!({ + "type": "authorized_user", "client_id": "user-client", "client_secret": "user-secret", + "refresh_token": "user-refresh", "quota_project_id": "quota-project", "token_uri": format!("http://{address}/token") + }); + if kind == "impersonated_service_account" { + serde_json::json!({ + "type": kind, "source_credentials": source, "quota_project_id": "quota-project", + "service_account_impersonation_url": format!("http://{address}/v1/projects/-/serviceAccounts/test-principal:generateAccessToken") + }) + } else { + source + } + }; + std::fs::write(file.path(), configured.to_string()).unwrap(); + let source = if inline { + CredentialSource::Trusted(SecretValue::new(configured.to_string())) + } else { + CredentialSource::ApplicationCredentials(file.path().to_str().unwrap().into()) + }; + let provider = NativeProviderLoader.load(source).await.unwrap(); + + assert!(matches!( + provider.project_id().await, + Err(Error::CredentialAcquisition(ErrorDetail::Unresolved { + subject: "GCP resource project" + })) + )); + let original = provider.token().await.unwrap(); + tokio::time::pause(); + tokio::time::advance(std::time::Duration::from_secs(3601)).await; + tokio::time::resume(); + let refreshed = provider.token().await.unwrap(); + assert_ne!(original, refreshed); + assert_eq!(requests.load(Ordering::SeqCst), expected_requests); + server.abort(); + } + + #[rstest::rstest] + #[case::canonical_endpoint(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#, true)] + #[case::noncanonical_endpoint(r#"{"token_uri":"http://127.0.0.1/token"}"#, false)] + #[case::missing_endpoint("{}", false)] + fn request_credentials_require_canonical_token_endpoint( + #[case] credentials: &str, + #[case] accepted: bool, + ) { + assert_eq!(validate_request_credentials(credentials).is_ok(), accepted); + } + + #[rstest::rstest] + #[case::missing_file("/nonexistent/vertex-key.json", "Google credentials")] + #[case::malformed_json(" {\"type\": \"service_account\",", "Google credentials")] + #[case::quoted_json("'{\"type\": \"service_account\"}'", "Google credentials")] + #[tokio::test] + async fn trusted_credentials_are_read_as_a_file_unless_they_are_json( + #[case] configured: &str, + #[case] expected: &str, + ) { + let Err(error) = NativeProviderLoader + .load(CredentialSource::Trusted(SecretValue::new(configured))) + .await + else { + panic!("{configured:?} should not load"); + }; + let report = litellm_error::Report::new(&error).to_string(); + assert!(report.contains(expected), "{report}"); + } + + #[rstest::rstest] + #[case::padded("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9")] + #[case::quoted_prefixed_unpadded("\"base64:eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9\"")] + #[case::url_safe_wrapped("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsICJr\nZXkiOiAiPz8_In0=")] + #[tokio::test] + async fn base64_encoded_credentials_are_named_as_base64(#[case] configured: &str) { + assert!( + trusted_load_error(configured) + .await + .contains("is base64-encoded JSON, not the JSON itself") + ); + } + + async fn trusted_load_error(configured: &str) -> String { + match NativeProviderLoader + .load(CredentialSource::Trusted(SecretValue::new(configured))) + .await + { + Ok(_) => panic!("{configured:?} should not load"), + Err(error) => litellm_error::Report::new(&error).to_string(), + } + } + + #[rstest::rstest] + #[tokio::test] + async fn a_relative_key_file_named_like_json_is_read_as_a_file() { + let key = tempfile::Builder::new() + .prefix("{vertex") + .suffix(".json") + .tempfile_in(".") + .unwrap(); + std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap(); + let relative_name = key.path().file_name().unwrap().to_str().unwrap(); + let absolute_path = std::fs::canonicalize(key.path()).unwrap(); + assert_eq!( + trusted_load_error(relative_name).await, + trusted_load_error(absolute_path.to_str().unwrap()).await + ); + } + + #[rstest::rstest] + #[tokio::test] + async fn unknown_application_credential_types_are_rejected_without_service_account_fallback() { + let file = tempfile::NamedTempFile::new().unwrap(); + std::fs::write( + file.path(), + r#"{"type":"unsupported","private_key":"credential-value"}"#, + ) + .unwrap(); + let result = GcpAuth::default() + .access_token(CredentialSource::ApplicationCredentials( + file.path().to_str().unwrap().into(), + )) + .await; + let error = result.unwrap_err(); + assert!(matches!(error, Error::InvalidConfiguration(_))); + assert!( + !litellm_error::Report::new(&error) + .to_string() + .contains("credential-value") + ); + } + + #[rstest::rstest] + fn inline_and_trusted_credentials_never_share_a_cache_entry() { + assert_ne!( + CredentialSource::Inline(SecretValue::new("same-value")).cache_key(), + CredentialSource::Trusted(SecretValue::new("same-value")).cache_key() + ); + } +} diff --git a/litellm-rust/crates/auth-gcp/src/lib.rs b/litellm-rust/crates/auth-gcp/src/lib.rs index 2b389d45926..03c660e6eb8 100644 --- a/litellm-rust/crates/auth-gcp/src/lib.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -1,16 +1,16 @@ mod auth; mod config; pub mod constants; +mod identity; #[cfg(feature = "google-sdk")] mod sdk; -pub use auth::{ - VertexAuth, VertexAuthFuture, VertexEnvironment, VertexProviderLoader, VertexTokenSource, -}; +pub use auth::{VertexAuth, VertexEnvironment}; pub use config::{ CredentialSource, VertexConfig, get_vertex_ai_location, get_vertex_ai_project, get_vertex_ai_project_from_credentials, }; pub use constants::secret_names; +pub use identity::{GcpAuth, GcpAuthFuture, GcpProviderLoader, GcpTokenSource}; #[cfg(feature = "google-sdk")] pub use sdk::GoogleCredentials; diff --git a/litellm-rust/crates/auth-gcp/tests/vertex_auth.rs b/litellm-rust/crates/auth-gcp/tests/vertex_auth.rs deleted file mode 100644 index 5991749ab8a..00000000000 --- a/litellm-rust/crates/auth-gcp/tests/vertex_auth.rs +++ /dev/null @@ -1,120 +0,0 @@ -use std::{ - collections::{BTreeMap, BTreeSet}, - sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, - }, -}; - -use litellm_auth_gcp::{ - CredentialSource, VertexAuth, VertexAuthFuture, VertexConfig, VertexProviderLoader, - VertexTokenSource, get_vertex_ai_location, secret_names, -}; -use rstest::rstest; -use serde_json::{Value, json}; - -struct FakeProvider { - calls: Arc, -} - -impl VertexTokenSource for FakeProvider { - fn project_id(&self) -> VertexAuthFuture<'_, String> { - self.calls.fetch_add(1, Ordering::SeqCst); - Box::pin(async { Ok("adc-project".into()) }) - } - - fn token(&self) -> VertexAuthFuture<'_, String> { - self.calls.fetch_add(1, Ordering::SeqCst); - Box::pin(async { Ok("adc-token".into()) }) - } -} - -struct FakeLoader { - loads: Arc, - provider: Arc, -} - -impl VertexProviderLoader for FakeLoader { - fn load(&self, _source: CredentialSource) -> VertexAuthFuture<'_, Arc> { - let loads = self.loads.clone(); - let provider = self.provider.clone(); - Box::pin(async move { - loads.fetch_add(1, Ordering::SeqCst); - Ok(provider) - }) - } -} - -fn config(value: Value) -> VertexConfig { - VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new()) - .unwrap() -} - -fn auth(calls: Arc, loads: Arc) -> VertexAuth { - let provider: Arc = Arc::new(FakeProvider { calls }); - VertexAuth::new(Arc::new(FakeLoader { loads, provider })) -} - -#[rstest] -#[tokio::test] -async fn secret_names_cover_environment_reads() { - let seen = Arc::new(std::sync::Mutex::new(BTreeSet::::new())); - let recorded = seen.clone(); - let env = |name: &str| { - recorded.lock().unwrap().insert(name.to_string()); - None - }; - let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0))); - auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env) - .await - .unwrap(); - get_vertex_ai_location(&VertexConfig::default(), &env); - assert!( - seen.lock() - .unwrap() - .iter() - .all(|name| secret_names().contains(&name.as_str())) - ); -} - -#[rstest] -#[tokio::test] -async fn explicit_token_and_header_do_not_acquire_adc() { - let loads = Arc::new(AtomicUsize::new(0)); - let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone()); - let configured = config(json!({"vertex_project":"project-1"})); - let explicit = auth - .validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None) - .await - .unwrap(); - assert_eq!(explicit.headers[0].1, "Bearer access-token"); - let existing = auth - .validate_environment( - vec![("authorization".into(), "Bearer existing".into())], - None, - &configured, - &|_| None, - ) - .await - .unwrap(); - assert_eq!(existing.headers[0].1, "Bearer existing"); - assert_eq!(loads.load(Ordering::SeqCst), 0); -} - -#[rstest] -#[tokio::test] -async fn provider_is_reused_across_authentication_calls() { - let calls = Arc::new(AtomicUsize::new(0)); - let loads = Arc::new(AtomicUsize::new(0)); - let auth = auth(calls.clone(), loads.clone()); - for _ in 0..2 { - let environment = auth - .validate_environment(Vec::new(), None, &VertexConfig::default(), &|_| None) - .await - .unwrap(); - assert_eq!(environment.project_id, "adc-project"); - assert_eq!(environment.headers[0].1, "Bearer adc-token"); - } - assert_eq!(loads.load(Ordering::SeqCst), 1); - assert_eq!(calls.load(Ordering::SeqCst), 4); -} diff --git a/litellm-rust/crates/auth-types/src/error.rs b/litellm-rust/crates/auth-types/src/error.rs index a8ba7b64e03..a2e5629bf3f 100644 --- a/litellm-rust/crates/auth-types/src/error.rs +++ b/litellm-rust/crates/auth-types/src/error.rs @@ -92,6 +92,11 @@ impl Error { #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum ErrorDetail { + #[error("{operation} returned HTTP {status}")] + UnexpectedStatus { + operation: &'static str, + status: u16, + }, #[error("{field} must be {expected}")] InvalidType { field: String, diff --git a/litellm-rust/crates/auth-types/tests/aws_params.rs b/litellm-rust/crates/auth-types/tests/aws_params.rs index ed9a4f78ec2..b438984c199 100644 --- a/litellm-rust/crates/auth-types/tests/aws_params.rs +++ b/litellm-rust/crates/auth-types/tests/aws_params.rs @@ -3,21 +3,13 @@ use rstest::rstest; use serde_json::{Map, Value}; #[rstest] -fn fields_name_every_param_once_and_in_declaration_order() { +fn every_named_param_roundtrips_without_loss() { let filled: Map = AwsParams::fields() .map(|name| (name.to_string(), Value::from(format!("value-of-{name}")))) .collect(); let typed = AwsParams::from_optional_params(&filled); let serialized = serde_json::to_value(&typed).unwrap(); - assert_eq!( - serialized - .as_object() - .unwrap() - .keys() - .map(String::as_str) - .collect::>(), - AwsParams::fields().collect::>() - ); + assert_eq!(filled.len(), AwsParams::fields().count()); assert_eq!(serialized, Value::Object(filled.clone())); assert_eq!( serde_json::from_value::(Value::Object(filled)).unwrap(), diff --git a/litellm-rust/crates/cache-gcs/src/cache.rs b/litellm-rust/crates/cache-gcs/src/cache.rs index a6b08445bc8..b8c6720c234 100644 --- a/litellm-rust/crates/cache-gcs/src/cache.rs +++ b/litellm-rust/crates/cache-gcs/src/cache.rs @@ -8,7 +8,7 @@ use litellm_cache::{ use litellm_http::Client; use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode}; -use crate::{GcpTokenSource, TokenSource}; +use crate::{GcsTokenSource, TokenSource}; pub const DEFAULT_ENDPOINT: &str = "https://storage.googleapis.com"; @@ -62,7 +62,7 @@ pub struct GcsCache { impl GcsCache { pub fn new(config: GcsConfig, client: Client, codec: S) -> Self { - let token = Arc::new(GcpTokenSource::new(config.path_service_account.clone())); + let token = Arc::new(GcsTokenSource::new(config.path_service_account.clone())); Self::with_token_source(config, client, codec, token) } diff --git a/litellm-rust/crates/cache-gcs/src/lib.rs b/litellm-rust/crates/cache-gcs/src/lib.rs index cbb61cf0685..d40bca081a4 100644 --- a/litellm-rust/crates/cache-gcs/src/lib.rs +++ b/litellm-rust/crates/cache-gcs/src/lib.rs @@ -2,4 +2,4 @@ mod cache; mod token; pub use cache::{DEFAULT_ENDPOINT, GcsCache, GcsConfig, key_prefix}; -pub use token::{GcpTokenSource, StaticTokenSource, TokenSource}; +pub use token::{GcsTokenSource, StaticTokenSource, TokenSource}; diff --git a/litellm-rust/crates/cache-gcs/src/token.rs b/litellm-rust/crates/cache-gcs/src/token.rs index 852d8bbf035..4f639e7b865 100644 --- a/litellm-rust/crates/cache-gcs/src/token.rs +++ b/litellm-rust/crates/cache-gcs/src/token.rs @@ -1,7 +1,6 @@ use std::{future::Future, pin::Pin}; -use litellm_auth_gcp::{VertexAuth, VertexConfig}; -use litellm_auth_types::{InputSource, Sourced}; +use litellm_auth_gcp::{CredentialSource, GcpAuth}; use litellm_secret_value::SecretValue; use litellm_cache::Error; @@ -10,28 +9,30 @@ pub trait TokenSource: Send + Sync + 'static { fn bearer_token(&self) -> Pin> + Send + '_>>; } -pub struct GcpTokenSource { - auth: VertexAuth, - config: VertexConfig, +pub struct GcsTokenSource { + auth: GcpAuth, + credentials: Option, } -impl GcpTokenSource { +impl GcsTokenSource { pub fn new(path_service_account: Option) -> Self { - let credentials = path_service_account - .map(|path| Sourced::new(SecretValue::new(path), InputSource::Deployment)); Self { - auth: VertexAuth::default(), - config: VertexConfig::new(credentials, None, None), + auth: GcpAuth::default(), + credentials: path_service_account.filter(|path| !path.trim().is_empty()), } } } -impl TokenSource for GcpTokenSource { +impl TokenSource for GcsTokenSource { fn bearer_token(&self) -> Pin> + Send + '_>> { Box::pin(async move { self.auth - .access_token(&self.config, &|name| std::env::var(name).ok()) + .access_token(match &self.credentials { + Some(path) => CredentialSource::Trusted(SecretValue::new(path.clone())), + None => CredentialSource::application_default(&|name| std::env::var(name).ok()), + }) .await + .map(|token| token.expose().to_owned()) .map_err(Error::unavailable) }) } diff --git a/litellm-rust/crates/cache-s3/src/auth.rs b/litellm-rust/crates/cache-s3/src/auth.rs index 659390a5659..02707c24881 100644 --- a/litellm-rust/crates/cache-s3/src/auth.rs +++ b/litellm-rust/crates/cache-s3/src/auth.rs @@ -36,9 +36,12 @@ impl ProvideCredentials for S3Credentials { self.config.secret_access_key.clone(), ) { return Ok(AwsCredentials::new( - access_key_id, - secret_access_key, - self.config.session_token.clone(), + access_key_id.expose(), + secret_access_key.expose(), + self.config + .session_token + .as_ref() + .map(|value| value.expose().to_owned()), None, "litellm-s3-cache", )); diff --git a/litellm-rust/crates/cache-s3/tests/auth.rs b/litellm-rust/crates/cache-s3/tests/auth.rs index 59c5aa2e091..5aa251ac749 100644 --- a/litellm-rust/crates/cache-s3/tests/auth.rs +++ b/litellm-rust/crates/cache-s3/tests/auth.rs @@ -5,9 +5,9 @@ use rstest::rstest; fn explicit(session_token: Option<&str>) -> AwsAuthConfig { AwsAuthConfig { - access_key_id: Some("key".to_string()), - secret_access_key: Some("secret".to_string()), - session_token: session_token.map(str::to_string), + access_key_id: Some("key".into()), + secret_access_key: Some("secret".into()), + session_token: session_token.map(Into::into), region_name: Some("us-east-1".to_string()), ..Default::default() } diff --git a/litellm-rust/crates/cache-s3/tests/support/mod.rs b/litellm-rust/crates/cache-s3/tests/support/mod.rs index 7378dde21bf..311d2402e72 100644 --- a/litellm-rust/crates/cache-s3/tests/support/mod.rs +++ b/litellm-rust/crates/cache-s3/tests/support/mod.rs @@ -23,8 +23,8 @@ pub fn config(endpoint: &str) -> S3CacheConfig { url: endpoint.to_string(), }), auth: AwsAuthConfig { - access_key_id: Some("key".to_string()), - secret_access_key: Some("secret".to_string()), + access_key_id: Some("key".into()), + secret_access_key: Some("secret".into()), region_name: Some("us-east-1".to_string()), ..Default::default() }, diff --git a/litellm-rust/crates/inference-ocr/tests/resources.rs b/litellm-rust/crates/inference-ocr/tests/resources.rs index d618647e89b..3132a837eaf 100644 --- a/litellm-rust/crates/inference-ocr/tests/resources.rs +++ b/litellm-rust/crates/inference-ocr/tests/resources.rs @@ -7,7 +7,7 @@ use std::sync::{ use litellm_auth::AuthServices; use litellm_auth_gcp::{ - CredentialSource, VertexAuth, VertexAuthFuture, VertexProviderLoader, VertexTokenSource, + CredentialSource, GcpAuthFuture, GcpProviderLoader, GcpTokenSource, VertexAuth, }; use litellm_http::{HttpSettings, Resolution}; use litellm_inference::resources::CoreResources; @@ -20,28 +20,28 @@ use support::{ReceivedRequest, json_response, upstream}; struct TokenSource(String); -impl VertexTokenSource for TokenSource { - fn project_id(&self) -> VertexAuthFuture<'_, String> { +impl GcpTokenSource for TokenSource { + fn project_id(&self) -> GcpAuthFuture<'_, String> { Box::pin(async { Ok(self.0.clone()) }) } - fn token(&self) -> VertexAuthFuture<'_, String> { - Box::pin(async { Ok(self.0.clone()) }) + fn token(&self) -> GcpAuthFuture<'_, litellm_secret_value::SecretValue> { + Box::pin(async { Ok(litellm_secret_value::SecretValue::new(self.0.clone())) }) } } #[derive(Default)] struct Loader(AtomicUsize); -impl VertexProviderLoader for Loader { - fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc> { +impl GcpProviderLoader for Loader { + fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc> { Box::pin(async move { self.0.fetch_add(1, Ordering::SeqCst); let identity = match source { CredentialSource::Trusted(secret) => secret.expose().to_string(), other => panic!("unexpected credential source: {other:?}"), }; - Ok(Arc::new(TokenSource(identity)) as Arc) + Ok(Arc::new(TokenSource(identity)) as Arc) }) } } diff --git a/litellm-rust/crates/secret-value/src/secret_value.rs b/litellm-rust/crates/secret-value/src/secret_value.rs index 27e4227080d..33c559b80ec 100644 --- a/litellm-rust/crates/secret-value/src/secret_value.rs +++ b/litellm-rust/crates/secret-value/src/secret_value.rs @@ -16,6 +16,18 @@ impl SecretValue { } } +impl From for SecretValue { + fn from(value: String) -> Self { + Self::new(value) + } +} + +impl From<&str> for SecretValue { + fn from(value: &str) -> Self { + Self::new(value) + } +} + impl PartialEq for SecretValue { fn eq(&self, other: &Self) -> bool { subtle::ConstantTimeEq::ct_eq(self.0.as_bytes(), other.0.as_bytes()).into() diff --git a/litellm-rust/crates/secrets-aws/src/auth.rs b/litellm-rust/crates/secrets-aws/src/auth.rs index 66f4f51bda9..6bec8ad96a1 100644 --- a/litellm-rust/crates/secrets-aws/src/auth.rs +++ b/litellm-rust/crates/secrets-aws/src/auth.rs @@ -39,30 +39,15 @@ impl Credentials { Self { auth, config: AwsAuthConfig { - access_key_id: context - .access_key_id - .as_ref() - .map(|value| value.expose().to_owned()), - secret_access_key: context - .secret_access_key - .as_ref() - .map(|value| value.expose().to_owned()), - session_token: context - .session_token - .as_ref() - .map(|value| value.expose().to_owned()), + access_key_id: context.access_key_id.as_ref().cloned(), + secret_access_key: context.secret_access_key.as_ref().cloned(), + session_token: context.session_token.as_ref().cloned(), region_name: region(settings, environment.as_ref()).ok(), role_name: settings.aws_role_name.clone(), session_name: settings.aws_session_name.clone(), - external_id: settings - .aws_external_id - .as_ref() - .map(|v| v.expose().to_owned()), + external_id: settings.aws_external_id.as_ref().cloned(), profile_name: settings.aws_profile_name.clone(), - web_identity_token: settings - .aws_web_identity_token - .as_ref() - .map(|v| v.expose().to_owned()), + web_identity_token: settings.aws_web_identity_token.as_ref().cloned(), sts_endpoint: settings.aws_sts_endpoint.clone(), }, environment,