scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129

test(rust): strengthen AWS auth contracts

Browse source
This commit is contained in:
Yujong Lee 2026-09-08 14:30:39 -07:00
parent 36eea7c356
commit 13820b17d3
13 changed files with 1267 additions and 796 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
litellm-rust/Cargo.lock generated
View file

@ -1434,7 +1434,9 @@ dependencies = [
"aws-smithy-runtime-api",
"aws-types",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"veil",
]
[[package]]
@ -2901,6 +2903,27 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "veil"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c7352f0bbf3ab98911b0c0277065094c1b1ec79bbc85fa3b7d16bf1859c3d96f"
dependencies = [
"once_cell",
"veil-macros",
]
[[package]]
name = "veil-macros"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "47a3f4f06d904eb789b935253752ba6bcc1dfa61349f8d5341c66abe070b44e5"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.119",
]
[[package]]
name = "version_check"
version = "0.9.5"

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

@ -13,7 +13,9 @@ aws-sigv4 = "1.5.1"
aws-smithy-runtime-api = "1.13.0"
aws-types = "1.4.0"
sha2.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["sync"] }
veil = "0.3"
[dev-dependencies]
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }

275
litellm-rust/crates/auth-aws/src/cache.rs Normal file
View file

@ -0,0 +1,275 @@
use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap};
use std::future::Future;
use std::hash::{Hash, Hasher};
use std::sync::Mutex;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use sha2::{Digest, Sha256};
use veil::Redact;
use crate::Credentials;
const CREDENTIAL_FETCH_LOCK_STRIPES: usize = 64;
struct CredentialCache {
entries: HashMap<CredentialScope, (Credentials, Duration)>,
expirations: BinaryHeap<Reverse<(Duration, CredentialScope)>>,
max_entries: usize,
}
impl Default for CredentialCache {
fn default() -> Self {
Self::new(200)
}
}
impl CredentialCache {
fn new(max_entries: usize) -> Self {
Self {
entries: HashMap::new(),
expirations: BinaryHeap::new(),
max_entries: max_entries.max(1),
}
}
fn get(&mut self, key: &CredentialScope, now: Duration) -> Option<Credentials> {
let (credentials, expiration) = self.entries.get(key)?;
if *expiration > now {
return Some(credentials.clone());
}
self.entries.remove(key);
None
}
fn insert(
&mut self,
key: CredentialScope,
credentials: Credentials,
ttl: Duration,
now: Duration,
) {
while let Some(Reverse((expiration, key))) = self.expirations.peek().cloned() {
if self.entries.get(&key).map(|(_, current)| *current) != Some(expiration) {
self.expirations.pop();
} else if expiration <= now || self.entries.len() >= self.max_entries {
self.expirations.pop();
self.entries.remove(&key);
} else {
break;
}
}
let expiration = now + ttl;
self.entries.insert(key.clone(), (credentials, expiration));
self.expirations.push(Reverse((expiration, key)));
}
}
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Redact)]
pub struct CredentialScope(#[redact(fixed = 8)] [u8; 32]);
impl CredentialScope {
pub fn from_optional_values<'a>(
namespace: &str,
values: impl IntoIterator<Item = Option<&'a str>>,
) -> Self {
let mut hasher = Sha256::new();
hasher.update(namespace.len().to_le_bytes());
hasher.update(namespace.as_bytes());
for value in values {
match value {
Some(value) => {
hasher.update([1]);
hasher.update(value.len().to_le_bytes());
hasher.update(value.as_bytes());
}
None => hasher.update([0]),
}
}
Self(hasher.finalize().into())
}
}
pub trait Clock: Send + Sync {
fn now(&self) -> Duration;
}
#[derive(Default)]
pub struct SystemClock;
impl Clock for SystemClock {
fn now(&self) -> Duration {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
}
}
pub struct CredentialState<R, C = SystemClock> {
runtime: R,
clock: C,
cache: Mutex<CredentialCache>,
fetch_locks: Box<[tokio::sync::Mutex<()>]>,
}
impl<R> CredentialState<R, SystemClock> {
pub fn new(runtime: R, max_entries: usize) -> Self {
Self::with_clock(runtime, max_entries, SystemClock)
}
}
impl<R, C> CredentialState<R, C>
where
C: Clock,
{
pub fn with_clock(runtime: R, max_entries: usize, clock: C) -> Self {
Self {
runtime,
clock,
cache: Mutex::new(CredentialCache::new(max_entries)),
fetch_locks: (0..CREDENTIAL_FETCH_LOCK_STRIPES)
.map(|_| tokio::sync::Mutex::new(()))
.collect(),
}
}
pub fn runtime(&self) -> &R {
&self.runtime
}
pub async fn get_or_acquire<E, F, Fut>(
&self,
scope: CredentialScope,
ttl: Duration,
acquire: F,
) -> Result<Credentials, E>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<Credentials, E>>,
{
let mut stripe_hasher = std::collections::hash_map::DefaultHasher::new();
scope.hash(&mut stripe_hasher);
let stripe = stripe_hasher.finish() as usize % self.fetch_locks.len();
let _fetch_guard = self.fetch_locks[stripe].lock().await;
if let Some(credentials) = self
.cache
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.get(&scope, self.clock.now())
{
return Ok(credentials);
}
let credentials = acquire().await?;
self.cache
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(scope, credentials.clone(), ttl, self.clock.now());
Ok(credentials)
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use super::*;
use crate::static_credentials;
#[test]
fn cache_expiry_and_bounds_are_clock_driven() {
let mut cache = CredentialCache::new(1);
let first = CredentialScope::from_optional_values("test", [Some("first")]);
let second = CredentialScope::from_optional_values("test", [Some("second")]);
cache.insert(
first.clone(),
static_credentials("ak1", "sk1"),
Duration::from_secs(10),
Duration::ZERO,
);
assert_eq!(
cache
.get(&first, Duration::from_secs(9))
.unwrap()
.access_key_id(),
"ak1"
);
cache.insert(
second.clone(),
static_credentials("ak2", "sk2"),
Duration::from_secs(10),
Duration::ZERO,
);
assert!(cache.get(&first, Duration::ZERO).is_none());
assert!(cache.get(&second, Duration::from_secs(11)).is_none());
}
#[tokio::test]
async fn credential_state_coordinates_concurrent_misses() {
let state = CredentialState::new((), 1);
let acquisitions = AtomicUsize::new(0);
let scope = CredentialScope::from_optional_values("test", [Some("identity")]);
let acquire = || async {
acquisitions.fetch_add(1, Ordering::SeqCst);
tokio::task::yield_now().await;
Ok::<_, ()>(static_credentials("ak", "sk"))
};
let (first, second) = tokio::join!(
state.get_or_acquire(scope.clone(), Duration::from_secs(10), acquire),
state.get_or_acquire(scope, Duration::from_secs(10), acquire),
);
assert_eq!(first.unwrap(), second.unwrap());
assert_eq!(acquisitions.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn credential_state_reacquires_at_the_injected_expiry_boundary() {
struct TestClock(Arc<AtomicU64>);
impl Clock for TestClock {
fn now(&self) -> Duration {
Duration::from_secs(self.0.load(Ordering::SeqCst))
}
}
let now = Arc::new(AtomicU64::new(0));
let state = CredentialState::with_clock((), 1, TestClock(now.clone()));
let acquisitions = AtomicUsize::new(0);
let scope = CredentialScope::from_optional_values("test", [Some("identity")]);
let acquire = || async {
acquisitions.fetch_add(1, Ordering::SeqCst);
Ok::<_, ()>(static_credentials("ak", "sk"))
};
state
.get_or_acquire(scope.clone(), Duration::from_secs(10), acquire)
.await
.unwrap();
now.store(9, Ordering::SeqCst);
state
.get_or_acquire(scope.clone(), Duration::from_secs(10), acquire)
.await
.unwrap();
now.store(10, Ordering::SeqCst);
state
.get_or_acquire(scope, Duration::from_secs(10), acquire)
.await
.unwrap();
assert_eq!(acquisitions.load(Ordering::SeqCst), 2);
}
#[test]
fn credential_scope_does_not_expose_key_material() {
let scope = CredentialScope::from_optional_values(
"test",
[Some("visible-id"), Some("never-print-secret"), None],
);
let debug = format!("{scope:?}");
assert!(!debug.contains("visible-id"));
assert!(!debug.contains("never-print-secret"));
}
}

84
litellm-rust/crates/auth-aws/src/credentials.rs Normal file
View file

@ -0,0 +1,84 @@
use std::time::SystemTime;
use aws_credential_types::Credentials as SdkCredentials;
use veil::Redact;
#[derive(Clone, Redact)]
pub struct Credentials(#[redact(fixed = 8)] pub(crate) SdkCredentials);
impl Credentials {
pub fn new(
access_key_id: impl Into<String>,
secret_access_key: impl Into<String>,
session_token: Option<String>,
expires_after: Option<SystemTime>,
provider_name: &'static str,
) -> Self {
Self(SdkCredentials::new(
access_key_id,
secret_access_key,
session_token,
expires_after,
provider_name,
))
}
pub fn access_key_id(&self) -> &str {
self.0.access_key_id()
}
pub fn session_token(&self) -> Option<&str> {
self.0.session_token()
}
}
impl PartialEq for Credentials {
fn eq(&self, other: &Self) -> bool {
self.0.access_key_id() == other.0.access_key_id()
&& self.0.secret_access_key() == other.0.secret_access_key()
&& self.0.session_token() == other.0.session_token()
}
}
impl Eq for Credentials {}
pub fn static_credentials(
access_key_id: impl Into<String>,
secret_access_key: impl Into<String>,
) -> Credentials {
Credentials::new(
access_key_id,
secret_access_key,
None,
None,
"litellm-static",
)
}
pub fn session_credentials(
access_key_id: impl Into<String>,
secret_access_key: impl Into<String>,
session_token: impl Into<String>,
provider_name: &'static str,
) -> Credentials {
Credentials::new(
access_key_id,
secret_access_key,
Some(session_token.into()),
None,
provider_name,
)
}
#[cfg(test)]
mod tests {
use super::static_credentials;
#[test]
fn credentials_are_redacted() {
let credentials = static_credentials("visible-id", "never-print-secret");
let debug = format!("{credentials:?}");
assert!(!debug.contains("visible-id"));
assert!(!debug.contains("never-print-secret"));
}
}

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

@ -0,0 +1,23 @@
use thiserror::Error as ThisError;
#[derive(Clone, Debug, PartialEq, Eq, ThisError)]
pub enum Error {
#[error("AWS profile credentials failed: {0}")]
ProfileCredentials(String),
#[error("AWS default credentials failed: {0}")]
DefaultCredentials(String),
#[error("AWS role credentials failed: {0}")]
RoleCredentials(String),
#[error("AWS web identity credentials failed: {0}")]
WebIdentityCredentials(String),
#[error("AWS web identity response had no credentials")]
MissingWebIdentityCredentials,
#[error("AWS web identity expiration was invalid: {0}")]
InvalidWebIdentityExpiration(String),
#[error("AWS signing parameters failed: {0}")]
SigningParameters(String),
#[error("AWS signable request failed: {0}")]
SignableRequest(String),
#[error("AWS request signing failed: {0}")]
RequestSigning(String),
}

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

@ -1,798 +1,17 @@
use std::cmp::Reverse;
use std::collections::{BTreeMap, BinaryHeap, HashMap};
use std::fmt;
use std::future::Future;
use std::hash::{Hash, Hasher};
use std::pin::Pin;
use std::sync::Mutex;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use aws_credential_types::Credentials as SdkCredentials;
use aws_credential_types::provider::ProvideCredentials;
use aws_sigv4::http_request::{
SignableBody, SignableRequest, SigningParams, SigningSettings, sign,
};
use aws_sigv4::sign::v4;
use aws_smithy_runtime_api::client::identity::Identity;
use sha2::{Digest, Sha256};
mod cache;
mod credentials;
mod error;
mod role;
mod runtime;
mod signing;
pub use cache::{Clock, CredentialScope, CredentialState, SystemClock};
pub use credentials::{Credentials, session_credentials, static_credentials};
pub use error::Error;
const CREDENTIAL_FETCH_LOCK_STRIPES: usize = 64;
#[derive(Clone)]
pub struct Credentials(SdkCredentials);
impl Credentials {
pub fn new(
access_key_id: impl Into<String>,
secret_access_key: impl Into<String>,
session_token: Option<String>,
expires_after: Option<SystemTime>,
provider_name: &'static str,
) -> Self {
Self(SdkCredentials::new(
access_key_id,
secret_access_key,
session_token,
expires_after,
provider_name,
))
}
pub fn access_key_id(&self) -> &str {
self.0.access_key_id()
}
pub fn session_token(&self) -> Option<&str> {
self.0.session_token()
}
}
impl fmt::Debug for Credentials {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("Credentials([REDACTED])")
}
}
struct CredentialCache {
entries: HashMap<CredentialScope, (Credentials, Duration)>,
expirations: BinaryHeap<Reverse<(Duration, CredentialScope)>>,
max_entries: usize,
}
impl Default for CredentialCache {
fn default() -> Self {
Self::new(200)
}
}
impl CredentialCache {
fn new(max_entries: usize) -> Self {
Self {
entries: HashMap::new(),
expirations: BinaryHeap::new(),
max_entries: max_entries.max(1),
}
}
fn get(&mut self, key: &CredentialScope, now: Duration) -> Option<Credentials> {
let (credentials, expiration) = self.entries.get(key)?;
if *expiration > now {
return Some(credentials.clone());
}
self.entries.remove(key);
None
}
fn insert(
&mut self,
key: CredentialScope,
credentials: Credentials,
ttl: Duration,
now: Duration,
) {
while let Some(Reverse((expiration, key))) = self.expirations.peek().cloned() {
if self.entries.get(&key).map(|(_, current)| *current) != Some(expiration) {
self.expirations.pop();
} else if expiration <= now || self.entries.len() >= self.max_entries {
self.expirations.pop();
self.entries.remove(&key);
} else {
break;
}
}
let expiration = now + ttl;
self.entries.insert(key.clone(), (credentials, expiration));
self.expirations.push(Reverse((expiration, key)));
}
}
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct CredentialScope([u8; 32]);
impl CredentialScope {
pub fn from_optional_values<'a>(
namespace: &str,
values: impl IntoIterator<Item = Option<&'a str>>,
) -> Self {
let mut hasher = Sha256::new();
hasher.update(namespace.len().to_le_bytes());
hasher.update(namespace.as_bytes());
for value in values {
match value {
Some(value) => {
hasher.update([1]);
hasher.update(value.len().to_le_bytes());
hasher.update(value.as_bytes());
}
None => hasher.update([0]),
}
}
Self(hasher.finalize().into())
}
}
impl fmt::Debug for CredentialScope {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("CredentialScope([REDACTED])")
}
}
pub trait Clock: Send + Sync {
fn now(&self) -> Duration;
}
#[derive(Default)]
pub struct SystemClock;
impl Clock for SystemClock {
fn now(&self) -> Duration {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
}
}
pub struct CredentialState<R, C = SystemClock> {
runtime: R,
clock: C,
cache: Mutex<CredentialCache>,
fetch_locks: Box<[tokio::sync::Mutex<()>]>,
}
impl<R> CredentialState<R, SystemClock> {
pub fn new(runtime: R, max_entries: usize) -> Self {
Self::with_clock(runtime, max_entries, SystemClock)
}
}
impl<R, C> CredentialState<R, C>
where
C: Clock,
{
pub fn with_clock(runtime: R, max_entries: usize, clock: C) -> Self {
Self {
runtime,
clock,
cache: Mutex::new(CredentialCache::new(max_entries)),
fetch_locks: (0..CREDENTIAL_FETCH_LOCK_STRIPES)
.map(|_| tokio::sync::Mutex::new(()))
.collect(),
}
}
pub fn runtime(&self) -> &R {
&self.runtime
}
pub async fn get_or_acquire<E, F, Fut>(
&self,
scope: CredentialScope,
ttl: Duration,
acquire: F,
) -> Result<Credentials, E>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<Credentials, E>>,
{
let mut stripe_hasher = std::collections::hash_map::DefaultHasher::new();
scope.hash(&mut stripe_hasher);
let stripe = stripe_hasher.finish() as usize % self.fetch_locks.len();
let _fetch_guard = self.fetch_locks[stripe].lock().await;
if let Some(credentials) = self
.cache
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.get(&scope, self.clock.now())
{
return Ok(credentials);
}
let credentials = acquire().await?;
self.cache
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(scope, credentials.clone(), ttl, self.clock.now());
Ok(credentials)
}
}
#[derive(Clone, PartialEq, Eq)]
pub struct AssumeRoleRequest {
pub role: String,
pub session_name: String,
pub region: Option<String>,
pub endpoint: Option<String>,
pub source_credentials: Option<Credentials>,
pub external_id: Option<String>,
}
impl fmt::Debug for AssumeRoleRequest {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("AssumeRoleRequest")
.field("role", &self.role)
.field("session_name", &self.session_name)
.field("region", &self.region)
.field("endpoint", &self.endpoint)
.field("source_credentials", &self.source_credentials.is_some())
.field("external_id", &self.external_id.is_some())
.finish()
}
}
#[derive(Clone, PartialEq, Eq)]
pub struct WebIdentityRequest {
pub token: String,
pub role: String,
pub session_name: String,
pub region: Option<String>,
pub endpoint: Option<String>,
}
impl fmt::Debug for WebIdentityRequest {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("WebIdentityRequest")
.field("token", &"[REDACTED]")
.field("role", &self.role)
.field("session_name", &self.session_name)
.field("region", &self.region)
.field("endpoint", &self.endpoint)
.finish()
}
}
pub type CredentialFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
pub trait CredentialRuntime: Send + Sync {
fn profile<'a>(&'a self, name: &'a str) -> CredentialFuture<'a, Credentials>;
fn ambient(&self) -> CredentialFuture<'_, Credentials>;
fn assume_role(&self, request: AssumeRoleRequest) -> CredentialFuture<'_, Credentials>;
fn web_identity(&self, request: WebIdentityRequest) -> CredentialFuture<'_, Credentials>;
fn caller_identity(
&self,
region: Option<String>,
endpoint: Option<String>,
) -> CredentialFuture<'_, Option<String>>;
}
#[derive(Default)]
pub struct NativeCredentialRuntime;
impl CredentialRuntime for NativeCredentialRuntime {
fn profile<'a>(&'a self, name: &'a str) -> CredentialFuture<'a, Credentials> {
Box::pin(profile_credentials(name))
}
fn ambient(&self) -> CredentialFuture<'_, Credentials> {
Box::pin(default_credentials())
}
fn assume_role(&self, request: AssumeRoleRequest) -> CredentialFuture<'_, Credentials> {
Box::pin(assume_role_credentials(request))
}
fn web_identity(&self, request: WebIdentityRequest) -> CredentialFuture<'_, Credentials> {
Box::pin(web_identity_credentials(request))
}
fn caller_identity(
&self,
region: Option<String>,
endpoint: Option<String>,
) -> CredentialFuture<'_, Option<String>> {
Box::pin(caller_identity(region, endpoint))
}
}
impl PartialEq for Credentials {
fn eq(&self, other: &Self) -> bool {
self.0.access_key_id() == other.0.access_key_id()
&& self.0.secret_access_key() == other.0.secret_access_key()
&& self.0.session_token() == other.0.session_token()
}
}
impl Eq for Credentials {}
pub fn static_credentials(
access_key_id: impl Into<String>,
secret_access_key: impl Into<String>,
) -> Credentials {
Credentials::new(
access_key_id,
secret_access_key,
None,
None,
"litellm-static",
)
}
pub fn session_credentials(
access_key_id: impl Into<String>,
secret_access_key: impl Into<String>,
session_token: impl Into<String>,
provider_name: &'static str,
) -> Credentials {
Credentials::new(
access_key_id,
secret_access_key,
Some(session_token.into()),
None,
provider_name,
)
}
pub async fn profile_credentials(name: &str) -> Result<Credentials, Error> {
let provider = aws_config::profile::ProfileFileCredentialsProvider::builder()
.profile_name(name)
.build();
provider
.provide_credentials()
.await
.map(Credentials)
.map_err(|error| Error::new(format!("AWS profile credentials failed: {error}")))
}
pub async fn default_credentials() -> Result<Credentials, Error> {
let provider = aws_config::default_provider::credentials::DefaultCredentialsChain::builder()
.build()
.await;
provider
.provide_credentials()
.await
.map(Credentials)
.map_err(|error| Error::new(format!("AWS default credentials failed: {error}")))
}
fn sdk_loader(region: Option<String>, endpoint: Option<String>) -> aws_config::ConfigLoader {
let loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
let loader = match region {
Some(region) => loader.region(aws_types::region::Region::new(region)),
None => loader,
};
match endpoint {
Some(endpoint) => loader.endpoint_url(endpoint),
None => loader,
}
}
pub async fn assume_role_credentials(request: AssumeRoleRequest) -> Result<Credentials, Error> {
let mut loader = sdk_loader(request.region, request.endpoint);
if let Some(credentials) = request.source_credentials {
loader = loader.credentials_provider(credentials.0);
}
let sdk_config = loader.load().await;
let builder = aws_config::sts::AssumeRoleProvider::builder(request.role)
.session_name(request.session_name);
let builder = match request.external_id {
Some(id) => builder.external_id(id),
None => builder,
};
builder
.configure(&sdk_config)
.build()
.await
.provide_credentials()
.await
.map(Credentials)
.map_err(|error| Error::new(format!("AWS role credentials failed: {error}")))
}
pub async fn web_identity_credentials(request: WebIdentityRequest) -> Result<Credentials, Error> {
let sdk_config = sdk_loader(request.region, request.endpoint).load().await;
let response = aws_sdk_sts::Client::new(&sdk_config)
.assume_role_with_web_identity()
.role_arn(request.role)
.role_session_name(request.session_name)
.web_identity_token(request.token)
.send()
.await
.map_err(|error| Error::new(format!("AWS web identity credentials failed: {error}")))?;
let credentials = response
.credentials()
.ok_or_else(|| Error::new("AWS web identity response had no credentials"))?;
let expiration = SystemTime::try_from(*credentials.expiration())
.map_err(|error| Error::new(format!("AWS web identity expiration was invalid: {error}")))?;
Ok(Credentials::new(
credentials.access_key_id(),
credentials.secret_access_key(),
Some(credentials.session_token().to_string()),
Some(expiration),
"litellm-web-identity",
))
}
pub async fn caller_identity(
region: Option<String>,
endpoint: Option<String>,
) -> Result<Option<String>, Error> {
let sdk_config = sdk_loader(region, endpoint).load().await;
match aws_sdk_sts::Client::new(&sdk_config)
.get_caller_identity()
.send()
.await
{
Ok(response) => Ok(response.arn().map(str::to_string)),
Err(_) => Ok(None),
}
}
pub fn role_identity(arn: &str) -> Option<(&str, &str, &str)> {
let mut parts = arn.splitn(6, ':');
let ("arn", partition, _, _, account, resource) = (
parts.next()?,
parts.next()?,
parts.next()?,
parts.next()?,
parts.next()?,
parts.next()?,
) else {
return None;
};
let role = if let Some(role) = resource.strip_prefix("role/") {
role.rsplit('/').next()?
} else {
resource.strip_prefix("assumed-role/")?.split('/').next()?
};
Some((partition, account, role))
}
pub fn same_role_arns(target: &str, caller: &str) -> bool {
role_identity(target) == role_identity(caller)
}
pub struct SigV4Request<'a> {
pub method: &'a str,
pub uri: &'a str,
pub body: &'a [u8],
pub headers: &'a BTreeMap<String, String>,
pub region: &'a str,
pub service: &'a str,
pub signing_time: SystemTime,
}
pub fn sign_v4(
request: SigV4Request<'_>,
credentials: &Credentials,
) -> Result<BTreeMap<String, String>, Error> {
let identity: Identity = credentials.0.clone().into();
let params = v4::SigningParams::builder()
.identity(&identity)
.region(request.region)
.name(request.service)
.time(request.signing_time)
.settings(SigningSettings::default())
.build()
.map(SigningParams::from)
.map_err(|error| Error::new(format!("AWS signing parameters failed: {error}")))?;
let header_refs = request
.headers
.iter()
.map(|(name, value)| (name.as_str(), value.as_str()));
let signable = SignableRequest::new(
request.method,
request.uri,
header_refs,
SignableBody::Bytes(request.body),
)
.map_err(|error| Error::new(format!("AWS signable request failed: {error}")))?;
let (instructions, _) = sign(signable, &params)
.map_err(|error| Error::new(format!("AWS request signing failed: {error}")))?
.into_parts();
Ok(instructions
.headers()
.map(|(name, value)| {
let normalized_name = match name {
"authorization" => "Authorization",
"x-amz-date" => "X-Amz-Date",
"x-amz-security-token" => "X-Amz-Security-Token",
_ => name,
};
(normalized_name.to_string(), value.to_string())
})
.collect())
}
#[cfg(test)]
mod tests {
use std::sync::{Arc, Mutex};
use super::*;
struct FixtureRuntime {
effects: Arc<Mutex<Vec<&'static str>>>,
}
impl CredentialRuntime for FixtureRuntime {
fn profile<'a>(&'a self, _name: &'a str) -> CredentialFuture<'a, Credentials> {
self.effects.lock().unwrap().push("profile");
Box::pin(std::future::ready(Ok(static_credentials(
"profile", "secret",
))))
}
fn ambient(&self) -> CredentialFuture<'_, Credentials> {
self.effects.lock().unwrap().push("ambient");
Box::pin(std::future::ready(Ok(static_credentials(
"ambient", "secret",
))))
}
fn assume_role(&self, _request: AssumeRoleRequest) -> CredentialFuture<'_, Credentials> {
self.effects.lock().unwrap().push("assume-role");
Box::pin(std::future::ready(Ok(static_credentials("role", "secret"))))
}
fn web_identity(&self, _request: WebIdentityRequest) -> CredentialFuture<'_, Credentials> {
self.effects.lock().unwrap().push("web-identity");
Box::pin(std::future::ready(Ok(static_credentials("web", "secret"))))
}
fn caller_identity(
&self,
_region: Option<String>,
_endpoint: Option<String>,
) -> CredentialFuture<'_, Option<String>> {
self.effects.lock().unwrap().push("caller-identity");
Box::pin(std::future::ready(Ok(None)))
}
}
fn signing_request<'a>(
uri: &'a str,
body: &'a [u8],
headers: &'a BTreeMap<String, String>,
service: &'a str,
) -> SigV4Request<'a> {
SigV4Request {
method: "POST",
uri,
body,
headers,
region: "us-east-1",
service,
signing_time: SystemTime::UNIX_EPOCH + Duration::from_secs(1_704_164_645),
}
}
#[test]
fn cache_expiry_and_bounds_are_clock_driven() {
let mut cache = CredentialCache::new(1);
let first = CredentialScope::from_optional_values("test", [Some("first")]);
let second = CredentialScope::from_optional_values("test", [Some("second")]);
cache.insert(
first.clone(),
static_credentials("ak1", "sk1"),
Duration::from_secs(10),
Duration::ZERO,
);
assert_eq!(
cache
.get(&first, Duration::from_secs(9))
.unwrap()
.access_key_id(),
"ak1"
);
cache.insert(
second.clone(),
static_credentials("ak2", "sk2"),
Duration::from_secs(10),
Duration::ZERO,
);
assert!(cache.get(&first, Duration::ZERO).is_none());
assert!(cache.get(&second, Duration::from_secs(11)).is_none());
}
#[tokio::test]
async fn credential_state_coordinates_concurrent_misses() {
use std::sync::atomic::{AtomicUsize, Ordering};
let state = CredentialState::new((), 1);
let acquisitions = AtomicUsize::new(0);
let scope = CredentialScope::from_optional_values("test", [Some("identity")]);
let acquire = || async {
acquisitions.fetch_add(1, Ordering::SeqCst);
tokio::task::yield_now().await;
Ok::<_, ()>(static_credentials("ak", "sk"))
};
let (first, second) = tokio::join!(
state.get_or_acquire(scope.clone(), Duration::from_secs(10), acquire),
state.get_or_acquire(scope, Duration::from_secs(10), acquire),
);
assert_eq!(first.unwrap(), second.unwrap());
assert_eq!(acquisitions.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn credential_state_reacquires_at_the_injected_expiry_boundary() {
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
struct TestClock(Arc<AtomicU64>);
impl Clock for TestClock {
fn now(&self) -> Duration {
Duration::from_secs(self.0.load(Ordering::SeqCst))
}
}
let now = Arc::new(AtomicU64::new(0));
let state = CredentialState::with_clock((), 1, TestClock(now.clone()));
let acquisitions = AtomicUsize::new(0);
let scope = CredentialScope::from_optional_values("test", [Some("identity")]);
let acquire = || async {
acquisitions.fetch_add(1, Ordering::SeqCst);
Ok::<_, ()>(static_credentials("ak", "sk"))
};
state
.get_or_acquire(scope.clone(), Duration::from_secs(10), acquire)
.await
.unwrap();
now.store(9, Ordering::SeqCst);
state
.get_or_acquire(scope.clone(), Duration::from_secs(10), acquire)
.await
.unwrap();
now.store(10, Ordering::SeqCst);
state
.get_or_acquire(scope, Duration::from_secs(10), acquire)
.await
.unwrap();
assert_eq!(acquisitions.load(Ordering::SeqCst), 2);
}
#[test]
fn credential_scope_does_not_expose_key_material() {
let scope = CredentialScope::from_optional_values(
"test",
[Some("visible-id"), Some("never-print-secret"), None],
);
let debug = format!("{scope:?}");
assert!(!debug.contains("visible-id"));
assert!(!debug.contains("never-print-secret"));
}
#[test]
fn signing_matches_the_bedrock_golden_vector() {
let uri = "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke";
let body = br#"{"input":"hello"}"#;
let headers = BTreeMap::from([("Content-Type".into(), "application/json".into())]);
let credentials = Credentials::new(
"AKIDEXAMPLE",
"wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
Some("session-token".into()),
None,
"test",
);
let signed = sign_v4(
signing_request(uri, body, &headers, "bedrock"),
&credentials,
)
.expect("signature");
assert_eq!(
signed.get("X-Amz-Date").map(String::as_str),
Some("20240102T030405Z")
);
assert_eq!(
signed.get("Authorization").map(String::as_str),
Some(
"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464"
)
);
}
#[test]
fn signer_is_generic_over_method_and_service() {
let uri = "https://sts.amazonaws.com/?Action=GetCallerIdentity&Version=2011-06-15";
let headers = BTreeMap::new();
let credentials = static_credentials("AKIDEXAMPLE", "secret");
let request = SigV4Request {
method: "GET",
uri,
body: b"",
headers: &headers,
region: "us-east-1",
service: "sts",
signing_time: SystemTime::UNIX_EPOCH,
};
let signed = sign_v4(request, &credentials).expect("signature");
assert!(signed["Authorization"].contains("/sts/aws4_request"));
}
#[test]
fn credentials_are_redacted() {
let credentials = static_credentials("visible-id", "never-print-secret");
let debug = format!("{credentials:?}");
assert!(!debug.contains("visible-id"));
assert!(!debug.contains("never-print-secret"));
}
#[test]
fn mechanism_inputs_are_redacted() {
let assume_role = AssumeRoleRequest {
role: "role".into(),
session_name: "session".into(),
region: None,
endpoint: None,
source_credentials: Some(static_credentials("visible-id", "secret")),
external_id: Some("external-secret".into()),
};
let web_identity = WebIdentityRequest {
token: "identity-secret".into(),
role: "role".into(),
session_name: "session".into(),
region: None,
endpoint: None,
};
let debug = format!("{assume_role:?} {web_identity:?}");
for secret in ["visible-id", "secret", "external-secret", "identity-secret"] {
assert!(!debug.contains(secret));
}
}
#[test]
fn role_matching_is_partition_account_and_role_aware() {
assert!(same_role_arns(
"arn:aws:iam::123456789012:role/path/demo",
"arn:aws:sts::123456789012:assumed-role/demo/session"
));
assert!(!same_role_arns(
"arn:aws:iam::123456789012:role/demo",
"arn:aws-cn:iam::123456789012:role/demo"
));
}
#[tokio::test]
async fn credential_io_is_injectable_for_policy_consumers() {
let effects = Arc::new(Mutex::new(Vec::new()));
let runtime = FixtureRuntime {
effects: effects.clone(),
};
assert_eq!(
runtime.profile("demo").await.unwrap().access_key_id(),
"profile"
);
assert_eq!(runtime.ambient().await.unwrap().access_key_id(), "ambient");
assert_eq!(
runtime
.caller_identity(Some("us-east-1".into()), None)
.await
.unwrap(),
None
);
assert_eq!(
effects.lock().unwrap().as_slice(),
&["profile", "ambient", "caller-identity"]
);
}
}
pub use role::{role_identity, same_role_arns};
pub use runtime::{
AssumeRoleRequest, CredentialFuture, CredentialRuntime, NativeCredentialRuntime,
WebIdentityRequest, assume_role_credentials, caller_identity, default_credentials,
profile_credentials, web_identity_credentials,
};
pub use signing::{SigV4Request, sign_v4};

43
litellm-rust/crates/auth-aws/src/role.rs Normal file
View file

@ -0,0 +1,43 @@
pub fn role_identity(arn: &str) -> Option<(&str, &str, &str)> {
let mut parts = arn.splitn(6, ':');
let ("arn", partition, _, _, account, resource) = (
parts.next()?,
parts.next()?,
parts.next()?,
parts.next()?,
parts.next()?,
parts.next()?,
) else {
return None;
};
let role = if let Some(role) = resource.strip_prefix("role/") {
role.rsplit('/').next()?
} else {
resource.strip_prefix("assumed-role/")?.split('/').next()?
};
Some((partition, account, role))
}
pub fn same_role_arns(target: &str, caller: &str) -> bool {
match (role_identity(target), role_identity(caller)) {
(Some(target), Some(caller)) => target == caller,
_ => false,
}
}
#[cfg(test)]
mod tests {
use super::same_role_arns;
#[test]
fn role_matching_is_partition_account_and_role_aware() {
assert!(same_role_arns(
"arn:aws:iam::123456789012:role/path/demo",
"arn:aws:sts::123456789012:assumed-role/demo/session"
));
assert!(!same_role_arns(
"arn:aws:iam::123456789012:role/demo",
"arn:aws-cn:iam::123456789012:role/demo"
));
}
}

262
litellm-rust/crates/auth-aws/src/runtime.rs Normal file
View file

@ -0,0 +1,262 @@
use std::future::Future;
use std::pin::Pin;
use std::time::SystemTime;
use aws_credential_types::provider::ProvideCredentials;
use veil::Redact;
use crate::{Credentials, Error};
#[derive(Clone, PartialEq, Eq, Redact)]
pub struct AssumeRoleRequest {
pub role: String,
pub session_name: String,
pub region: Option<String>,
pub endpoint: Option<String>,
pub source_credentials: Option<Credentials>,
#[redact(fixed = 8)]
pub external_id: Option<String>,
}
#[derive(Clone, PartialEq, Eq, Redact)]
pub struct WebIdentityRequest {
#[redact(fixed = 8)]
pub token: String,
pub role: String,
pub session_name: String,
pub region: Option<String>,
pub endpoint: Option<String>,
}
pub type CredentialFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
pub trait CredentialRuntime: Send + Sync {
fn profile<'a>(&'a self, name: &'a str) -> CredentialFuture<'a, Credentials>;
fn ambient(&self) -> CredentialFuture<'_, Credentials>;
fn assume_role(&self, request: AssumeRoleRequest) -> CredentialFuture<'_, Credentials>;
fn web_identity(&self, request: WebIdentityRequest) -> CredentialFuture<'_, Credentials>;
fn caller_identity(
&self,
region: Option<String>,
endpoint: Option<String>,
) -> CredentialFuture<'_, Option<String>>;
}
#[derive(Default)]
pub struct NativeCredentialRuntime;
impl CredentialRuntime for NativeCredentialRuntime {
fn profile<'a>(&'a self, name: &'a str) -> CredentialFuture<'a, Credentials> {
Box::pin(profile_credentials(name))
}
fn ambient(&self) -> CredentialFuture<'_, Credentials> {
Box::pin(default_credentials())
}
fn assume_role(&self, request: AssumeRoleRequest) -> CredentialFuture<'_, Credentials> {
Box::pin(assume_role_credentials(request))
}
fn web_identity(&self, request: WebIdentityRequest) -> CredentialFuture<'_, Credentials> {
Box::pin(web_identity_credentials(request))
}
fn caller_identity(
&self,
region: Option<String>,
endpoint: Option<String>,
) -> CredentialFuture<'_, Option<String>> {
Box::pin(caller_identity(region, endpoint))
}
}
pub async fn profile_credentials(name: &str) -> Result<Credentials, Error> {
let provider = aws_config::profile::ProfileFileCredentialsProvider::builder()
.profile_name(name)
.build();
provider
.provide_credentials()
.await
.map(Credentials)
.map_err(|error| Error::ProfileCredentials(error.to_string()))
}
pub async fn default_credentials() -> Result<Credentials, Error> {
let provider = aws_config::default_provider::credentials::DefaultCredentialsChain::builder()
.build()
.await;
provider
.provide_credentials()
.await
.map(Credentials)
.map_err(|error| Error::DefaultCredentials(error.to_string()))
}
fn sdk_loader(region: Option<String>, endpoint: Option<String>) -> aws_config::ConfigLoader {
let loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
let loader = match region {
Some(region) => loader.region(aws_types::region::Region::new(region)),
None => loader,
};
match endpoint {
Some(endpoint) => loader.endpoint_url(endpoint),
None => loader,
}
}
pub async fn assume_role_credentials(request: AssumeRoleRequest) -> Result<Credentials, Error> {
let mut loader = sdk_loader(request.region, request.endpoint);
if let Some(credentials) = request.source_credentials {
loader = loader.credentials_provider(credentials.0);
}
let sdk_config = loader.load().await;
let builder = aws_config::sts::AssumeRoleProvider::builder(request.role)
.session_name(request.session_name);
let builder = match request.external_id {
Some(id) => builder.external_id(id),
None => builder,
};
builder
.configure(&sdk_config)
.build()
.await
.provide_credentials()
.await
.map(Credentials)
.map_err(|error| Error::RoleCredentials(error.to_string()))
}
pub async fn web_identity_credentials(request: WebIdentityRequest) -> Result<Credentials, Error> {
let sdk_config = sdk_loader(request.region, request.endpoint).load().await;
let response = aws_sdk_sts::Client::new(&sdk_config)
.assume_role_with_web_identity()
.role_arn(request.role)
.role_session_name(request.session_name)
.web_identity_token(request.token)
.send()
.await
.map_err(|error| Error::WebIdentityCredentials(error.to_string()))?;
let credentials = response
.credentials()
.ok_or(Error::MissingWebIdentityCredentials)?;
let expiration = SystemTime::try_from(*credentials.expiration())
.map_err(|error| Error::InvalidWebIdentityExpiration(error.to_string()))?;
Ok(Credentials::new(
credentials.access_key_id(),
credentials.secret_access_key(),
Some(credentials.session_token().to_string()),
Some(expiration),
"litellm-web-identity",
))
}
pub async fn caller_identity(
region: Option<String>,
endpoint: Option<String>,
) -> Result<Option<String>, Error> {
let sdk_config = sdk_loader(region, endpoint).load().await;
match aws_sdk_sts::Client::new(&sdk_config)
.get_caller_identity()
.send()
.await
{
Ok(response) => Ok(response.arn().map(str::to_string)),
Err(_) => Ok(None),
}
}
#[cfg(test)]
mod tests {
use std::sync::{Arc, Mutex};
use super::*;
use crate::static_credentials;
struct FixtureRuntime {
effects: Arc<Mutex<Vec<&'static str>>>,
}
impl CredentialRuntime for FixtureRuntime {
fn profile<'a>(&'a self, _name: &'a str) -> CredentialFuture<'a, Credentials> {
self.effects.lock().unwrap().push("profile");
Box::pin(std::future::ready(Ok(static_credentials(
"profile", "secret",
))))
}
fn ambient(&self) -> CredentialFuture<'_, Credentials> {
self.effects.lock().unwrap().push("ambient");
Box::pin(std::future::ready(Ok(static_credentials(
"ambient", "secret",
))))
}
fn assume_role(&self, _request: AssumeRoleRequest) -> CredentialFuture<'_, Credentials> {
self.effects.lock().unwrap().push("assume-role");
Box::pin(std::future::ready(Ok(static_credentials("role", "secret"))))
}
fn web_identity(&self, _request: WebIdentityRequest) -> CredentialFuture<'_, Credentials> {
self.effects.lock().unwrap().push("web-identity");
Box::pin(std::future::ready(Ok(static_credentials("web", "secret"))))
}
fn caller_identity(
&self,
_region: Option<String>,
_endpoint: Option<String>,
) -> CredentialFuture<'_, Option<String>> {
self.effects.lock().unwrap().push("caller-identity");
Box::pin(std::future::ready(Ok(None)))
}
}
#[test]
fn mechanism_inputs_are_redacted() {
let assume_role = AssumeRoleRequest {
role: "role".into(),
session_name: "session".into(),
region: None,
endpoint: None,
source_credentials: Some(static_credentials("visible-id", "secret")),
external_id: Some("external-secret".into()),
};
let web_identity = WebIdentityRequest {
token: "identity-secret".into(),
role: "role".into(),
session_name: "session".into(),
region: None,
endpoint: None,
};
let debug = format!("{assume_role:?} {web_identity:?}");
for secret in ["visible-id", "secret", "external-secret", "identity-secret"] {
assert!(!debug.contains(secret));
}
}
#[tokio::test]
async fn credential_io_is_injectable_for_policy_consumers() {
let effects = Arc::new(Mutex::new(Vec::new()));
let runtime = FixtureRuntime {
effects: effects.clone(),
};
assert_eq!(
runtime.profile("demo").await.unwrap().access_key_id(),
"profile"
);
assert_eq!(runtime.ambient().await.unwrap().access_key_id(), "ambient");
assert_eq!(
runtime
.caller_identity(Some("us-east-1".into()), None)
.await
.unwrap(),
None
);
assert_eq!(
effects.lock().unwrap().as_slice(),
&["profile", "ambient", "caller-identity"]
);
}
}

