scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-02 07:13:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 49

fix(auth): align cloud credential loading and expiry with the Python SDKs

Browse source
This commit is contained in:
Devin AI 2026-11-02 00:17:12 +00:00
parent 60fc121ce8
commit 28cc4ee570
25 changed files with 2055 additions and 813 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm-rust/Cargo.lock generated
View file

@ -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",
]

8
litellm-rust/crates/auth-aws/Cargo.toml
View file

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

1073
litellm-rust/crates/auth-aws/src/aws.rs
View file

File diff suppressed because it is too large Load diff

2
litellm-rust/crates/auth-aws/src/error.rs
View file

@ -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,

3
litellm-rust/crates/auth-aws/src/signer.rs
View file

@ -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),

5
litellm-rust/crates/auth-azure/Cargo.toml
View file

@ -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"] }

433
litellm-rust/crates/auth-azure/src/native.rs
View file

@ -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<ResolvedCredential, Error> {
match self {
Self::Resolved(credential) => Ok(credential),
Self::Unavailable(error) | Self::Failed(error) => Err(error),
}
}
}
pub(crate) struct NativeAzureTokenAcquirer {
cache: AzureCredentialProviderCache,
transport: Option<azure_core::http::Transport>,
imds_available: tokio::sync::Mutex<ImdsAvailability>,
}
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<ResolvedCredential, Error> {
@ -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::<azure_core::Error>()
.or_else(|| {
cause
.downcast_ref::<litellm_error::Source>()?
.downcast_ref::<azure_core::Error>()
})
.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<usize>,
response_status: Option<StatusCode>,
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<dyn std::future::Future<Output = Result<AsyncRawResponse>> + 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);

171
litellm-rust/crates/auth-azure/src/resolve.rs
View file

@ -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<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>>;
) -> Pin<Box<dyn Future<Output = CredentialAttempt> + Send + '_>>;
}
impl AzureTokenAcquirer for NativeAzureTokenAcquirer {
fn acquire(
&self,
request: ValidatedAzureRequest,
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>> {
) -> Pin<Box<dyn Future<Output = CredentialAttempt> + 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<Vec<&'static str>>,
succeed_on: Option<&'static str>,
}
@ -518,20 +545,25 @@ mod tests {
fn acquire(
&self,
request: ValidatedAzureRequest,
) -> std::pin::Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>>
{
) -> std::pin::Pin<Box<dyn Future<Output = CredentialAttempt> + 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<Output = azure_core::Result<azure_core::http::AsyncRawResponse>>
+ 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<Box<dyn Future<Output = CredentialAttempt> + 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);

8
litellm-rust/crates/auth-gcp/Cargo.toml
View file

@ -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"] }

428
litellm-rust/crates/auth-gcp/src/auth.rs
View file

@ -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<CredentialCacheKey, Arc<dyn VertexTokenSource>>,
loader: Arc<dyn VertexProviderLoader>,
}
impl Default for VertexAuth {
fn default() -> Self {
Self::new(Arc::new(GcpProviderLoader))
}
identity: GcpAuth,
}
impl VertexAuth {
pub fn new(loader: Arc<dyn VertexProviderLoader>) -> Self {
pub fn new(loader: Arc<dyn GcpProviderLoader>) -> 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<String> + Sync),
) -> Result<String, Error> {
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<String> + Sync),
) -> Result<VertexAccessToken, Error> {
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<String> + Sync),
) -> Result<Arc<dyn VertexTokenSource>, Error> {
let source = credential_source(config, env_lookup);
let key = source.cache_key();
self.providers
.try_get_with(key, self.loader.load(source))
) -> Result<Arc<dyn GcpTokenSource>, 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<dyn VertexTokenSource>>;
}
pub type VertexAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
struct GcpTokenSource(Arc<dyn TokenProvider>);
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<dyn VertexTokenSource>> {
Box::pin(async move {
let provider: Arc<dyn TokenProvider> = 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<dyn VertexTokenSource>)
})
}
}
fn validate_request_credentials(configured: &str) -> Result<&str, Error> {
let token_uri = serde_json::from_str::<Value>(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<CustomServiceAccount, Error> {
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<AtomicUsize>,
}
#[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<AtomicUsize>,
provider: Arc<dyn GcpTokenSource>,
}
#[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<dyn GcpTokenSource>> {
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<AtomicUsize>, loads: Arc<AtomicUsize>) -> VertexAuth {
let provider: Arc<dyn GcpTokenSource> = 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::<String>::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::<BTreeSet<_>>();
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");
}
}

39
litellm-rust/crates/auth-gcp/src/config.rs
View file

