diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index c404b93f283..6fb0c2a3c75 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -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" diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml index 03769eaf0c3..18c4daab082 100644 --- a/litellm-rust/crates/auth-aws/Cargo.toml +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -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"] } diff --git a/litellm-rust/crates/auth-aws/src/cache.rs b/litellm-rust/crates/auth-aws/src/cache.rs new file mode 100644 index 00000000000..1b05f638f63 --- /dev/null +++ b/litellm-rust/crates/auth-aws/src/cache.rs @@ -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, + expirations: BinaryHeap>, + 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 { + 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>, + ) -> 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 { + runtime: R, + clock: C, + cache: Mutex, + fetch_locks: Box<[tokio::sync::Mutex<()>]>, +} + +impl CredentialState { + pub fn new(runtime: R, max_entries: usize) -> Self { + Self::with_clock(runtime, max_entries, SystemClock) + } +} + +impl CredentialState +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( + &self, + scope: CredentialScope, + ttl: Duration, + acquire: F, + ) -> Result + where + F: FnOnce() -> Fut, + Fut: Future>, + { + 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); + + 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")); + } +} diff --git a/litellm-rust/crates/auth-aws/src/credentials.rs b/litellm-rust/crates/auth-aws/src/credentials.rs new file mode 100644 index 00000000000..0bb6e3f8615 --- /dev/null +++ b/litellm-rust/crates/auth-aws/src/credentials.rs @@ -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, + secret_access_key: impl Into, + session_token: Option, + expires_after: Option, + 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, + secret_access_key: impl Into, +) -> Credentials { + Credentials::new( + access_key_id, + secret_access_key, + None, + None, + "litellm-static", + ) +} + +pub fn session_credentials( + access_key_id: impl Into, + secret_access_key: impl Into, + session_token: impl Into, + 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")); + } +} diff --git a/litellm-rust/crates/auth-aws/src/error.rs b/litellm-rust/crates/auth-aws/src/error.rs new file mode 100644 index 00000000000..7048eaa4790 --- /dev/null +++ b/litellm-rust/crates/auth-aws/src/error.rs @@ -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), +} diff --git a/litellm-rust/crates/auth-aws/src/lib.rs b/litellm-rust/crates/auth-aws/src/lib.rs index d208bb32627..a50c0ccd10f 100644 --- a/litellm-rust/crates/auth-aws/src/lib.rs +++ b/litellm-rust/crates/auth-aws/src/lib.rs @@ -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, - secret_access_key: impl Into, - session_token: Option, - expires_after: Option, - 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, - expirations: BinaryHeap>, - 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 { - 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>, - ) -> 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 { - runtime: R, - clock: C, - cache: Mutex, - fetch_locks: Box<[tokio::sync::Mutex<()>]>, -} - -impl CredentialState { - pub fn new(runtime: R, max_entries: usize) -> Self { - Self::with_clock(runtime, max_entries, SystemClock) - } -} - -impl CredentialState -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( - &self, - scope: CredentialScope, - ttl: Duration, - acquire: F, - ) -> Result - where - F: FnOnce() -> Fut, - Fut: Future>, - { - 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, - pub endpoint: Option, - pub source_credentials: Option, - pub external_id: Option, -} - -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, - pub endpoint: Option, -} - -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> + 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, - endpoint: Option, - ) -> CredentialFuture<'_, Option>; -} - -#[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, - endpoint: Option, - ) -> CredentialFuture<'_, Option> { - 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, - secret_access_key: impl Into, -) -> Credentials { - Credentials::new( - access_key_id, - secret_access_key, - None, - None, - "litellm-static", - ) -} - -pub fn session_credentials( - access_key_id: impl Into, - secret_access_key: impl Into, - session_token: impl Into, - 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 { - 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 { - 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, endpoint: Option) -> 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 { - 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 { - 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, - endpoint: Option, -) -> Result, 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, - pub region: &'a str, - pub service: &'a str, - pub signing_time: SystemTime, -} - -pub fn sign_v4( - request: SigV4Request<'_>, - credentials: &Credentials, -) -> Result, 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, ¶ms) - .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>>, - } - - 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, - _endpoint: Option, - ) -> CredentialFuture<'_, Option> { - 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, - 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); - - 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}; diff --git a/litellm-rust/crates/auth-aws/src/role.rs b/litellm-rust/crates/auth-aws/src/role.rs new file mode 100644 index 00000000000..4d28626adbc --- /dev/null +++ b/litellm-rust/crates/auth-aws/src/role.rs @@ -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" + )); + } +} diff --git a/litellm-rust/crates/auth-aws/src/runtime.rs b/litellm-rust/crates/auth-aws/src/runtime.rs new file mode 100644 index 00000000000..ab361581423 --- /dev/null +++ b/litellm-rust/crates/auth-aws/src/runtime.rs @@ -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, + pub endpoint: Option, + pub source_credentials: Option, + #[redact(fixed = 8)] + pub external_id: Option, +} + +#[derive(Clone, PartialEq, Eq, Redact)] +pub struct WebIdentityRequest { + #[redact(fixed = 8)] + pub token: String, + pub role: String, + pub session_name: String, + pub region: Option, + pub endpoint: Option, +} + +pub type CredentialFuture<'a, T> = Pin> + 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, + endpoint: Option, + ) -> CredentialFuture<'_, Option>; +} + +#[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, + endpoint: Option, + ) -> CredentialFuture<'_, Option> { + Box::pin(caller_identity(region, endpoint)) + } +} + +pub async fn profile_credentials(name: &str) -> Result { + 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 { + 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, endpoint: Option) -> 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 { + 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 { + 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, + endpoint: Option, +) -> Result, 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>>, + } + + 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, + _endpoint: Option, + ) -> CredentialFuture<'_, Option> { + 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"] + ); + } +} diff --git a/litellm-rust/crates/auth-aws/src/signing.rs b/litellm-rust/crates/auth-aws/src/signing.rs new file mode 100644 index 00000000000..4286a47dbfe --- /dev/null +++ b/litellm-rust/crates/auth-aws/src/signing.rs @@ -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, + pub region: &'a str, + pub service: &'a str, + pub signing_time: SystemTime, +} + +pub fn sign_v4( + request: SigV4Request<'_>, + credentials: &Credentials, +) -> Result, 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, ¶ms) + .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, + 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")); + } +} diff --git a/litellm-rust/crates/auth-aws/tests/credential_state.rs b/litellm-rust/crates/auth-aws/tests/credential_state.rs new file mode 100644 index 00000000000..5d242b0322d --- /dev/null +++ b/litellm-rust/crates/auth-aws/tests/credential_state.rs @@ -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); + +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); +} diff --git a/litellm-rust/crates/auth-aws/tests/signing.rs b/litellm-rust/crates/auth-aws/tests/signing.rs new file mode 100644 index 00000000000..881150880f2 --- /dev/null +++ b/litellm-rust/crates/auth-aws/tests/signing.rs @@ -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, + 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); +} diff --git a/litellm-rust/crates/auth-aws/tests/sts.rs b/litellm-rust/crates/auth-aws/tests/sts.rs new file mode 100644 index 00000000000..826f3160c7d --- /dev/null +++ b/litellm-rust/crates/auth-aws/tests/sts.rs @@ -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, 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::().expect("content length")) + }); + request.len() >= header_end + 4 + content_length.unwrap_or(0) +} + +fn credentials_xml() -> String { + format!( + "acquired-access-keyacquired-secret-keyacquired-session-token{EXPIRATION}" + ) +} + +#[tokio::test] +async fn assume_role_sends_the_exact_policy_inputs_and_parses_credentials() { + let response_body = format!( + "{}arn:aws:sts::123456789012:assumed-role/demo/test-sessionid:test-sessionrequest-id", + 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!( + "{}arn:aws:sts::123456789012:assumed-role/demo/web-sessionid:web-sessionaudienceprovidersubjectrequest-id", + 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 = "request-id".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); +} diff --git a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs index e616fc93b95..f9cb7253372 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs @@ -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(