134
litellm-rust/crates/auth-aws/src/signing.rs Normal file
View file

@ -0,0 +1,134 @@
use std::collections::BTreeMap;
use std::time::SystemTime;
use aws_sigv4::http_request::{
SignableBody, SignableRequest, SigningParams, SigningSettings, sign,
};
use aws_sigv4::sign::v4;
use aws_smithy_runtime_api::client::identity::Identity;
use crate::{Credentials, Error};
pub struct SigV4Request<'a> {
pub method: &'a str,
pub uri: &'a str,
pub body: &'a [u8],
pub headers: &'a BTreeMap<String, String>,
pub region: &'a str,
pub service: &'a str,
pub signing_time: SystemTime,
}
pub fn sign_v4(
request: SigV4Request<'_>,
credentials: &Credentials,
) -> Result<BTreeMap<String, String>, Error> {
let identity: Identity = credentials.0.clone().into();
let params = v4::SigningParams::builder()
.identity(&identity)
.region(request.region)
.name(request.service)
.time(request.signing_time)
.settings(SigningSettings::default())
.build()
.map(SigningParams::from)
.map_err(|error| Error::SigningParameters(error.to_string()))?;
let header_refs = request
.headers
.iter()
.map(|(name, value)| (name.as_str(), value.as_str()));
let signable = SignableRequest::new(
request.method,
request.uri,
header_refs,
SignableBody::Bytes(request.body),
)
.map_err(|error| Error::SignableRequest(error.to_string()))?;
let (instructions, _) = sign(signable, &params)
.map_err(|error| Error::RequestSigning(error.to_string()))?
.into_parts();
Ok(instructions
.headers()
.map(|(name, value)| {
let normalized_name = match name {
"authorization" => "Authorization",
"x-amz-date" => "X-Amz-Date",
"x-amz-security-token" => "X-Amz-Security-Token",
_ => name,
};
(normalized_name.to_string(), value.to_string())
})
.collect())
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::*;
use crate::static_credentials;
fn signing_request<'a>(
uri: &'a str,
body: &'a [u8],
headers: &'a BTreeMap<String, String>,
service: &'a str,
) -> SigV4Request<'a> {
SigV4Request {
method: "POST",
uri,
body,
headers,
region: "us-east-1",
service,
signing_time: SystemTime::UNIX_EPOCH + Duration::from_secs(1_704_164_645),
}
}
#[test]
fn signing_matches_the_bedrock_golden_vector() {
let uri = "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke";
let body = br#"{"input":"hello"}"#;
let headers = BTreeMap::from([("Content-Type".into(), "application/json".into())]);
let credentials = Credentials::new(
"AKIDEXAMPLE",
"wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
Some("session-token".into()),
None,
"test",
);
let signed = sign_v4(
signing_request(uri, body, &headers, "bedrock"),
&credentials,
)
.expect("signature");
assert_eq!(
signed.get("X-Amz-Date").map(String::as_str),
Some("20240102T030405Z")
);
assert_eq!(
signed.get("Authorization").map(String::as_str),
Some(
"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464"
)
);
}
#[test]
fn signer_is_generic_over_method_and_service() {
let uri = "https://sts.amazonaws.com/?Action=GetCallerIdentity&Version=2011-06-15";
let headers = BTreeMap::new();
let credentials = static_credentials("AKIDEXAMPLE", "secret");
let request = SigV4Request {
method: "GET",
uri,
body: b"",
headers: &headers,
region: "us-east-1",
service: "sts",
signing_time: SystemTime::UNIX_EPOCH,
};
let signed = sign_v4(request, &credentials).expect("signature");
assert!(signed["Authorization"].contains("/sts/aws4_request"));
}
}