@ -146,6 +146,20 @@ pub enum CredentialSource {
Adc,
}
impl CredentialSource {
pub fn from_vertex_environment(env_lookup: &dyn Fn(&str) -> Option<String>) -> 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<String>) -> 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<String>,
@ -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 {

445
litellm-rust/crates/auth-gcp/src/identity.rs Normal file
View file

@ -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<CredentialCacheKey, Arc<dyn GcpTokenSource>>,
loader: Arc<dyn GcpProviderLoader>,
}
impl Default for GcpAuth {
fn default() -> Self {
Self::new(Arc::new(NativeProviderLoader))
}
}
impl GcpAuth {
pub fn new(loader: Arc<dyn GcpProviderLoader>) -> Self {
Self {
providers: Cache::builder().max_capacity(64).build(),
loader,
}
}
pub async fn access_token(&self, source: CredentialSource) -> Result<SecretValue, Error> {
self.load_provider(source).await?.token().await
}
pub(crate) async fn load_provider(
&self,
source: CredentialSource,
) -> Result<Arc<dyn GcpTokenSource>, 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<dyn GcpTokenSource>>;
}
pub type GcpAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
struct ServiceAccountTokenSource(Arc<dyn TokenProvider>);
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<String>,
}
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<Arc<dyn GcpTokenSource>, 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<Arc<dyn GcpTokenSource>, 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<dyn GcpTokenSource>> {
Box::pin(async move {
let provider: Arc<dyn TokenProvider> = 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<dyn GcpTokenSource>)
})
}
}
fn validate_request_credentials(configured: &str) -> Result<&str, Error> {
let token_uri = serde_json::from_str::<Value>(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<Arc<dyn GcpTokenSource>, 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()
);
}
}

6
litellm-rust/crates/auth-gcp/src/lib.rs
View file

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

120
litellm-rust/crates/auth-gcp/tests/vertex_auth.rs
View file

@ -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<AtomicUsize>,
}
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<AtomicUsize>,
provider: Arc<dyn VertexTokenSource>,
}
impl VertexProviderLoader for FakeLoader {
fn load(&self, _source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>> {
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<AtomicUsize>, loads: Arc<AtomicUsize>) -> VertexAuth {
let provider: Arc<dyn VertexTokenSource> = 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::<String>::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);
}

5
litellm-rust/crates/auth-types/src/error.rs
View file

@ -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,

12
litellm-rust/crates/auth-types/tests/aws_params.rs
View file

@ -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<String, Value> = 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::<Vec<_>>(),
AwsParams::fields().collect::<Vec<_>>()
);
assert_eq!(filled.len(), AwsParams::fields().count());
assert_eq!(serialized, Value::Object(filled.clone()));
assert_eq!(
serde_json::from_value::<AwsParams>(Value::Object(filled)).unwrap(),

4
litellm-rust/crates/cache-gcs/src/cache.rs
View file

@ -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<S: CacheCodec> {
impl<S: CacheCodec> GcsCache<S> {
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)
}

2
litellm-rust/crates/cache-gcs/src/lib.rs
View file

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

25
litellm-rust/crates/cache-gcs/src/token.rs
View file

@ -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<Box<dyn Future<Output = Result<String, Error>> + Send + '_>>;
}
pub struct GcpTokenSource {
auth: VertexAuth,
config: VertexConfig,
pub struct GcsTokenSource {
auth: GcpAuth,
credentials: Option<String>,
}
impl GcpTokenSource {
impl GcsTokenSource {
pub fn new(path_service_account: Option<String>) -> 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<Box<dyn Future<Output = Result<String, Error>> + 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)
})
}

9
litellm-rust/crates/cache-s3/src/auth.rs
View file

@ -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",
));

6
litellm-rust/crates/cache-s3/tests/auth.rs
View file

@ -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()
}

4
litellm-rust/crates/cache-s3/tests/support/mod.rs
View file

@ -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()
},

16
litellm-rust/crates/inference-ocr/tests/resources.rs
View file

@ -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<dyn VertexTokenSource>> {
impl GcpProviderLoader for Loader {
fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc<dyn GcpTokenSource>> {
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<dyn VertexTokenSource>)
Ok(Arc::new(TokenSource(identity)) as Arc<dyn GcpTokenSource>)
})
}
}

12
litellm-rust/crates/secret-value/src/secret_value.rs
View file

@ -16,6 +16,18 @@ impl SecretValue {
}
}
impl From<String> 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()

25
litellm-rust/crates/secrets-aws/src/auth.rs
View file

@ -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,