108
litellm-rust/crates/auth-aws/tests/credential_state.rs Normal file
View file

@ -0,0 +1,108 @@
use std::collections::BTreeSet;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::time::Duration;
use litellm_auth_aws::{Clock, CredentialScope, CredentialState, static_credentials};
struct TestClock(Arc<AtomicU64>);
impl Clock for TestClock {
fn now(&self) -> Duration {
Duration::from_secs(self.0.load(Ordering::SeqCst))
}
}
#[test]
fn scopes_isolate_every_input_and_preserve_value_boundaries() {
let scopes = BTreeSet::from([
CredentialScope::from_optional_values("bedrock", [Some("ab"), Some("c")]),
CredentialScope::from_optional_values("bedrock", [Some("a"), Some("bc")]),
CredentialScope::from_optional_values("bedrock", [Some("ab"), None]),
CredentialScope::from_optional_values("bedrock", [Some("ab"), Some("")]),
CredentialScope::from_optional_values("federation", [Some("ab"), Some("c")]),
]);
assert_eq!(scopes.len(), 5);
}
#[tokio::test]
async fn failed_acquisitions_are_not_cached_or_shared_as_success() {
let state = CredentialState::new((), 1);
let acquisitions = AtomicUsize::new(0);
let scope = CredentialScope::from_optional_values("test", [Some("identity")]);
let acquire = || async {
let attempt = acquisitions.fetch_add(1, Ordering::SeqCst);
tokio::task::yield_now().await;
if attempt == 0 {
Err("synthetic acquisition failure")
} else {
Ok(static_credentials("second-attempt", "secret"))
}
};
let (first, second) = tokio::join!(
state.get_or_acquire(scope.clone(), Duration::from_secs(60), acquire),
state.get_or_acquire(scope, Duration::from_secs(60), acquire),
);
assert!(first.is_err() ^ second.is_err());
assert_eq!(
first
.ok()
.or_else(|| second.ok())
.expect("one waiter reacquires")
.access_key_id(),
"second-attempt"
);
assert_eq!(acquisitions.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn concurrent_waiters_singleflight_an_expired_refresh() {
let now = Arc::new(AtomicU64::new(0));
let state = CredentialState::with_clock((), 1, TestClock(now.clone()));
let acquisitions = AtomicUsize::new(0);
let scope = CredentialScope::from_optional_values("test", [Some("identity")]);
let acquire = || async {
let attempt = acquisitions.fetch_add(1, Ordering::SeqCst) + 1;
tokio::task::yield_now().await;
Ok::<_, ()>(static_credentials(format!("attempt-{attempt}"), "secret"))
};
let initial = state
.get_or_acquire(scope.clone(), Duration::from_secs(10), acquire)
.await
.expect("initial acquisition");
assert_eq!(initial.access_key_id(), "attempt-1");
now.store(10, Ordering::SeqCst);
let (first, second) = tokio::join!(
state.get_or_acquire(scope.clone(), Duration::from_secs(10), acquire),
state.get_or_acquire(scope, Duration::from_secs(10), acquire),
);
assert_eq!(first.expect("first waiter").access_key_id(), "attempt-2");
assert_eq!(second.expect("second waiter").access_key_id(), "attempt-2");
assert_eq!(acquisitions.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn cache_entries_are_isolated_by_scope() {
let state = CredentialState::new((), 2);
let acquisitions = AtomicUsize::new(0);
let first_scope = CredentialScope::from_optional_values("test", [Some("first")]);
let second_scope = CredentialScope::from_optional_values("test", [Some("second")]);
for scope in [first_scope.clone(), second_scope.clone(), first_scope] {
state
.get_or_acquire(scope, Duration::from_secs(60), || async {
let attempt = acquisitions.fetch_add(1, Ordering::SeqCst) + 1;
Ok::<_, ()>(static_credentials(format!("attempt-{attempt}"), "secret"))
})
.await
.expect("credential acquisition");
}
assert_eq!(acquisitions.load(Ordering::SeqCst), 2);
}

153
litellm-rust/crates/auth-aws/tests/signing.rs Normal file
View file

@ -0,0 +1,153 @@
use std::collections::{BTreeMap, BTreeSet};
use std::time::SystemTime;
use litellm_auth_aws::{Credentials, SigV4Request, sign_v4, static_credentials};
const URI: &str = "https://sts.amazonaws.com/?Action=GetCallerIdentity&Version=2011-06-15";
const SIGNING_TIME: SystemTime = SystemTime::UNIX_EPOCH;
fn authorization(
method: &str,
uri: &str,
body: &[u8],
headers: &BTreeMap<String, String>,
region: &str,
service: &str,
credentials: &Credentials,
) -> String {
sign_v4(
SigV4Request {
method,
uri,
body,
headers,
region,
service,
signing_time: SIGNING_TIME,
},
credentials,
)
.expect("signature")
.remove("Authorization")
.expect("authorization header")
}
#[test]
fn every_signed_request_dimension_changes_the_authorization() {
let headers = BTreeMap::from([(
"Content-Type".into(),
"application/x-www-form-urlencoded".into(),
)]);
let changed_headers = BTreeMap::from([("Content-Type".into(), "application/json".into())]);
let credentials = static_credentials("AKIDEXAMPLE", "secret");
let baseline = authorization(
"POST",
URI,
b"payload",
&headers,
"us-east-1",
"sts",
&credentials,
);
let signatures = BTreeSet::from([
baseline.clone(),
authorization(
"GET",
URI,
b"payload",
&headers,
"us-east-1",
"sts",
&credentials,
),
authorization(
"POST",
"https://sts.amazonaws.com/?Action=GetCallerIdentity&Version=2011-06-14",
b"payload",
&headers,
"us-east-1",
"sts",
&credentials,
),
authorization(
"POST",
URI,
b"payloaD",
&headers,
"us-east-1",
"sts",
&credentials,
),
authorization(
"POST",
URI,
b"payload",
&changed_headers,
"us-east-1",
"sts",
&credentials,
),
authorization(
"POST",
URI,
b"payload",
&headers,
"us-west-2",
"sts",
&credentials,
),
authorization(
"POST",
URI,
b"payload",
&headers,
"us-east-1",
"bedrock",
&credentials,
),
]);
assert_eq!(signatures.len(), 7);
}
#[test]
fn session_token_is_emitted_and_covered_by_the_signature() {
let headers = BTreeMap::new();
let first = Credentials::new(
"AKIDEXAMPLE",
"secret",
Some("first-session".into()),
None,
"test",
);
let second = Credentials::new(
"AKIDEXAMPLE",
"secret",
Some("second-session".into()),
None,
"test",
);
let signed = sign_v4(
SigV4Request {
method: "GET",
uri: URI,
body: b"",
headers: &headers,
region: "us-east-1",
service: "sts",
signing_time: SIGNING_TIME,
},
&first,
)
.expect("signature");
let first_authorization = signed.get("Authorization").expect("authorization");
let second_authorization =
authorization("GET", URI, b"", &headers, "us-east-1", "sts", &second);
assert_eq!(
signed.get("X-Amz-Security-Token").map(String::as_str),
Some("first-session")
);
assert!(first_authorization.contains("x-amz-security-token"));
assert_ne!(first_authorization, &second_authorization);
}

145
litellm-rust/crates/auth-aws/tests/sts.rs Normal file
View file

@ -0,0 +1,145 @@
use std::io::{Read, Write};
use std::net::TcpListener;
use std::sync::mpsc::{self, Receiver};
use std::thread::JoinHandle;
use std::time::Duration;
use litellm_auth_aws::{
AssumeRoleRequest, Error, WebIdentityRequest, assume_role_credentials, static_credentials,
web_identity_credentials,
};
const EXPIRATION: &str = "2035-01-02T03:04:05Z";
fn sts_server(response_body: String) -> (String, Receiver<String>, JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").expect("listener");
let endpoint = format!("http://{}", listener.local_addr().expect("address"));
let (sender, receiver) = mpsc::channel();
let handle = std::thread::spawn(move || {
let (mut stream, _) = listener.accept().expect("request");
stream
.set_read_timeout(Some(Duration::from_secs(5)))
.expect("read timeout");
let mut request = Vec::new();
let mut buffer = [0; 4096];
loop {
let count = stream.read(&mut buffer).expect("request bytes");
assert_ne!(count, 0, "request ended before its body arrived");
request.extend_from_slice(&buffer[..count]);
if request_is_complete(&request) {
break;
}
}
sender
.send(String::from_utf8(request).expect("utf-8 request"))
.expect("captured request");
let response = format!(
"HTTP/1.1 200 OK\r\nContent-Type: text/xml\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
response_body.len(),
response_body
);
stream.write_all(response.as_bytes()).expect("response");
});
(endpoint, receiver, handle)
}
fn request_is_complete(request: &[u8]) -> bool {
let Some(header_end) = request.windows(4).position(|window| window == b"\r\n\r\n") else {
return false;
};
let headers = String::from_utf8_lossy(&request[..header_end]);
let content_length = headers.lines().find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().expect("content length"))
});
request.len() >= header_end + 4 + content_length.unwrap_or(0)
}
fn credentials_xml() -> String {
format!(
"<Credentials><AccessKeyId>acquired-access-key</AccessKeyId><SecretAccessKey>acquired-secret-key</SecretAccessKey><SessionToken>acquired-session-token</SessionToken><Expiration>{EXPIRATION}</Expiration></Credentials>"
)
}
#[tokio::test]
async fn assume_role_sends_the_exact_policy_inputs_and_parses_credentials() {
let response_body = format!(
"<AssumeRoleResponse xmlns=\"https://sts.amazonaws.com/doc/2011-06-15/\"><AssumeRoleResult>{}<AssumedRoleUser><Arn>arn:aws:sts::123456789012:assumed-role/demo/test-session</Arn><AssumedRoleId>id:test-session</AssumedRoleId></AssumedRoleUser></AssumeRoleResult><ResponseMetadata><RequestId>request-id</RequestId></ResponseMetadata></AssumeRoleResponse>",
credentials_xml()
);
let (endpoint, request, server) = sts_server(response_body);
let credentials = assume_role_credentials(AssumeRoleRequest {
role: "arn:aws:iam::123456789012:role/path/demo".into(),
session_name: "test-session".into(),
region: Some("us-west-2".into()),
endpoint: Some(endpoint),
source_credentials: Some(static_credentials("source-access-key", "source-secret-key")),
external_id: Some("external-value".into()),
})
.await
.expect("assume role credentials");
let request = request.recv().expect("captured request");
server.join().expect("server");
assert!(request.starts_with("POST / HTTP/1.1\r\n"));
assert!(request.contains("Action=AssumeRole"));
assert!(request.contains("RoleArn=arn%3Aaws%3Aiam%3A%3A123456789012%3Arole%2Fpath%2Fdemo"));
assert!(request.contains("RoleSessionName=test-session"));
assert!(request.contains("ExternalId=external-value"));
assert!(
request
.to_ascii_lowercase()
.contains("authorization: aws4-hmac-sha256 credential=source-access-key/")
);
assert_eq!(credentials.access_key_id(), "acquired-access-key");
assert_eq!(credentials.session_token(), Some("acquired-session-token"));
}
#[tokio::test]
async fn web_identity_sends_the_token_and_parses_credentials() {
let response_body = format!(
"<AssumeRoleWithWebIdentityResponse xmlns=\"https://sts.amazonaws.com/doc/2011-06-15/\"><AssumeRoleWithWebIdentityResult>{}<AssumedRoleUser><Arn>arn:aws:sts::123456789012:assumed-role/demo/web-session</Arn><AssumedRoleId>id:web-session</AssumedRoleId></AssumedRoleUser><Audience>audience</Audience><Provider>provider</Provider><SubjectFromWebIdentityToken>subject</SubjectFromWebIdentityToken></AssumeRoleWithWebIdentityResult><ResponseMetadata><RequestId>request-id</RequestId></ResponseMetadata></AssumeRoleWithWebIdentityResponse>",
credentials_xml()
);
let (endpoint, request, server) = sts_server(response_body);
let credentials = web_identity_credentials(WebIdentityRequest {
token: "header.payload.signature".into(),
role: "arn:aws:iam::123456789012:role/demo".into(),
session_name: "web-session".into(),
region: Some("us-east-1".into()),
endpoint: Some(endpoint),
})
.await
.expect("web identity credentials");
let request = request.recv().expect("captured request");
server.join().expect("server");
assert!(request.starts_with("POST / HTTP/1.1\r\n"));
assert!(request.contains("Action=AssumeRoleWithWebIdentity"));
assert!(request.contains("RoleArn=arn%3Aaws%3Aiam%3A%3A123456789012%3Arole%2Fdemo"));
assert!(request.contains("RoleSessionName=web-session"));
assert!(request.contains("WebIdentityToken=header.payload.signature"));
assert!(!request.to_ascii_lowercase().contains("\r\nauthorization:"));
assert_eq!(credentials.access_key_id(), "acquired-access-key");
assert_eq!(credentials.session_token(), Some("acquired-session-token"));
}
#[tokio::test]
async fn web_identity_rejects_a_success_response_without_credentials() {
let response_body = "<AssumeRoleWithWebIdentityResponse xmlns=\"https://sts.amazonaws.com/doc/2011-06-15/\"><AssumeRoleWithWebIdentityResult></AssumeRoleWithWebIdentityResult><ResponseMetadata><RequestId>request-id</RequestId></ResponseMetadata></AssumeRoleWithWebIdentityResponse>".to_string();
let (endpoint, request, server) = sts_server(response_body);
let error = web_identity_credentials(WebIdentityRequest {
token: "header.payload.signature".into(),
role: "arn:aws:iam::123456789012:role/demo".into(),
session_name: "web-session".into(),
region: Some("us-east-1".into()),
endpoint: Some(endpoint),
})
.await
.expect_err("missing credentials");
request.recv().expect("captured request");
server.join().expect("server");
assert_eq!(error, Error::MissingWebIdentityCredentials);
}

2
litellm-rust/crates/core/src/providers/bedrock/aws_base.rs
View file

@ -335,7 +335,7 @@ where
}
fn auth_error(error: litellm_auth_aws::Error) -> Error {
Error::Auth(error.message().to_string())
Error::Auth(error.to_string())
}
async fn is_already_running_as_role(