From b0ea7e5527799042392203a644b3d7170d1a229f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 2 Nov 2026 00:17:32 +0000 Subject: [PATCH] feat(redis): rotating cloud credentials for Redis --- litellm-rust/Cargo.lock | 14 + litellm-rust/crates/auth-aws/src/aws.rs | 27 +- litellm-rust/crates/auth-azure/src/lib.rs | 2 +- litellm-rust/crates/auth-azure/src/resolve.rs | 170 ++++- litellm-rust/crates/gateway/Cargo.toml | 1 + litellm-rust/crates/gateway/src/lib.rs | 16 +- litellm-rust/crates/gateway/src/main.rs | 3 +- litellm-rust/crates/gateway/src/redis.rs | 56 +- litellm-rust/crates/redis-client/AGENTS.md | 3 +- litellm-rust/crates/redis-client/Cargo.toml | 1 + .../crates/redis-client/src/config.rs | 41 +- .../crates/redis-client/src/connection.rs | 52 +- .../crates/redis-client/src/events.rs | 4 + litellm-rust/crates/redis-client/src/lib.rs | 6 +- .../crates/redis-client/tests/client.rs | 2 + .../crates/redis-client/tests/config.rs | 2 +- .../crates/redis-client/tests/connection.rs | 78 ++- .../crates/redis-client/tests/support/mod.rs | 42 +- litellm-rust/crates/redis-connect/AGENTS.md | 13 +- litellm-rust/crates/redis-connect/Cargo.toml | 21 +- .../crates/redis-connect/src/connect.rs | 74 +- .../crates/redis-connect/src/error.rs | 92 ++- .../crates/redis-connect/src/events.rs | 129 ++++ litellm-rust/crates/redis-connect/src/lib.rs | 9 +- .../crates/redis-connect/src/refresh.rs | 636 ++++++++++++++++++ .../crates/redis-connect/src/sources/aws.rs | 188 ++++++ .../crates/redis-connect/src/sources/azure.rs | 64 ++ .../crates/redis-connect/src/sources/gcp.rs | 126 ++++ .../crates/redis-connect/src/sources/mod.rs | 7 + .../crates/redis-connect/tests/connect.rs | 116 +++- .../crates/redis-connect/tests/live/aws.rs | 17 + .../crates/redis-connect/tests/live/azure.rs | 17 + .../crates/redis-connect/tests/live/gcp.rs | 17 + .../crates/redis-connect/tests/live/main.rs | 93 +++ 34 files changed, 2007 insertions(+), 132 deletions(-) create mode 100644 litellm-rust/crates/redis-connect/src/events.rs create mode 100644 litellm-rust/crates/redis-connect/src/refresh.rs create mode 100644 litellm-rust/crates/redis-connect/src/sources/aws.rs create mode 100644 litellm-rust/crates/redis-connect/src/sources/azure.rs create mode 100644 litellm-rust/crates/redis-connect/src/sources/gcp.rs create mode 100644 litellm-rust/crates/redis-connect/src/sources/mod.rs create mode 100644 litellm-rust/crates/redis-connect/tests/live/aws.rs create mode 100644 litellm-rust/crates/redis-connect/tests/live/azure.rs create mode 100644 litellm-rust/crates/redis-connect/tests/live/gcp.rs create mode 100644 litellm-rust/crates/redis-connect/tests/live/main.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 1c9b2d7c472..6ff3ec496c0 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4438,16 +4438,30 @@ dependencies = [ name = "litellm-redis-connect" version = "0.1.0" dependencies = [ + "futures-util", + "litellm-auth-aws", + "litellm-auth-azure", + "litellm-auth-gcp", + "litellm-auth-types", "litellm-circuit-breaker", "litellm-config", + "litellm-core-utils", + "litellm-error", + "litellm-http", "litellm-redis-client", "litellm-redis-testing", + "litellm-secret-value", "rstest", "rustls 0.23.42", + "serde", "serde_json", + "strum", "thiserror 2.0.19", + "time", "tokio", "tracing", + "url", + "veil", ] [[package]] diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index 113d6d3fbab..d32fc76d3e4 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -268,16 +268,35 @@ impl AwsAuthService { ) -> Result { let resolved = config.clone().with_environment(env_lookup); let flow = classify_auth(config, env_lookup); - let cached = self.cached_credentials(&resolved, &flow).await?; + self.fresh_credentials(&resolved, &flow).await + } + + /// The credentials botocore's default chain finds in `region`, without LiteLLM's Bedrock settings + /// (`AWS_ROLE_NAME`, `AWS_PROFILE_NAME`, `AWS_SESSION_NAME`) that `resolve_credentials` also reads. + pub async fn default_chain_credentials(&self, region: &str) -> Result { + let config = AwsAuthConfig { + region_name: Some(region.to_owned()), + ..Default::default() + }; + self.fresh_credentials(&config, &AwsAuthFlow::DefaultChain) + .await + } + + async fn fresh_credentials( + &self, + resolved: &AwsAuthConfig, + flow: &AwsAuthFlow, + ) -> Result { + let cached = self.cached_credentials(resolved, flow).await?; if fresh(&cached) { return Ok(cached); } - self.credentials.invalidate(&resolved).await; - let refetched = self.cached_credentials(&resolved, &flow).await?; + self.credentials.invalidate(resolved).await; + let refetched = self.cached_credentials(resolved, flow).await?; if fresh(&refetched) { return Ok(refetched); } - self.credentials.invalidate(&resolved).await; + self.credentials.invalidate(resolved).await; Err(Error::ExpiredCredentials) } diff --git a/litellm-rust/crates/auth-azure/src/lib.rs b/litellm-rust/crates/auth-azure/src/lib.rs index 9ed505e4160..6319406ce9c 100644 --- a/litellm-rust/crates/auth-azure/src/lib.rs +++ b/litellm-rust/crates/auth-azure/src/lib.rs @@ -3,5 +3,5 @@ mod native; mod resolve; mod types; -pub use resolve::{AzureAuthService, SECRET_NAMES}; +pub use resolve::{AzureAuthService, AzureIdentity, SECRET_NAMES}; pub use types::{AzureAuthInputs, ConfigValue}; diff --git a/litellm-rust/crates/auth-azure/src/resolve.rs b/litellm-rust/crates/auth-azure/src/resolve.rs index 8e52bed5941..764055ab263 100644 --- a/litellm-rust/crates/auth-azure/src/resolve.rs +++ b/litellm-rust/crates/auth-azure/src/resolve.rs @@ -84,12 +84,36 @@ impl AzureAuthService { Self { native } } + pub async fn identity_token( + &self, + identity: &AzureIdentity, + scope: &str, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result, Error> { + let plan = identity_plan(identity, scope, env_lookup)?; + let inputs = AzureAuthInputs::default_credential_for_scope(scope); + Ok(self + .run_plan(plan, &inputs, env_lookup) + .await? + .map(Sourced::into_value)) + } + pub async fn get_azure_ad_token( &self, inputs: &AzureAuthInputs, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result>, Error> { - match select_auth_plan(inputs, env_lookup)? { + self.run_plan(select_auth_plan(inputs, env_lookup)?, inputs, env_lookup) + .await + } + + async fn run_plan( + &self, + plan: AzureCredentialPlan, + inputs: &AzureAuthInputs, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result>, Error> { + match plan { AzureCredentialPlan::Supplied(credential) => Ok(Some(credential)), AzureCredentialPlan::Caller(caller) => { let credential = caller.acquire().await?; @@ -166,6 +190,105 @@ impl AzureAuthService { } } +/// The identity azure-identity builds for one explicit credential class, without LiteLLM's provider settings +/// (`AZURE_AD_TOKEN`, `AZURE_CREDENTIAL`, `AZURE_SCOPE`) that `get_azure_ad_token` also reads. +#[derive(Clone, Debug)] +pub enum AzureIdentity { + ClientSecret { + tenant_id: String, + client_id: String, + client_secret: SecretValue, + }, + ManagedIdentity { + client_id: String, + }, + Default, +} + +fn identity_plan( + identity: &AzureIdentity, + scope: &str, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result { + let deployment = |value: &str| Sourced::new(value.to_owned(), InputSource::Deployment); + let environment = |name: &str| { + env_lookup(name) + .filter(|value| !value.trim().is_empty()) + .map(|value| Sourced::new(value, InputSource::Environment)) + }; + let scope = deployment(scope); + let authority = environment(AZURE_AUTHORITY_HOST_ENV); + match identity { + AzureIdentity::ClientSecret { + tenant_id, + client_id, + client_secret, + } => Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new( + NativeAzureRequest::ClientSecret { + tenant_id: deployment(tenant_id), + client_id: deployment(client_id), + client_secret: Sourced::new(client_secret.clone(), InputSource::Deployment), + scope, + authority, + }, + )?)), + AzureIdentity::ManagedIdentity { client_id } => Ok(AzureCredentialPlan::Native( + ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id: Some(deployment(client_id)), + scope, + selection_source: InputSource::Deployment, + })?, + )), + AzureIdentity::Default => { + let tenant_id = environment(AZURE_TENANT_ID_ENV); + let client_id = environment(AZURE_CLIENT_ID_ENV); + let from_environment = match ( + tenant_id.clone(), + client_id.clone(), + environment(AZURE_CLIENT_SECRET_ENV), + ) { + (Some(tenant_id), Some(client_id), Some(client_secret)) => { + Some(NativeAzureRequest::ClientSecret { + tenant_id, + client_id, + client_secret: client_secret.map(SecretValue::new), + scope: scope.clone(), + authority: authority.clone(), + }) + } + _ => None, + }; + let probe_imds = + env_lookup("IDENTITY_ENDPOINT").is_none() && env_lookup("MSI_ENDPOINT").is_none(); + let AzureCredentialPlan::Chain(chain) = select_native_plan( + Some(Sourced::new( + AzureCredentialType::DefaultAzureCredential, + InputSource::Deployment, + )), + tenant_id, + client_id, + environment(AZURE_FEDERATED_TOKEN_FILE_ENV), + scope, + authority, + InputSource::Environment, + )? + else { + unreachable!("DefaultAzureCredential always plans a chain") + }; + let environment_request = from_environment + .map(ValidatedAzureRequest::new) + .transpose()?; + Ok(AzureCredentialPlan::Chain( + environment_request + .into_iter() + .chain(chain) + .map(|request| request.with_imds_probe(probe_imds)) + .collect(), + )) + } + } +} + pub(crate) fn select_auth_plan( inputs: &AzureAuthInputs, env_lookup: &dyn Fn(&str) -> Option, @@ -808,6 +931,51 @@ mod tests { } } + #[rstest::rstest] + #[case::litellm_provider_settings_are_ignored( + &[("AZURE_AD_TOKEN", "provider-token"), ("AZURE_CREDENTIAL", "ClientSecretCredential")], + &["managed-identity", "developer-tools"] + )] + #[case::environment_client_secret_comes_first( + &[("AZURE_TENANT_ID", "tenant"), ("AZURE_CLIENT_ID", "client"), ("AZURE_CLIENT_SECRET", "secret")], + &["client-secret", "managed-identity", "developer-tools"] + )] + #[case::aks_workload_identity( + &[("AZURE_TENANT_ID", "tenant"), ("AZURE_CLIENT_ID", "client"), ("AZURE_FEDERATED_TOKEN_FILE", "/var/run/token")], + &["workload-identity", "managed-identity", "developer-tools"] + )] + #[tokio::test] + async fn the_default_identity_follows_azure_identity_defaults( + #[case] environment: &[(&str, &str)], + #[case] expected: &[&str], + ) { + let acquirer = Arc::new(ChainAcquirer { + terminal_on: None, + requests: Mutex::new(Vec::new()), + succeed_on: Some("developer-tools"), + }); + let service = AzureAuthService::with_acquirer(acquirer.clone()); + let lookup = |name: &str| { + environment + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| (*value).to_owned()) + }; + + let credential = service + .identity_token( + &super::AzureIdentity::Default, + "https://redis.azure.com/.default", + &lookup, + ) + .await + .unwrap() + .unwrap(); + + assert_eq!(credential.secret().expose(), "chain-token"); + assert_eq!(*acquirer.requests.lock().unwrap(), expected); + } + #[rstest::rstest] #[tokio::test] async fn default_chain_uses_declared_order_and_stops_after_success() { diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index e9f6de2faab..68460afbcb3 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -38,6 +38,7 @@ tower-sessions-moka-store = "0.15.0" [dev-dependencies] litellm-redis-testing.workspace = true +litellm-http = { workspace = true, features = ["test-support"] } base64.workspace = true rstest.workspace = true tempfile.workspace = true diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index ed6cd64d69c..0a0191b0499 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -74,15 +74,23 @@ pub fn install_crypto_provider() { } } -pub async fn build_redis(config: &Config) -> Result { - let environment = secrets::environment_values(&config.environment_variables)?; - let lookup = |name: &str| { +pub async fn build_redis( + config: &Config, + pool: &HttpClientPool, + http: &litellm_http::HttpClientConfig, +) -> Result { + let environment = Arc::new(secrets::environment_values(&config.environment_variables)?); + let lookup: litellm_redis_connect::SharedEnvironment = Arc::new(move |name: &str| { environment .get(name) .map(|value| value.expose().to_owned()) .or_else(|| std::env::var(name).ok()) + }); + let resources = redis::Resources { + http: pool.client(http, ClientVariant::Provider)?, + lookup: lookup.clone(), }; - redis::connect(config, &lookup).await + redis::connect(config, &*lookup, &resources).await } pub async fn build_mcp( diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index 589103191ae..988de4c819a 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -47,7 +47,8 @@ async fn main() -> Result<(), Box> { Logger::new(StderrSink { level }).install_global()?; let config = Config::load(settings.litellm_config)?; let inference = litellm_gateway::build_inference(&config)?; - let _redis = litellm_gateway::build_redis(&config).await?; + let _redis = + litellm_gateway::build_redis(&config, &inference.resources.pool, &inference.http).await?; let ui = match std::env::var_os("LITELLM_UI_PATH") { Some(directory) => { let ui_settings: UiSettings = settings::from_iter(std::env::vars_os())?; diff --git a/litellm-rust/crates/gateway/src/redis.rs b/litellm-rust/crates/gateway/src/redis.rs index 77184d6b299..c15dacfc9cc 100644 --- a/litellm-rust/crates/gateway/src/redis.rs +++ b/litellm-rust/crates/gateway/src/redis.rs @@ -1,15 +1,21 @@ -use std::time::Duration; +use std::{sync::Arc, time::Duration}; -use litellm_circuit_breaker::{Breaker, Config as BreakerConfig, State}; +use litellm_circuit_breaker::Config as BreakerConfig; use litellm_config::{ CacheSettings, Config, Expected, RedisError, RedisSettings, ResolvedRedis, SecretSnapshot, Value, environment_has_target, }; use litellm_core_utils::settings::Lookup; -use litellm_redis_client::RedisClient; +use litellm_redis_client::{RedisClient, RedisEvents}; +use litellm_redis_connect::{HostServices, SharedEnvironment, TracingEvents}; use crate::Error; +pub(crate) struct Resources { + pub(crate) http: litellm_http::Client, + pub(crate) lookup: SharedEnvironment, +} + const ENVIRONMENT_PING_TIMEOUT: Duration = Duration::from_secs(2); type Params = ResolvedRedis; @@ -60,13 +66,23 @@ pub(crate) fn plan(config: &Config, lookup: &dyn Lookup) -> Result }) } -pub async fn connect(config: &Config, lookup: &dyn Lookup) -> Result { +pub(crate) async fn connect( + config: &Config, + lookup: &dyn Lookup, + resources: &Resources, +) -> Result { let plan = plan(config, lookup)?; let breaker = breaker_config(lookup)?; + let events: Arc = Arc::new(TracingEvents::new(lookup)); let client = |params: Params, breaker: Option| -> Result { Ok(litellm_redis_connect::connect( params, - breaker.map(|config| Breaker::new(config, Box::new(log_transition))), + HostServices { + events: Some(events.clone()), + breaker, + http: resources.http.clone(), + environment: resources.lookup.clone(), + }, )?) }; let coordination = match plan.coordination { @@ -162,23 +178,6 @@ async fn reachable(client: RedisClient) -> Option { } } -fn log_transition(transition: &litellm_circuit_breaker::Transition) { - match transition.to { - State::Open { .. } => tracing::warn!( - generation = transition.generation, - "redis circuit breaker opened, fast-failing redis calls" - ), - State::HalfOpen => tracing::info!( - generation = transition.generation, - "redis circuit breaker half open, probing redis" - ), - State::Closed { .. } => tracing::info!( - generation = transition.generation, - "redis circuit breaker closed, redis recovered" - ), - } -} - fn parse_seconds(text: &str) -> Option { text.parse::() .ok() @@ -237,6 +236,7 @@ fn environment_value( #[cfg(test)] mod tests { + use litellm_core_utils::settings::Lookup; use std::{collections::HashMap, time::Duration}; use litellm_config::{Config, RedisAuth, RedisError}; @@ -244,9 +244,19 @@ mod tests { use litellm_redis_testing::RedisServer; use rstest::rstest; - use super::{Params, Plan, Target, connect, plan}; + use std::sync::Arc; + + use super::{Params, Plan, Resources, Target, plan}; use crate::Error; + async fn connect(config: &Config, lookup: &dyn Lookup) -> Result { + let resources = Resources { + http: litellm_http::Client::for_test(), + lookup: Arc::new(|_: &str| None), + }; + super::connect(config, lookup, &resources).await + } + fn host(name: &str, port: u16) -> Params { Params { endpoint: Endpoint::Host(HostEndpoint { diff --git a/litellm-rust/crates/redis-client/AGENTS.md b/litellm-rust/crates/redis-client/AGENTS.md index abb39fb4cd4..8be4af65eab 100644 --- a/litellm-rust/crates/redis-client/AGENTS.md +++ b/litellm-rust/crates/redis-client/AGENTS.md @@ -49,7 +49,8 @@ ## TLS, credentials and Sentinel - `TlsOptions` follows redis-py 5.3.1. It verifies the certificate by default, skips the hostname check by default (`ssl_check_hostname=False`), and trusts the system roots plus `ca_certs` - redis-rs loads the system trust store for every TLS connection and fails if that fails, even with `ca_certs` set. Where the macOS keychain is blocked (agent sandboxes), run the TLS tests with `SSL_CERT_FILE=/etc/ssl/cert.pem` -- `Credentials` are static and sent at connect. URL userinfo wins over them, like redis-py's `from_url` +- `Credentials::Static` are sent at connect, and URL userinfo wins over them, like redis-py's `from_url` +- `Credentials::Provider` wraps a redis-rs `StreamingCredentialsProvider`. redis-rs takes the first item at connect and re-`AUTH`s live connections on every later item, on each cluster node. Providers push, redis-rs never polls. With a provider, redis-rs overwrites URL userinfo at connect, which matches Python stripping `user:pass@` from the URL - Sentinel discovers the master on each new connection and drops it on a transport failure or READONLY, like redis-py's `SentinelConnectionPool` - Sentinel monitors get `sentinel_password`, the data username when that password is set, and the TLS mode and certificates. Python's `_sentinel_auth_kwargs` forwards every other connection kwarg too, and it sends the username even without a sentinel password - A `Debug` impl never shows a secret. Follow `.agents/skills/rust-redaction` diff --git a/litellm-rust/crates/redis-client/Cargo.toml b/litellm-rust/crates/redis-client/Cargo.toml index 3a7699a70d8..eb0ba02576c 100644 --- a/litellm-rust/crates/redis-client/Cargo.toml +++ b/litellm-rust/crates/redis-client/Cargo.toml @@ -15,6 +15,7 @@ redis = { workspace = true, features = [ "connection-manager", "sentinel", "tls-rustls-insecure", + "token-based-authentication", "tokio-rustls-comp", ] } pem = "4.0.0" diff --git a/litellm-rust/crates/redis-client/src/config.rs b/litellm-rust/crates/redis-client/src/config.rs index fa4d1bcbe10..01af9d38879 100644 --- a/litellm-rust/crates/redis-client/src/config.rs +++ b/litellm-rust/crates/redis-client/src/config.rs @@ -1,8 +1,11 @@ -use std::{path::PathBuf, sync::Arc, time::Duration}; +use std::{fmt, path::PathBuf, pin::Pin, sync::Arc, time::Duration}; + +use futures_util::Stream; use redis::{ - ClientTlsConfig, ConnectionAddr, ConnectionInfo, IntoConnectionInfo, RedisConnectionInfo, - TlsCertificates, TlsMode, io::tcp::TcpSettings, + BasicAuth, ClientTlsConfig, ConnectionAddr, ConnectionInfo, IntoConnectionInfo, + RedisConnectionInfo, RedisResult, StreamingCredentialsProvider, TlsCertificates, TlsMode, + io::tcp::TcpSettings, }; use litellm_circuit_breaker::Breaker; @@ -45,13 +48,33 @@ impl Default for Timeouts { } #[derive(Clone, Debug)] -pub struct Credentials { - pub username: Option, - pub password: Option, +pub enum Credentials { + Static { + username: Option, + password: Option, + }, + Provider(CredentialsProvider), +} + +#[derive(Clone)] +pub struct CredentialsProvider(pub Arc); + +impl fmt::Debug for CredentialsProvider { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("CredentialsProvider") + .finish_non_exhaustive() + } +} + +impl StreamingCredentialsProvider for CredentialsProvider { + fn subscribe(&self) -> Pin> + Send + 'static>> { + self.0.subscribe() + } } impl Credentials { - pub const NONE: Self = Self { + pub const NONE: Self = Self::Static { username: None, password: None, }; @@ -173,7 +196,9 @@ pub(crate) fn connection_info( } fn with_static_credentials(info: ConnectionInfo, credentials: &Credentials) -> ConnectionInfo { - let Credentials { username, password } = credentials; + let Credentials::Static { username, password } = credentials else { + return info; + }; let settings = info.redis_settings().clone(); let settings = match (username, settings.username()) { (Some(username), None) => settings.set_username(username.expose()), diff --git a/litellm-rust/crates/redis-client/src/connection.rs b/litellm-rust/crates/redis-client/src/connection.rs index 1292ae5beb0..4563868e47e 100644 --- a/litellm-rust/crates/redis-client/src/connection.rs +++ b/litellm-rust/crates/redis-client/src/connection.rs @@ -44,9 +44,11 @@ impl Connection { timeouts: Timeouts, ) -> Result { match endpoint { - Endpoint::Url { .. } | Endpoint::Host(_) => { - Self::standalone(config::connection_info(endpoint, credentials)?, timeouts) - } + Endpoint::Url { .. } | Endpoint::Host(_) => Self::standalone( + config::connection_info(endpoint, credentials)?, + credentials, + timeouts, + ), Endpoint::Cluster(cluster) => Self::cluster(cluster, credentials, timeouts), Endpoint::Sentinel(sentinel) => Ok(Self::Sentinel(SentinelConnection::new( sentinel, @@ -64,11 +66,19 @@ impl Connection { } } - fn standalone(info: ConnectionInfo, timeouts: Timeouts) -> Result { + fn standalone( + info: ConnectionInfo, + credentials: &Credentials, + timeouts: Timeouts, + ) -> Result { let config = ConnectionManagerConfig::new() .set_connection_timeout(Some(timeouts.connect)) .set_response_timeout(Some(timeouts.response)) .set_number_of_retries(RECONNECT_RETRIES); + let config = match credentials { + Credentials::Provider(provider) => config.set_credentials_provider(provider.clone()), + Credentials::Static { .. } => config, + }; let manager = ConnectionManager::new_lazy_with_config(redis::Client::open(info)?, config)?; Ok(Self::Standalone(manager)) } @@ -97,14 +107,18 @@ impl Connection { .danger_accept_invalid_hostnames(!tls.check_hostname), None => builder, }; - let Credentials { username, password } = credentials; - let builder = match username { - Some(username) => builder.username(username.expose()), - None => builder, - }; - let builder = match password { - Some(password) => builder.password(password.expose()), - None => builder, + let builder = match credentials { + Credentials::Static { username, password } => { + let builder = match username { + Some(username) => builder.username(username.expose()), + None => builder, + }; + match password { + Some(password) => builder.password(password.expose()), + None => builder, + } + } + Credentials::Provider(provider) => builder.set_credentials_provider(provider.clone()), }; // The builder's response timeout bounds the whole request, but with_config overwrites each node's // response timeout with this one even when unset, so it has to be repeated here @@ -176,8 +190,14 @@ impl SentinelConnection { &endpoint.service_name, SentinelServerType::Master, )?; - let builder = match (&endpoint.sentinel_password, &credentials.username) { - (Some(password), Some(username)) => builder + let builder = match (&endpoint.sentinel_password, credentials) { + ( + Some(password), + Credentials::Static { + username: Some(username), + .. + }, + ) => builder .set_client_to_sentinel_username(username.expose()) .set_client_to_sentinel_password(password.expose()), (Some(password), _) => builder.set_client_to_sentinel_password(password.expose()), @@ -192,6 +212,10 @@ impl SentinelConnection { let config = AsyncConnectionConfig::new() .set_connection_timeout(Some(timeouts.connect)) .set_response_timeout(Some(timeouts.response)); + let config = match credentials { + Credentials::Provider(provider) => config.set_credentials_provider(provider.clone()), + Credentials::Static { .. } => config, + }; Ok(Self(Arc::new(SentinelState { discovery: tokio::sync::Mutex::new(builder.build()?), current: Mutex::new(None), diff --git a/litellm-rust/crates/redis-client/src/events.rs b/litellm-rust/crates/redis-client/src/events.rs index e6a400cf87e..46a71085cc4 100644 --- a/litellm-rust/crates/redis-client/src/events.rs +++ b/litellm-rust/crates/redis-client/src/events.rs @@ -1,9 +1,13 @@ use std::time::Duration; +use litellm_circuit_breaker::{State, Transition}; + use crate::Error; pub trait RedisEvents: Send + Sync { fn call(&self, event: &CallEvent); + + fn breaker(&self, transition: &Transition); } #[derive(Clone, Copy, Debug, PartialEq, Eq)] diff --git a/litellm-rust/crates/redis-client/src/lib.rs b/litellm-rust/crates/redis-client/src/lib.rs index d4a2385a143..91e6530cec9 100644 --- a/litellm-rust/crates/redis-client/src/lib.rs +++ b/litellm-rust/crates/redis-client/src/lib.rs @@ -10,7 +10,7 @@ mod script; pub use client::{RedisClient, Ttl}; pub use commands::Expiry; -pub use config::{Credentials, DEFAULT_TIMEOUT, RedisConfig, Timeouts}; +pub use config::{Credentials, CredentialsProvider, DEFAULT_TIMEOUT, RedisConfig, Timeouts}; pub use error::{ConnectError, Error}; pub use events::{CallEvent, FailureKind, Operation, Outcome, RedisEvents}; pub use litellm_redis_types::{ @@ -18,5 +18,7 @@ pub use litellm_redis_types::{ }; pub use namespace::Namespace; pub use pipeline::Pipeline; -pub use redis::{ErrorKind as RedisErrorKind, RedisError, Value}; +pub use redis::{ + BasicAuth, ErrorKind as RedisErrorKind, RedisError, StreamingCredentialsProvider, Value, +}; pub use script::Script; diff --git a/litellm-rust/crates/redis-client/tests/client.rs b/litellm-rust/crates/redis-client/tests/client.rs index 5ba20002aed..ea4e59b6922 100644 --- a/litellm-rust/crates/redis-client/tests/client.rs +++ b/litellm-rust/crates/redis-client/tests/client.rs @@ -292,6 +292,8 @@ impl RedisEvents for Recorded { fn call(&self, event: &CallEvent) { self.0.lock().unwrap().push(*event); } + + fn breaker(&self, _transition: &litellm_circuit_breaker::Transition) {} } #[rstest] diff --git a/litellm-rust/crates/redis-client/tests/config.rs b/litellm-rust/crates/redis-client/tests/config.rs index 206cd906e05..eb118af7195 100644 --- a/litellm-rust/crates/redis-client/tests/config.rs +++ b/litellm-rust/crates/redis-client/tests/config.rs @@ -5,7 +5,7 @@ use litellm_secret_value::SecretValue; #[rstest] #[case::url(format!("{:?}", Endpoint::Url { url: SecretValue::new("redis://user:credential-value@cache:6379/0"), tls: TlsOptions::verified() }))] -#[case::static_password(format!("{:?}", Credentials { username: None, password: Some(SecretValue::new("credential-value")) }))] +#[case::static_password(format!("{:?}", Credentials::Static { username: None, password: Some(SecretValue::new("credential-value")) }))] #[case::sentinel_password(format!("{:?}", Endpoint::Sentinel(SentinelEndpoint { sentinels: Vec::new(), service_name: "primary".to_owned(), diff --git a/litellm-rust/crates/redis-client/tests/connection.rs b/litellm-rust/crates/redis-client/tests/connection.rs index c8e414a4856..462dd630134 100644 --- a/litellm-rust/crates/redis-client/tests/connection.rs +++ b/litellm-rust/crates/redis-client/tests/connection.rs @@ -1,18 +1,84 @@ mod support; -use std::time::Duration; +use std::{collections::BTreeSet, sync::Arc, time::Duration}; use litellm_redis_testing::{ - RedisDeployment, RedisSentinel, RedisServer, SERVICE_NAME, Topology, eventually, + ClientConnection, RedisDeployment, RedisSentinel, RedisServer, SERVICE_NAME, Topology, + eventually, }; use rstest::rstest; use litellm_redis_client::{ - ClusterNode, Credentials, Endpoint, Error, HostEndpoint, Namespace, RedisClient, RedisConfig, - SentinelEndpoint, Timeouts, + ClusterNode, Credentials, CredentialsProvider, Endpoint, Error, HostEndpoint, Namespace, + RedisClient, RedisConfig, SentinelEndpoint, Timeouts, }; use litellm_secret_value::SecretValue; -use support::connect_as; +use support::{Rotating, connect_as, connections}; + +#[rstest] +#[tokio::test] +async fn a_pushed_credential_reauthenticates_the_live_connection_and_a_reconnect_uses_the_latest( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + server + .cli_on_every_node(&[ + "ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + server + .cli_on_every_node(&[ + "ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + let provider = Arc::new(Rotating::new("alice", "first")); + let client = connect_as( + &server, + Credentials::Provider(CredentialsProvider(provider.clone())), + ); + client.set("k", "1", Duration::from_secs(60)).await.unwrap(); + let before = connections(&server).await; + + provider.rotate("bob", "second"); + let rotated = eventually(async || { + let now = connections(&server).await; + (users(&now) == BTreeSet::from(["bob"])).then_some(now) + }) + .await; + server + .cli_on_every_node(&["ACL", "DELUSER", "alice"]) + .await + .unwrap(); + server + .cli_on_every_node(&["CLIENT", "KILL", "USER", "bob"]) + .await + .unwrap(); + let after_reconnect = eventually(async || client.get::("k").await.ok().flatten()).await; + + let reconnected = connections(&server).await; + + assert_eq!(users(&before), BTreeSet::from(["alice"])); + assert_eq!(ids(&rotated), ids(&before)); + assert_eq!(after_reconnect, "1"); + assert!(!reconnected.is_empty()); + assert_eq!(users(&reconnected), BTreeSet::from(["bob"])); +} + +fn users(connections: &[ClientConnection]) -> BTreeSet<&str> { + connections + .iter() + .map(|connection| connection.user.as_str()) + .collect() +} + +fn ids(connections: &[ClientConnection]) -> Vec<(u16, u64)> { + connections + .iter() + .map(|connection| (connection.node, connection.id)) + .collect() +} #[rstest] #[tokio::test] @@ -48,7 +114,7 @@ async fn a_rejected_credential_is_a_transport_failure() { let with_password = |password: &str| { connect_as( &server, - Credentials { + Credentials::Static { username: None, password: Some(SecretValue::new(password)), }, diff --git a/litellm-rust/crates/redis-client/tests/support/mod.rs b/litellm-rust/crates/redis-client/tests/support/mod.rs index 258fdcc9bc3..dc7d95e4d26 100644 --- a/litellm-rust/crates/redis-client/tests/support/mod.rs +++ b/litellm-rust/crates/redis-client/tests/support/mod.rs @@ -1,14 +1,16 @@ #![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset -use std::time::Duration; +use std::{pin::Pin, time::Duration}; +use futures_util::Stream; use litellm_circuit_breaker::Breaker; use litellm_redis_client::{ - ClusterEndpoint, ClusterNode, Credentials, Endpoint, Expiry, Namespace, RedisClient, - RedisConfig, Timeouts, TlsOptions, + BasicAuth, ClusterEndpoint, ClusterNode, Credentials, Endpoint, Expiry, Namespace, RedisClient, + RedisConfig, RedisError, StreamingCredentialsProvider, Timeouts, TlsOptions, }; use litellm_redis_testing::{ClientConnection, RedisDeployment}; use litellm_secret_value::SecretValue; +use tokio::sync::watch; pub const KEEP: Expiry = Expiry::KeepExisting(Duration::from_secs(60)); @@ -67,3 +69,37 @@ pub fn connect_as(deployment: &RedisDeployment, credentials: Credentials) -> Red pub async fn connections(deployment: &RedisDeployment) -> Vec { deployment.clients().await.unwrap() } + +pub struct Rotating(watch::Sender); + +impl Rotating { + pub fn new(username: &str, password: &str) -> Self { + Self(watch::Sender::new(BasicAuth::new( + username.to_owned(), + password.to_owned(), + ))) + } + + pub fn rotate(&self, username: &str, password: &str) { + self.0 + .send_replace(BasicAuth::new(username.to_owned(), password.to_owned())); + } +} + +impl StreamingCredentialsProvider for Rotating { + fn subscribe( + &self, + ) -> Pin> + Send + 'static>> { + let receiver = self.0.subscribe(); + Box::pin(futures_util::stream::unfold( + (receiver, true), + |(mut receiver, first)| async move { + if !first { + receiver.changed().await.ok()?; + } + let credentials = receiver.borrow_and_update().clone(); + Some((Ok(credentials), (receiver, false))) + }, + )) + } +} diff --git a/litellm-rust/crates/redis-connect/AGENTS.md b/litellm-rust/crates/redis-connect/AGENTS.md index 5df3d299e6e..e1d3cdbf723 100644 --- a/litellm-rust/crates/redis-connect/AGENTS.md +++ b/litellm-rust/crates/redis-connect/AGENTS.md @@ -1,3 +1,10 @@ -- Turns a `litellm_config::ResolvedRedis` into a `RedisClient`, the same way for every host. The endpoint passes through unchanged because config already resolves to `litellm-redis-types`. `redis-client` stays cloud-agnostic, so credentials are mapped here -- `RedisAuth::Host` (a Python `credential_provider` object) fails at connect with an error naming the key, since no Rust host can build one. Cloud auth (GCP, Azure, AWS) fails the same way until its token sources land, so a config that needs it never silently runs without Redis -- Tests run against real Redis from `litellm-redis-testing` +- Turns a `litellm_config::ResolvedRedis` into a `RedisClient`, the same way for every host. The endpoint passes through unchanged because config already resolves to `litellm-redis-types`. `redis-client` stays cloud-agnostic, so credential providers and cloud token sources are built here +- The host passes only what it alone can build in `HostServices`: the event sink, the breaker config, the HTTP client and the environment lookup. Each cloud source asks its auth crate for one explicit identity, the way the Python SDK builds it: `AwsAuthService::default_chain_credentials` (botocore's default chain), `CredentialSource::application_default` (`google.auth.default`) and `AzureIdentity` (the class `_build_azure_credential` picks). LiteLLM's provider settings such as `AWS_ROLE_NAME`, `AZURE_AD_TOKEN` or `VERTEXAI_CREDENTIALS` belong to Bedrock, Azure OpenAI and Vertex and never reach those constructors. `RedisAuth::Host` (a Python `credential_provider` object) fails at connect with an error naming the key, since no host can supply one yet +- Every expiring credential goes through the one refresh loop in `refresh.rs`. A cloud source only fetches. Never add a second loop or let a source sleep +- AWS SDK identity settings must be in the process environment. `aws-config` cannot take the host lookup, so config-only overrides for profiles, role/web-identity settings, credentials, files and metadata sources fail before connecting rather than silently using another identity. Equal process/config values are allowed. Bedrock-only settings remain unrelated to Redis +- redis-rs subscribes once for the connection handshake and again for re-authentication, so a new connection sends a redundant initial AUTH. Unchanged refreshes do not trigger another AUTH. Do not skip a provider's first or alternating subscriptions, because concurrent connections and retries each need the current token. Removing the initial duplicate requires an upstream connection-level fix +- The loop starts on the first `subscribe()`, which redis-rs calls inside its runtime, so `connect` is safe off a runtime +- A token-only credential (Azure without `REDIS_USERNAME`, GCP) is sent as `AUTH default `, because redis-rs always sends a username with a credentials provider. On Redis 6+ that is the same as Python's `AUTH ` +- Tests for the refresh loop use a fake source and real Redis ACL users. Cloud sources are proven against the live service, unit tests only pin request shapes we own +- Rotation tests must retain the same server-side connection IDs across re-AUTH, then force a reconnect after removing the old user and verify the new identity still works +- Live cloud tests live in `tests/live/`. Set `LITELLM_REDIS_AWS_LIVE_CONFIG`, `LITELLM_REDIS_AZURE_LIVE_CONFIG` or `LITELLM_REDIS_GCP_LIVE_CONFIG` to a JSON file containing `redis` settings and `hold_seconds`. Set the waiting period from the deployment's token lifetime when checking expiry. Optional fixtures return early when unset, and configured tests run under plain `cargo test`. An unconfigured run never counts as live proof diff --git a/litellm-rust/crates/redis-connect/Cargo.toml b/litellm-rust/crates/redis-connect/Cargo.toml index 4300c9ea088..7484182e440 100644 --- a/litellm-rust/crates/redis-connect/Cargo.toml +++ b/litellm-rust/crates/redis-connect/Cargo.toml @@ -6,15 +6,32 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-auth-aws.workspace = true +litellm-auth-azure.workspace = true +litellm-auth-gcp.workspace = true +litellm-auth-types.workspace = true +litellm-error.workspace = true +litellm-secret-value.workspace = true +strum.workspace = true litellm-circuit-breaker.workspace = true litellm-config.workspace = true +litellm-core-utils.workspace = true +litellm-http.workspace = true litellm-redis-client.workspace = true +futures-util.workspace = true +serde.workspace = true +serde_json.workspace = true thiserror.workspace = true +veil.workspace = true +time.workspace = true +url.workspace = true +tokio = { workspace = true, features = ["sync", "time", "macros"] } tracing.workspace = true [dev-dependencies] +litellm-http = { workspace = true, features = ["test-support"] } litellm-redis-testing.workspace = true rustls.workspace = true -serde_json.workspace = true rstest.workspace = true -tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } +serde_json.workspace = true +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "test-util"] } diff --git a/litellm-rust/crates/redis-connect/src/connect.rs b/litellm-rust/crates/redis-connect/src/connect.rs index 63edf9ab381..f4b9f4eb06a 100644 --- a/litellm-rust/crates/redis-connect/src/connect.rs +++ b/litellm-rust/crates/redis-connect/src/connect.rs @@ -1,38 +1,72 @@ -use litellm_circuit_breaker::Breaker; +use std::sync::Arc; + +use litellm_circuit_breaker::{Breaker, Config as BreakerConfig}; use litellm_config::{RedisAuth, ResolvedRedis}; +use litellm_core_utils::settings::Lookup; use litellm_redis_client::{ - Credentials, DEFAULT_TIMEOUT, Namespace, RedisClient, RedisConfig, Timeouts, + Credentials, CredentialsProvider, DEFAULT_TIMEOUT, Namespace, RedisClient, RedisConfig, + RedisEvents, Timeouts, }; -use crate::Error; +use crate::{ + Error, RefreshingCredentials, + sources::{AzureAd, ElastiCacheIam, GcpIam}, +}; -pub fn connect(resolved: ResolvedRedis, breaker: Option) -> Result { +pub type SharedEnvironment = Arc; + +pub struct HostServices { + pub events: Option>, + pub breaker: Option, + pub http: litellm_http::Client, + pub environment: SharedEnvironment, +} + +pub fn connect(resolved: ResolvedRedis, host: HostServices) -> Result { tracing::info!(endpoint = ?resolved.endpoint, namespace = ?resolved.namespace, "redis client configured"); + let credentials = credentials(resolved.auth, &host)?; + let events = host.events; + let breaker = host.breaker.map(|config| { + let observer = events.clone(); + Breaker::new( + config, + Box::new(move |transition| { + if let Some(events) = &observer { + events.breaker(transition); + } + }), + ) + }); Ok(RedisClient::new(RedisConfig { endpoint: resolved.endpoint, - credentials: credentials(resolved.auth)?, + credentials, namespace: resolved.namespace.map(Namespace::new).unwrap_or_default(), timeouts: Timeouts::new( resolved.timeout.unwrap_or(DEFAULT_TIMEOUT), resolved.connect_timeout, ), breaker, - events: None, + events, })?) } -fn credentials(auth: RedisAuth) -> Result { - match auth { - RedisAuth::Static { username, password } => Ok(Credentials { username, password }), - RedisAuth::Host => Err(Error::NoHostCredentialProvider), - RedisAuth::GcpIam { .. } => Err(Error::CloudAuth { - setting: "gcp_service_account", - }), - RedisAuth::AzureAd(_) => Err(Error::CloudAuth { - setting: "azure_redis_ad_token", - }), - RedisAuth::AwsIam(_) => Err(Error::CloudAuth { - setting: "aws_iam_auth", - }), - } +fn credentials(auth: RedisAuth, host: &HostServices) -> Result { + let refreshing = |provider: RefreshingCredentials| { + Credentials::Provider(CredentialsProvider(Arc::new(provider))) + }; + Ok(match auth { + RedisAuth::Static { username, password } => Credentials::Static { username, password }, + RedisAuth::Host => return Err(Error::NoHostCredentialProvider), + RedisAuth::AzureAd(auth) => refreshing(RefreshingCredentials::new(AzureAd::new( + auth, + host.environment.clone(), + ))), + RedisAuth::AwsIam(auth) => refreshing(RefreshingCredentials::new(ElastiCacheIam::new( + auth, + &*host.environment, + )?)), + RedisAuth::GcpIam { service_account } => refreshing(RefreshingCredentials::new( + GcpIam::new(host.http.clone(), service_account, host.environment.clone()), + )), + }) } diff --git a/litellm-rust/crates/redis-connect/src/error.rs b/litellm-rust/crates/redis-connect/src/error.rs index 779dd47924b..b9e3f1ad5b5 100644 --- a/litellm-rust/crates/redis-connect/src/error.rs +++ b/litellm-rust/crates/redis-connect/src/error.rs @@ -1,4 +1,4 @@ -#[derive(Debug, thiserror::Error)] +#[derive(veil::Redact, thiserror::Error)] pub enum Error { #[error(transparent)] Connect(#[from] litellm_redis_client::ConnectError), @@ -7,7 +7,93 @@ pub enum Error { )] NoHostCredentialProvider, #[error( - "Redis setting {setting} asks for cloud authentication, which this host does not support yet" + "Redis credential setting {setting} must be set in the process environment, not only environment_variables" )] - CloudAuth { setting: &'static str }, + ProcessEnvironmentRequired { setting: &'static str }, + #[error("{stage} failed{}", http_status_suffix(*http_status))] + Credential { + stage: CredentialStage, + http_status: Option, + #[source] + #[redact(with = "[REDACTED]")] + cause: litellm_error::Source, + }, + #[error("credential acquisition exceeded {timeout:?}")] + CredentialTimeout { + timeout: std::time::Duration, + #[source] + cause: tokio::time::error::Elapsed, + }, + #[error("Redis credential has expired")] + ExpiredCredential, + #[error("no {provider} credential is available for Redis")] + MissingCredential { provider: &'static str }, +} + +#[derive(Clone, Copy, Debug, strum::Display, strum::VariantArray)] +pub enum CredentialStage { + #[strum(serialize = "caller identity")] + Identity, + #[strum(serialize = "Redis token signing")] + Signing, + #[strum(serialize = "Redis token exchange")] + Exchange, + #[strum(serialize = "Redis token response")] + Response, +} + +impl Error { + pub(crate) fn credential( + stage: CredentialStage, + ) -> impl FnOnce(E) -> Self { + move |cause| Self::Credential { + stage, + http_status: None, + cause: litellm_error::Source::new(cause), + } + } + + pub(crate) fn credential_response( + stage: CredentialStage, + status: u16, + ) -> impl FnOnce(litellm_http::Error) -> Self { + move |cause| Self::Credential { + stage, + http_status: Some(status), + cause: litellm_error::Source::new(cause), + } + } + + pub(crate) fn upstream_status(&self) -> Option { + match self { + Self::Credential { http_status, .. } => *http_status, + Self::Connect(_) + | Self::NoHostCredentialProvider + | Self::ProcessEnvironmentRequired { .. } + | Self::CredentialTimeout { .. } + | Self::ExpiredCredential + | Self::MissingCredential { .. } => None, + } + } + + pub(crate) fn kind(&self) -> &'static str { + match self { + Self::Connect(_) => "connection", + Self::NoHostCredentialProvider => "host-credential-provider", + Self::ProcessEnvironmentRequired { .. } => "process-environment-required", + Self::Credential { stage, .. } => match stage { + CredentialStage::Identity => "caller-identity", + CredentialStage::Signing => "token-signing", + CredentialStage::Exchange => "token-exchange", + CredentialStage::Response => "token-response", + }, + Self::CredentialTimeout { .. } => "acquisition-timeout", + Self::ExpiredCredential => "expired-credential", + Self::MissingCredential { .. } => "missing-credential", + } + } +} + +fn http_status_suffix(status: Option) -> String { + status.map_or_else(String::new, |status| format!(" (HTTP {status})")) } diff --git a/litellm-rust/crates/redis-connect/src/events.rs b/litellm-rust/crates/redis-connect/src/events.rs new file mode 100644 index 00000000000..a556b7d370b --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/events.rs @@ -0,0 +1,129 @@ +use std::{ + sync::Mutex, + time::{Duration, Instant}, +}; + +use litellm_circuit_breaker::{State, Transition}; +use litellm_core_utils::settings::Lookup; +use litellm_redis_client::{CallEvent, FailureKind, Outcome, RedisEvents}; + +const DEFAULT_TIMEOUT_LOG_INTERVAL: Duration = Duration::from_secs(5); + +pub struct TracingEvents { + throttle: TimeoutLogThrottle, +} + +impl TracingEvents { + pub fn new(lookup: &dyn Lookup) -> Self { + let interval = lookup + .get("REDIS_TIMEOUT_LOG_INTERVAL") + .and_then(|seconds| seconds.parse::().ok()) + .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) + .unwrap_or(DEFAULT_TIMEOUT_LOG_INTERVAL); + Self { + throttle: TimeoutLogThrottle::new(interval), + } + } +} + +impl RedisEvents for TracingEvents { + fn call(&self, event: &CallEvent) { + let Outcome::Failed(kind) = event.outcome else { + return; + }; + match kind { + FailureKind::BreakerOpen => { + tracing::debug!(operation = ?event.operation, "redis call rejected by the open circuit breaker") + } + FailureKind::Timeout => match self.throttle.admit(Instant::now()) { + None => { + tracing::debug!(operation = ?event.operation, duration = ?event.duration, "redis call timed out") + } + Some(0) => { + tracing::error!(operation = ?event.operation, duration = ?event.duration, "redis call timed out") + } + Some(suppressed) => tracing::error!( + operation = ?event.operation, + duration = ?event.duration, + suppressed, + "redis call timed out, more redis timeouts since the previous line were logged at debug" + ), + }, + FailureKind::Transport | FailureKind::Server | FailureKind::Decode => { + tracing::error!(operation = ?event.operation, failure = ?kind, "redis call failed") + } + } + } + + fn breaker(&self, transition: &Transition) { + match transition.to { + State::Open { .. } => tracing::warn!( + generation = transition.generation, + "redis circuit breaker opened, fast-failing redis calls" + ), + State::HalfOpen => tracing::info!( + generation = transition.generation, + "redis circuit breaker half open, probing redis" + ), + State::Closed { .. } => tracing::info!( + generation = transition.generation, + "redis circuit breaker closed, redis recovered" + ), + } + } +} + +struct TimeoutLogThrottle { + interval: Duration, + state: Mutex<(Option, u64)>, +} + +impl TimeoutLogThrottle { + fn new(interval: Duration) -> Self { + Self { + interval, + state: Mutex::new((None, 0)), + } + } + + fn admit(&self, now: Instant) -> Option { + let mut state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let (last_logged_at, suppressed) = &mut *state; + if last_logged_at.is_some_and(|last| now.duration_since(last) < self.interval) { + *suppressed += 1; + return None; + } + *last_logged_at = Some(now); + Some(std::mem::take(suppressed)) + } +} + +#[cfg(test)] +mod tests { + use std::time::{Duration, Instant}; + + use rstest::rstest; + + use super::TimeoutLogThrottle; + + #[rstest] + fn admits_one_line_per_interval_and_counts_the_rest() { + let throttle = TimeoutLogThrottle::new(Duration::from_secs(5)); + let start = Instant::now(); + let at = |seconds: u64| start + Duration::from_secs(seconds); + + let admitted = [ + throttle.admit(at(0)), + throttle.admit(at(1)), + throttle.admit(at(4)), + throttle.admit(at(5)), + throttle.admit(at(6)), + throttle.admit(at(20)), + ]; + + assert_eq!(admitted, [Some(0), None, None, Some(2), None, Some(1)]); + } +} diff --git a/litellm-rust/crates/redis-connect/src/lib.rs b/litellm-rust/crates/redis-connect/src/lib.rs index 733c98069b6..08d7af91fb7 100644 --- a/litellm-rust/crates/redis-connect/src/lib.rs +++ b/litellm-rust/crates/redis-connect/src/lib.rs @@ -1,5 +1,10 @@ mod connect; mod error; +mod events; +mod refresh; +mod sources; -pub use connect::connect; -pub use error::Error; +pub use connect::{HostServices, SharedEnvironment, connect}; +pub use error::{CredentialStage, Error}; +pub use events::TracingEvents; +pub(crate) use refresh::{CredentialSource, Expiring, RefreshingCredentials, SourceFuture}; diff --git a/litellm-rust/crates/redis-connect/src/refresh.rs b/litellm-rust/crates/redis-connect/src/refresh.rs new file mode 100644 index 00000000000..13c03b3f8ce --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/refresh.rs @@ -0,0 +1,636 @@ +use std::{ + future::{Future, pending}, + pin::Pin, + sync::{Arc, Mutex}, + time::{Duration, SystemTime}, +}; + +use futures_util::{Stream, stream}; +use litellm_redis_client::{BasicAuth, RedisError, StreamingCredentialsProvider}; +use litellm_secret_value::SecretValue; +use tokio::{sync::watch, time::Instant}; + +use crate::Error; + +const TOKEN_ONLY_USERNAME: &str = "default"; +const UNKNOWN_LIFETIME: Duration = Duration::from_secs(300); +const MIN_REFRESH: Duration = Duration::from_secs(1); +const ACQUISITION_TIMEOUT: Duration = Duration::from_secs(30); +const FIRST_RETRY: Duration = Duration::from_secs(1); +const MAX_RETRY: Duration = Duration::from_secs(60); + +#[derive(Clone, PartialEq, Eq)] +pub struct Expiring { + pub username: Option, + pub password: SecretValue, + pub expires_at: Option, +} + +pub type SourceFuture<'a> = Pin> + Send + 'a>>; + +pub trait CredentialSource: Send + Sync + 'static { + fn fetch(&self) -> SourceFuture<'_>; +} + +#[derive(Clone)] +enum Latest { + Pending, + Ready { + username: String, + password: SecretValue, + }, + Unavailable(Arc), +} + +impl Latest { + fn is_same_token(&self, other: &Self) -> bool { + matches!( + (self, other), + ( + Self::Ready { username, password }, + Self::Ready { username: other_username, password: other_password }, + ) if username == other_username && password == other_password + ) + } +} + +type Start = Box; + +#[derive(Clone)] +pub struct RefreshingCredentials { + latest: watch::Receiver, + start: Arc>>, +} + +impl RefreshingCredentials { + pub fn new(source: impl CredentialSource) -> Self { + let (sender, latest) = watch::channel(Latest::Pending); + let start: Start = Box::new(move || drop(tokio::spawn(refresh(source, sender)))); + Self { + latest, + start: Arc::new(Mutex::new(Some(start))), + } + } + + fn start_once(&self) { + let start = self + .start + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .take(); + if let Some(start) = start { + start(); + } + } +} + +impl StreamingCredentialsProvider for RefreshingCredentials { + fn subscribe( + &self, + ) -> Pin> + Send + 'static>> { + self.start_once(); + Box::pin(stream::unfold( + (self.latest.clone(), true), + |(mut receiver, first)| async move { + if !first { + receiver.changed().await.ok()?; + } + let latest = receiver + .wait_for(|latest| !matches!(latest, Latest::Pending)) + .await + .ok()? + .clone(); + let item = match latest { + Latest::Ready { username, password } => { + Ok(BasicAuth::new(username, password.expose().to_owned())) + } + Latest::Unavailable(error) => Err(public_error(&error)), + Latest::Pending => unreachable!("wait_for skips Pending"), + }; + Some((item, (receiver, false))) + }, + )) + } +} + +fn public_error(error: &Error) -> RedisError { + RedisError::from(( + litellm_redis_client::RedisErrorKind::AuthenticationFailed, + "fetching Redis credentials failed", + error.to_string(), + )) +} + +struct Fresh { + latest: Latest, + expires_at: Option, + refresh_in: Duration, +} + +fn fresh(credentials: Expiring, now: SystemTime) -> Result { + let remaining = credentials + .expires_at + .map(|expires_at| expires_at.duration_since(now).unwrap_or_default()); + if remaining == Some(Duration::ZERO) { + return Err(Error::ExpiredCredential); + } + Ok(Fresh { + latest: Latest::Ready { + username: credentials + .username + .unwrap_or_else(|| TOKEN_ONLY_USERNAME.to_owned()), + password: credentials.password, + }, + expires_at: remaining.map(|remaining| Instant::now() + remaining), + refresh_in: refresh_after(credentials.expires_at, now), + }) +} + +async fn until_expiry( + future: impl Future, + sender: &watch::Sender, + expires_at: Option, +) -> Option { + let expiration = async { + match expires_at { + Some(deadline) => tokio::time::sleep_until(deadline).await, + None => pending().await, + } + }; + tokio::pin!(future, expiration); + let mut expired = false; + loop { + tokio::select! { + biased; + () = sender.closed() => return None, + () = &mut expiration, if !expired => { + expired = true; + sender.send_replace(Latest::Unavailable(Arc::new(Error::ExpiredCredential))); + } + result = &mut future => return Some(result), + } + } +} + +async fn refresh(source: impl CredentialSource, sender: watch::Sender) { + let mut retry = FIRST_RETRY; + let mut expires_at = None; + loop { + let acquisition = async { + tokio::time::timeout(ACQUISITION_TIMEOUT, source.fetch()) + .await + .map_err(|cause| Error::CredentialTimeout { + timeout: ACQUISITION_TIMEOUT, + cause, + })? + .and_then(|credentials| fresh(credentials, SystemTime::now())) + }; + let Some(result) = until_expiry(acquisition, &sender, expires_at).await else { + return; + }; + let wait = match result { + Ok(Fresh { + latest, + expires_at: deadline, + refresh_in, + }) => { + retry = FIRST_RETRY; + expires_at = deadline; + sender.send_if_modified(|current| { + let changed = !current.is_same_token(&latest); + if changed { + *current = latest; + } + changed + }); + refresh_in + } + Err(error) => { + tracing::warn!(failure = error.kind(), upstream_status = error.upstream_status(), retry_in = ?retry, "fetching Redis credentials failed"); + if !matches!(*sender.borrow(), Latest::Ready { .. }) { + sender.send_replace(Latest::Unavailable(Arc::new(error))); + } + let wait = retry; + retry = (retry * 2).min(MAX_RETRY); + wait + } + }; + if until_expiry(tokio::time::sleep(wait), &sender, expires_at) + .await + .is_none() + { + return; + } + } +} + +fn refresh_after(expires_at: Option, now: SystemTime) -> Duration { + let remaining = expires_at + .map(|expires_at| expires_at.duration_since(now).unwrap_or_default()) + .unwrap_or(UNKNOWN_LIFETIME); + (remaining * 3 / 4).max(MIN_REFRESH) +} + +#[cfg(test)] +mod tests { + use std::{ + collections::BTreeSet, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + time::{Duration, SystemTime}, + }; + + use futures_util::StreamExt; + use litellm_redis_client::StreamingCredentialsProvider; + use litellm_secret_value::SecretValue; + use rstest::rstest; + + use super::{CredentialSource, Expiring, RefreshingCredentials, SourceFuture, refresh_after}; + use litellm_redis_client::{ + ClusterEndpoint, ClusterNode, Credentials, CredentialsProvider, Endpoint, Namespace, + RedisClient, RedisConfig, Timeouts, TlsOptions, + }; + use litellm_redis_testing::{ClientConnection, RedisDeployment, Topology, eventually}; + + #[rstest] + #[case::three_quarters_of_the_lifetime( + Some(Duration::from_secs(900)), + Duration::from_secs(675) + )] + #[case::already_expired(Some(Duration::ZERO), Duration::from_secs(1))] + #[case::unknown_expiry(None, Duration::from_secs(225))] + fn refreshes_well_before_expiry( + #[case] lifetime: Option, + #[case] expected: Duration, + ) { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_000); + + assert_eq!( + refresh_after(lifetime.map(|lifetime| now + lifetime), now), + expected + ); + } + + struct Scripted { + replies: Mutex>>, + calls: Arc, + } + + impl CredentialSource for Scripted { + fn fetch(&self) -> SourceFuture<'_> { + self.calls.fetch_add(1, Ordering::SeqCst); + let reply = self.replies.lock().unwrap().remove(0); + Box::pin(async move { + reply.map_err(|reason| { + crate::Error::credential(crate::error::CredentialStage::Identity)( + std::io::Error::other(reason), + ) + }) + }) + } + } + + fn token(password: &str, lifetime: Duration) -> Result { + Ok(Expiring { + username: None, + password: SecretValue::new(password), + expires_at: Some(SystemTime::now() + lifetime), + }) + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn pushes_each_refresh_and_keeps_the_last_token_through_failures() { + let calls = Arc::new(AtomicUsize::new(0)); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + token("first", Duration::from_secs(8)), + Err("token endpoint down"), + token("second", Duration::from_secs(3600)), + ]), + calls: calls.clone(), + }); + let mut stream = provider.subscribe(); + + let first = stream.next().await.unwrap().unwrap(); + let second = stream.next().await.unwrap().unwrap(); + let late = provider.subscribe().next().await.unwrap().unwrap(); + + assert_eq!((first.username(), first.password()), ("default", "first")); + assert_eq!(second.password(), "second"); + assert_eq!(late.password(), "second"); + assert_eq!(calls.load(Ordering::SeqCst), 3); + } + + #[rstest] + fn building_outside_a_runtime_fetches_nothing_until_a_connection_subscribes() { + let calls = Arc::new(AtomicUsize::new(0)); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![token("first", Duration::from_secs(3600))]), + calls: calls.clone(), + }); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + let before = calls.load(Ordering::SeqCst); + let first = runtime.block_on(async { provider.subscribe().next().await.unwrap().unwrap() }); + + assert_eq!(before, 0); + assert_eq!(first.password(), "first"); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn a_failure_before_any_token_fails_the_connection_instead_of_hanging() { + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + Err("no identity"), + token("late", Duration::from_secs(3600)), + ]), + calls: Arc::new(AtomicUsize::new(0)), + }); + let mut stream = provider.subscribe(); + + assert!(stream.next().await.unwrap().is_err()); + assert_eq!(stream.next().await.unwrap().unwrap().password(), "late"); + } + + struct Blocking { + first: Mutex>, + started: Arc, + cancelled: Arc, + } + + struct Cancellation(Arc); + + impl Drop for Cancellation { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + impl CredentialSource for Blocking { + fn fetch(&self) -> SourceFuture<'_> { + let first = self.first.lock().unwrap().take(); + Box::pin(async move { + if let Some(first) = first { + return Ok(first); + } + let _guard = Cancellation(self.cancelled.clone()); + self.started.notify_one(); + std::future::pending().await + }) + } + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn a_refresh_that_returns_the_same_token_is_not_pushed_again() { + let calls = Arc::new(AtomicUsize::new(0)); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + token("same", Duration::from_secs(8)), + token("same", Duration::from_secs(3600)), + ]), + calls: calls.clone(), + }); + let mut stream = provider.subscribe(); + assert_eq!(stream.next().await.unwrap().unwrap().password(), "same"); + + let pushed = tokio::time::timeout(Duration::from_secs(60), stream.next()).await; + + assert_eq!(calls.load(Ordering::SeqCst), 2); + assert!(pushed.is_err(), "an unchanged token was pushed again"); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn expiration_is_published_while_refresh_is_stalled() { + let started = Arc::new(tokio::sync::Notify::new()); + let provider = RefreshingCredentials::new(Blocking { + first: Mutex::new(Some(token("first", Duration::from_secs(4)).unwrap())), + started: started.clone(), + cancelled: Arc::new(AtomicUsize::new(0)), + }); + let mut stream = provider.subscribe(); + assert_eq!(stream.next().await.unwrap().unwrap().password(), "first"); + let start = tokio::time::Instant::now(); + started.notified().await; + let expired = stream.next().await.unwrap().unwrap_err(); + assert!(start.elapsed() <= Duration::from_secs(5)); + assert!(expired.to_string().contains("expired")); + let fresh = provider.subscribe().next().await.unwrap().unwrap_err(); + assert!(fresh.to_string().contains("expired")); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn initial_acquisition_has_a_bounded_wait() { + let provider = RefreshingCredentials::new(Blocking { + first: Mutex::new(None), + started: Arc::new(tokio::sync::Notify::new()), + cancelled: Arc::new(AtomicUsize::new(0)), + }); + let error = provider.subscribe().next().await.unwrap().unwrap_err(); + assert!(error.to_string().contains("exceeded")); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn dropping_all_subscribers_cancels_inflight_acquisition() { + let started = Arc::new(tokio::sync::Notify::new()); + let cancelled = Arc::new(AtomicUsize::new(0)); + let provider = RefreshingCredentials::new(Blocking { + first: Mutex::new(None), + started: started.clone(), + cancelled: cancelled.clone(), + }); + let stream = provider.subscribe(); + started.notified().await; + drop(stream); + drop(provider); + tokio::task::yield_now().await; + assert_eq!(cancelled.load(Ordering::SeqCst), 1); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn an_expired_fetch_is_rejected_and_a_later_valid_token_recovers() { + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + Ok(Expiring { + username: None, + password: SecretValue::new("expired"), + expires_at: Some(SystemTime::UNIX_EPOCH), + }), + token("recovered", Duration::from_secs(3600)), + ]), + calls: Arc::new(AtomicUsize::new(0)), + }); + let mut stream = provider.subscribe(); + assert!( + stream + .next() + .await + .unwrap() + .unwrap_err() + .to_string() + .contains("expired") + ); + assert_eq!( + stream.next().await.unwrap().unwrap().password(), + "recovered" + ); + } + + #[rstest] + fn credential_failure_keeps_its_typed_cause_but_hides_it_at_the_redis_boundary() { + let error = crate::Error::credential(crate::error::CredentialStage::Identity)( + std::io::Error::new(std::io::ErrorKind::PermissionDenied, "credential-value"), + ); + assert_eq!( + litellm_error::find::(&error) + .unwrap() + .kind(), + std::io::ErrorKind::PermissionDenied + ); + let debug = format!("{error:?}"); + assert!(!debug.contains("credential-value")); + assert!(debug.contains("[REDACTED]")); + let public = super::public_error(&error).to_string(); + assert!(public.contains("caller identity")); + assert!(!public.contains("credential-value")); + } + + #[rstest] + fn credential_http_failures_expose_status_without_exposing_the_response_body() { + let error = crate::Error::credential_response(crate::CredentialStage::Exchange, 503)( + litellm_http::Error::upstream(503, "credential-value".into()), + ); + let public = super::public_error(&error).to_string(); + assert!(public.contains("HTTP 503")); + assert!(!public.contains("credential-value")); + assert_eq!( + litellm_error::find::(&error) + .unwrap() + .upstream_response() + .unwrap() + .0, + 503 + ); + } + + fn endpoint(deployment: &RedisDeployment) -> Endpoint { + match deployment { + RedisDeployment::Standalone(server) => Endpoint::Url { + url: SecretValue::new(server.url()), + tls: TlsOptions::verified(), + }, + RedisDeployment::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint { + nodes: cluster + .nodes() + .map(|(host, port)| ClusterNode { + host: host.to_owned(), + port, + }) + .collect(), + tls: None, + }), + } + } + + fn user(username: &str, password: &str, lifetime: Duration) -> Result { + Ok(Expiring { + username: Some(username.to_owned()), + password: SecretValue::new(password), + expires_at: Some(SystemTime::now() + lifetime), + }) + } + + fn connection_users(connections: &[ClientConnection]) -> BTreeSet<&str> { + connections + .iter() + .map(|connection| connection.user.as_str()) + .collect() + } + + fn connection_ids(connections: &[ClientConnection]) -> BTreeSet<(u16, u64)> { + connections + .iter() + .map(|connection| (connection.node, connection.id)) + .collect() + } + + #[rstest] + #[tokio::test] + async fn a_refreshed_token_reauthenticates_live_connections_before_the_old_one_expires( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = RedisDeployment::start(topology).await.unwrap(); + server + .cli_on_every_node(&[ + "ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + server + .cli_on_every_node(&[ + "ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + user("alice", "first", Duration::from_secs(2)), + user("bob", "second", Duration::from_secs(3600)), + ]), + calls: Arc::new(AtomicUsize::new(0)), + }); + let client = RedisClient::new(RedisConfig { + endpoint: endpoint(&server), + credentials: Credentials::Provider(CredentialsProvider(Arc::new(provider))), + namespace: Namespace::default(), + timeouts: Timeouts::default(), + breaker: None, + events: None, + }) + .unwrap(); + client.set("k", "1", Duration::from_secs(60)).await.unwrap(); + let before = server.clients().await.unwrap(); + + let after = eventually(async || { + let connections = server.clients().await.ok()?; + (!connections.is_empty() && connection_users(&connections) == BTreeSet::from(["bob"])) + .then_some(connections) + }) + .await; + server + .cli_on_every_node(&["ACL", "DELUSER", "alice"]) + .await + .unwrap(); + + assert!(!before.is_empty()); + assert_eq!(connection_users(&before), BTreeSet::from(["alice"])); + assert_eq!(connection_ids(&after), connection_ids(&before)); + assert_eq!( + client.get::("k").await.unwrap().as_deref(), + Some("1") + ); + server + .cli_on_every_node(&["CLIENT", "KILL", "USER", "bob"]) + .await + .unwrap(); + let value = eventually(async || client.get::("k").await.ok().flatten()).await; + let reconnected = server.clients().await.unwrap(); + + assert_eq!(value, "1"); + assert!(!reconnected.is_empty()); + assert_eq!(connection_users(&reconnected), BTreeSet::from(["bob"])); + assert!(connection_ids(&reconnected).is_disjoint(&connection_ids(&after))); + } +} diff --git a/litellm-rust/crates/redis-connect/src/sources/aws.rs b/litellm-rust/crates/redis-connect/src/sources/aws.rs new file mode 100644 index 00000000000..bcc08749710 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/sources/aws.rs @@ -0,0 +1,188 @@ +use std::time::{Duration, SystemTime}; + +use litellm_auth_aws::{AwsAuthService, SigningScope, presign_url}; +use litellm_config::AwsIamAuth; +use litellm_core_utils::settings::Lookup; + +use crate::Error; +use crate::error::CredentialStage; +use litellm_secret_value::SecretValue; + +use crate::{CredentialSource, Expiring, SourceFuture}; + +const SERVICE: &str = "elasticache"; +const TOKEN_LIFETIME: Duration = Duration::from_secs(900); + +const SDK_IDENTITY_ENVIRONMENT: &[&str] = &[ + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "AWS_PROFILE", + "AWS_DEFAULT_PROFILE", + "AWS_SHARED_CREDENTIALS_FILE", + "AWS_CONFIG_FILE", + "AWS_WEB_IDENTITY_TOKEN_FILE", + "AWS_ROLE_ARN", + "AWS_ROLE_SESSION_NAME", + "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", + "AWS_CONTAINER_CREDENTIALS_FULL_URI", + "AWS_CONTAINER_AUTHORIZATION_TOKEN", + "AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE", + "AWS_EC2_METADATA_DISABLED", + "AWS_EC2_METADATA_SERVICE_ENDPOINT", + "AWS_EC2_METADATA_SERVICE_ENDPOINT_MODE", + "AWS_ENDPOINT_URL", + "AWS_ENDPOINT_URL_STS", +]; + +pub(crate) struct ElastiCacheIam { + service: AwsAuthService, + auth: AwsIamAuth, +} + +impl ElastiCacheIam { + pub(crate) fn new(auth: AwsIamAuth, environment: &dyn Lookup) -> Result { + validate_environment(environment, &|name: &str| std::env::var(name).ok())?; + Ok(Self { + service: AwsAuthService::default(), + auth, + }) + } +} + +fn validate_environment(environment: &dyn Lookup, process: &dyn Lookup) -> Result<(), Error> { + match SDK_IDENTITY_ENVIRONMENT.iter().find(|name| { + environment.get(name).is_some_and(|value| { + Some(SecretValue::new(value)) != process.get(name).map(SecretValue::new) + }) + }) { + Some(setting) => Err(Error::ProcessEnvironmentRequired { setting }), + None => Ok(()), + } +} + +pub(crate) fn connect_url(auth: &AwsIamAuth) -> String { + let query = [ + Some(("Action", "connect")), + Some(("User", auth.user_name.as_str())), + auth.serverless + .then_some(("ResourceType", "ServerlessCache")), + ]; + let mut url = String::from("https://"); + url.push_str(&auth.cache_name.to_lowercase()); + url.push_str("/?"); + url.push_str( + &url::form_urlencoded::Serializer::new(String::new()) + .extend_pairs(query.into_iter().flatten()) + .finish(), + ); + url +} + +impl CredentialSource for ElastiCacheIam { + fn fetch(&self) -> SourceFuture<'_> { + Box::pin(async move { + let credentials = self + .service + .default_chain_credentials(&self.auth.region) + .await + .map_err(Error::credential(CredentialStage::Identity))?; + let now = SystemTime::now(); + let signed = presign_url( + &connect_url(&self.auth), + TOKEN_LIFETIME, + SigningScope { + region: &self.auth.region, + service: SERVICE, + credentials: &credentials, + time: now, + }, + ) + .map_err(Error::credential(CredentialStage::Signing))?; + let token_expiry = now + TOKEN_LIFETIME; + Ok(Expiring { + username: Some(self.auth.user_name.clone()), + password: SecretValue::new(signed.trim_start_matches("https://")), + expires_at: Some( + credentials + .expiry() + .map_or(token_expiry, |expiry| expiry.min(token_expiry)), + ), + }) + }) + } +} + +#[cfg(test)] +mod tests { + use litellm_config::AwsIamAuth; + use rstest::rstest; + + use super::{connect_url, validate_environment}; + + #[rstest] + #[case::profile("AWS_PROFILE")] + #[case::role("AWS_ROLE_ARN")] + #[case::web_identity("AWS_WEB_IDENTITY_TOKEN_FILE")] + #[case::source_key("AWS_ACCESS_KEY_ID")] + #[case::source_secret("AWS_SECRET_ACCESS_KEY")] + #[case::session("AWS_SESSION_TOKEN")] + #[case::profile_path("AWS_SHARED_CREDENTIALS_FILE")] + #[case::container("AWS_CONTAINER_CREDENTIALS_FULL_URI")] + fn configured_identity_overrides_cannot_silently_use_an_ambient_identity( + #[case] setting: &'static str, + ) { + let configured = |name: &str| (name == setting).then(|| "credential-value".to_owned()); + let process = |name: &str| (name == setting).then(|| "ambient".to_owned()); + let error = validate_environment(&configured, &process).unwrap_err(); + + assert!(matches!( + error, + crate::Error::ProcessEnvironmentRequired { setting: rejected } if rejected == setting + )); + assert!(!error.to_string().contains("credential-value")); + assert!(!format!("{error:?}").contains("credential-value")); + } + + #[rstest] + #[case::process_environment("AWS_PROFILE", Some("profile"), Some("profile"))] + #[case::ambient_only("AWS_PROFILE", None, Some("profile"))] + #[case::unconfigured("AWS_PROFILE", None, None)] + #[case::unrelated_bedrock_setting("AWS_ROLE_NAME", Some("bedrock"), None)] + fn unchanged_or_unrelated_settings_preserve_the_sdk_chain( + #[case] setting: &str, + #[case] configured: Option<&str>, + #[case] ambient: Option<&str>, + ) { + let configured = |name: &str| { + (name == setting) + .then_some(configured) + .flatten() + .map(str::to_owned) + }; + let ambient = |name: &str| { + (name == setting) + .then_some(ambient) + .flatten() + .map(str::to_owned) + }; + assert!(validate_environment(&configured, &ambient).is_ok()); + } + + #[rstest] + #[case::provisioned(false, "https://cache-1/?Action=connect&User=iam+user%2F1")] + #[case::serverless( + true, + "https://cache-1/?Action=connect&User=iam+user%2F1&ResourceType=ServerlessCache" + )] + fn the_connect_url_matches_python_urlencode(#[case] serverless: bool, #[case] expected: &str) { + let auth = AwsIamAuth { + user_name: "iam user/1".to_owned(), + cache_name: "Cache-1".to_owned(), + region: "us-east-1".to_owned(), + serverless, + }; + + assert_eq!(connect_url(&auth), expected); + } +} diff --git a/litellm-rust/crates/redis-connect/src/sources/azure.rs b/litellm-rust/crates/redis-connect/src/sources/azure.rs new file mode 100644 index 00000000000..f7aa3e60c85 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/sources/azure.rs @@ -0,0 +1,64 @@ +use litellm_auth_azure::{AzureAuthService, AzureIdentity}; +use litellm_auth_types::ResolvedCredential; +use litellm_config::AzureAdAuth; + +use crate::Error; +use crate::error::CredentialStage; + +use crate::{CredentialSource, Expiring, SharedEnvironment, SourceFuture}; + +const AZURE_REDIS_SCOPE: &str = "https://redis.azure.com/.default"; + +pub(crate) struct AzureAd { + service: AzureAuthService, + identity: AzureIdentity, + username: Option, + lookup: SharedEnvironment, +} + +impl AzureAd { + pub(crate) fn new(auth: AzureAdAuth, lookup: SharedEnvironment) -> Self { + // _build_azure_credential in litellm/_redis.py picks the credential class the same way + let identity = match (auth.client_id, auth.tenant_id, auth.client_secret) { + (Some(client_id), Some(tenant_id), Some(client_secret)) => { + AzureIdentity::ClientSecret { + tenant_id, + client_id, + client_secret, + } + } + (Some(client_id), _, _) => AzureIdentity::ManagedIdentity { client_id }, + (None, _, _) => AzureIdentity::Default, + }; + Self { + service: AzureAuthService::default(), + identity, + username: auth.username, + lookup, + } + } +} + +impl CredentialSource for AzureAd { + fn fetch(&self) -> SourceFuture<'_> { + Box::pin(async move { + let credential = self + .service + .identity_token(&self.identity, AZURE_REDIS_SCOPE, &|name| { + self.lookup.get(name) + }) + .await + .map_err(Error::credential(CredentialStage::Identity))? + .ok_or(Error::MissingCredential { provider: "Azure" })?; + let (token, expires_at) = match credential { + ResolvedCredential::AccessToken { token, expires_on } => (token, expires_on), + ResolvedCredential::Static(token) => (token, None), + }; + Ok(Expiring { + username: self.username.clone(), + password: token, + expires_at, + }) + }) + } +} diff --git a/litellm-rust/crates/redis-connect/src/sources/gcp.rs b/litellm-rust/crates/redis-connect/src/sources/gcp.rs new file mode 100644 index 00000000000..9da60fe9c08 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/sources/gcp.rs @@ -0,0 +1,126 @@ +use std::time::SystemTime; + +use litellm_auth_gcp::{CredentialSource as GcpCredentialSource, GcpAuth}; +use serde::Deserialize; + +use crate::Error; +use crate::error::CredentialStage; +use litellm_secret_value::SecretValue; + +use crate::{CredentialSource, Expiring, SharedEnvironment, SourceFuture}; + +const IAM_CREDENTIALS: &str = "https://iamcredentials.googleapis.com/v1"; +const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform"; + +pub(crate) struct GcpIam { + http: litellm_http::Client, + identity: GcpAuth, + service_account: String, + lookup: SharedEnvironment, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct AccessToken { + access_token: String, + expire_time: String, +} + +impl GcpIam { + pub(crate) fn new( + http: litellm_http::Client, + service_account: String, + lookup: SharedEnvironment, + ) -> Self { + Self { + http, + identity: GcpAuth::default(), + service_account, + lookup, + } + } +} + +impl CredentialSource for GcpIam { + fn fetch(&self) -> SourceFuture<'_> { + Box::pin(async move { + let lookup = |name: &str| self.lookup.get(name); + let caller = self + .identity + .access_token(GcpCredentialSource::application_default(&lookup)) + .await + .map_err(Error::credential(CredentialStage::Identity))?; + let response = self + .http + .post(format!( + "{IAM_CREDENTIALS}/{}:generateAccessToken", + self.service_account + )) + .bearer_auth(caller.expose()) + .json(&serde_json::json!({ "scope": [CLOUD_PLATFORM_SCOPE] })) + .send() + .await + .map_err(Error::credential(CredentialStage::Exchange))?; + let status = response.status().as_u16(); + let reply: AccessToken = response + .error_for_status() + .map_err(litellm_http::Error::from) + .map_err(Error::credential_response( + CredentialStage::Exchange, + status, + ))? + .json() + .await + .map_err(litellm_http::Error::from) + .map_err(Error::credential(CredentialStage::Response))?; + let expires_at = time::OffsetDateTime::parse( + &reply.expire_time, + &time::format_description::well_known::Rfc3339, + ) + .map(SystemTime::from) + .map_err(Error::credential(CredentialStage::Response))?; + Ok(Expiring { + username: None, + password: SecretValue::new(reply.access_token), + expires_at: Some(expires_at), + }) + }) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::GcpIam; + use crate::CredentialSource; + + #[rstest] + #[tokio::test] + async fn the_caller_identity_comes_from_application_default_credentials_not_vertex_settings() { + let environment = |name: &str| { + match name { + "VERTEXAI_CREDENTIALS" => { + Some(r#"{"type":"service_account","project_id":"vertex"}"#) + } + "GOOGLE_APPLICATION_CREDENTIALS" => Some("/nonexistent/redis-adc.json"), + _ => None, + } + .map(str::to_owned) + }; + let source = GcpIam::new( + litellm_http::Client::for_test(), + "projects/-/serviceAccounts/redis@project.iam.gserviceaccount.com".to_owned(), + std::sync::Arc::new(environment), + ); + + let Err(error) = source.fetch().await else { + panic!("a missing credentials file must fail"); + }; + + assert_eq!( + litellm_error::find::(&error).map(std::io::Error::kind), + Some(std::io::ErrorKind::NotFound) + ); + } +} diff --git a/litellm-rust/crates/redis-connect/src/sources/mod.rs b/litellm-rust/crates/redis-connect/src/sources/mod.rs new file mode 100644 index 00000000000..988ad3e1d23 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/sources/mod.rs @@ -0,0 +1,7 @@ +mod aws; +mod azure; +mod gcp; + +pub(crate) use aws::ElastiCacheIam; +pub(crate) use azure::AzureAd; +pub(crate) use gcp::GcpIam; diff --git a/litellm-rust/crates/redis-connect/tests/connect.rs b/litellm-rust/crates/redis-connect/tests/connect.rs index 3fe910f5bcb..d8a4a3c325d 100644 --- a/litellm-rust/crates/redis-connect/tests/connect.rs +++ b/litellm-rust/crates/redis-connect/tests/connect.rs @@ -1,8 +1,10 @@ -use std::{collections::HashMap, time::Duration}; +use std::{collections::HashMap, sync::Arc, time::Duration}; use litellm_config::{RedisSettings, ResolvedRedis, SecretSnapshot}; -use litellm_redis_connect::{Error, connect}; -use litellm_redis_testing::{RedisSentinel, RedisServer, RedisTlsServer, SERVICE_NAME, eventually}; +use litellm_redis_connect::{Error, HostServices, connect}; +use litellm_redis_testing::{ + RedisDeployment, RedisSentinel, RedisTlsServer, SERVICE_NAME, Topology, eventually, +}; use rstest::rstest; fn resolved(kwargs: serde_json::Value, environment: &[(&str, &str)]) -> ResolvedRedis { @@ -17,6 +19,15 @@ fn resolved(kwargs: serde_json::Value, environment: &[(&str, &str)]) -> Resolved .unwrap() } +fn services() -> HostServices { + HostServices { + events: None, + breaker: None, + http: litellm_http::Client::for_test(), + environment: Arc::new(|_: &str| None), + } +} + fn port(url: &str) -> u16 { url.rsplit(':').next().unwrap().parse().unwrap() } @@ -38,7 +49,10 @@ async fn python_tls_kwargs_reach_the_server(#[case] with_ca: bool, #[case] conne "socket_timeout": 2, }); - let pinged = connect(resolved(kwargs, &[]), None).unwrap().ping().await; + let pinged = connect(resolved(kwargs, &[]), services()) + .unwrap() + .ping() + .await; assert_eq!(pinged.is_ok(), connects, "{pinged:?}"); } @@ -52,20 +66,31 @@ async fn static_credentials_from_kwargs_or_environment_authenticate( #[case] password: Option<&str>, #[case] environment: &[(&str, &str)], #[case] connects: bool, + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, ) { - let server = RedisServer::start().await.unwrap(); + let server = RedisDeployment::start(topology).await.unwrap(); server - .cli(&["CONFIG", "SET", "requirepass", "right"]) + .cli_on_every_node(&["ACL", "SETUSER", "default", "on", ">right"]) .await .unwrap(); - let kwargs = serde_json::json!({ - "host": "127.0.0.1", - "port": port(server.url()), - "password": password, - "socket_timeout": 2, - }); + let kwargs = match &server { + RedisDeployment::Standalone(server) => serde_json::json!({ + "host": "127.0.0.1", + "port": port(server.url()), + "password": password, + "socket_timeout": 2, + }), + RedisDeployment::Cluster(cluster) => serde_json::json!({ + "startup_nodes": cluster.nodes().map(|(host, port)| serde_json::json!({ + "host": host, + "port": port, + })).collect::>(), + "password": password, + "socket_timeout": 2, + }), + }; - let pinged = connect(resolved(kwargs, environment), None) + let pinged = connect(resolved(kwargs, environment), services()) .unwrap() .ping() .await; @@ -83,7 +108,7 @@ async fn python_sentinel_kwargs_follow_the_master_through_a_failover() { "service_name": SERVICE_NAME, "socket_timeout": 1, }); - let client = connect(resolved(kwargs, &[]), None).unwrap(); + let client = connect(resolved(kwargs, &[]), services()).unwrap(); client .set("before", "1", Duration::from_secs(600)) .await @@ -107,24 +132,49 @@ async fn python_sentinel_kwargs_follow_the_master_through_a_failover() { } #[rstest] -#[case::credential_provider(serde_json::json!({"host": "redis", "credential_provider": true}), &[], None)] -#[case::gcp(serde_json::json!({"host": "redis", "gcp_service_account": "account"}), &[], Some("gcp_service_account"))] -#[case::azure(serde_json::json!({"host": "redis"}), &[("REDIS_AZURE_AD_TOKEN", "true")], Some("azure_redis_ad_token"))] -#[case::aws( - serde_json::json!({"host": "cache", "ssl": true, "aws_iam_auth": true, "aws_iam_user_name": "user", "aws_iam_cache_name": "cache", "aws_iam_region": "us-east-1"}), - &[], - Some("aws_iam_auth") -)] -fn auth_this_host_cannot_build_fails_at_boot_naming_the_setting( - #[case] kwargs: serde_json::Value, - #[case] environment: &[(&str, &str)], - #[case] cloud_setting: Option<&str>, -) { - let error = connect(resolved(kwargs, environment), None).err().unwrap(); +fn a_python_credential_provider_fails_at_boot() { + let error = connect( + resolved( + serde_json::json!({"host": "redis", "credential_provider": true}), + &[], + ), + services(), + ) + .err() + .unwrap(); - match (error, cloud_setting) { - (Error::NoHostCredentialProvider, None) => {} - (Error::CloudAuth { setting }, Some(expected)) => assert_eq!(setting, expected), - (error, _) => panic!("unexpected {error:?}"), - } + assert!( + matches!(error, Error::NoHostCredentialProvider), + "{error:?}" + ); +} + +#[rstest] +#[case::profile("AWS_PROFILE")] +#[case::role("AWS_ROLE_ARN")] +#[case::web_identity("AWS_WEB_IDENTITY_TOKEN_FILE")] +fn config_only_aws_identities_fail_before_connecting(#[case] setting: &'static str) { + let configured = format!("{}-override", std::env::var(setting).unwrap_or_default()); + let host = HostServices { + environment: Arc::new(move |name: &str| (name == setting).then(|| configured.clone())), + ..services() + }; + let error = connect( + resolved( + serde_json::json!({ + "host": "redis", "ssl": true, "aws_iam_auth": true, + "aws_iam_user_name": "user", "aws_iam_cache_name": "cache", + "aws_iam_region": "us-east-1", + }), + &[], + ), + host, + ) + .err() + .expect("must not select ambient credentials"); + + assert!(matches!( + error, + Error::ProcessEnvironmentRequired { setting: rejected } if rejected == setting + )); } diff --git a/litellm-rust/crates/redis-connect/tests/live/aws.rs b/litellm-rust/crates/redis-connect/tests/live/aws.rs new file mode 100644 index 00000000000..da5ec323392 --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/live/aws.rs @@ -0,0 +1,17 @@ +use rstest::{fixture, rstest}; + +use super::{LiveConfig, config, roundtrip}; + +#[fixture] +fn live_config() -> Option { + config("LITELLM_REDIS_AWS_LIVE_CONFIG") +} + +#[rstest] +#[tokio::test] +async fn credentials_authenticate_existing_and_fresh_connections(live_config: Option) { + let Some(config) = live_config else { + return; + }; + roundtrip(config).await; +} diff --git a/litellm-rust/crates/redis-connect/tests/live/azure.rs b/litellm-rust/crates/redis-connect/tests/live/azure.rs new file mode 100644 index 00000000000..a7f678edef4 --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/live/azure.rs @@ -0,0 +1,17 @@ +use rstest::{fixture, rstest}; + +use super::{LiveConfig, config, roundtrip}; + +#[fixture] +fn live_config() -> Option { + config("LITELLM_REDIS_AZURE_LIVE_CONFIG") +} + +#[rstest] +#[tokio::test] +async fn credentials_authenticate_existing_and_fresh_connections(live_config: Option) { + let Some(config) = live_config else { + return; + }; + roundtrip(config).await; +} diff --git a/litellm-rust/crates/redis-connect/tests/live/gcp.rs b/litellm-rust/crates/redis-connect/tests/live/gcp.rs new file mode 100644 index 00000000000..0965a7a9e9a --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/live/gcp.rs @@ -0,0 +1,17 @@ +use rstest::{fixture, rstest}; + +use super::{LiveConfig, config, roundtrip}; + +#[fixture] +fn live_config() -> Option { + config("LITELLM_REDIS_GCP_LIVE_CONFIG") +} + +#[rstest] +#[tokio::test] +async fn credentials_authenticate_existing_and_fresh_connections(live_config: Option) { + let Some(config) = live_config else { + return; + }; + roundtrip(config).await; +} diff --git a/litellm-rust/crates/redis-connect/tests/live/main.rs b/litellm-rust/crates/redis-connect/tests/live/main.rs new file mode 100644 index 00000000000..bd142c0d005 --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/live/main.rs @@ -0,0 +1,93 @@ +mod aws; +mod azure; +mod gcp; + +use std::{path::PathBuf, sync::Arc, time::Duration}; + +use litellm_config::{RedisSettings, SecretSnapshot}; +use litellm_http::{ + ClientVariant, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, + media::PublicDnsResolver, +}; +use litellm_redis_connect::{HostServices, connect}; +use serde::Deserialize; + +#[derive(Deserialize)] +struct LiveConfig { + redis: RedisSettings, + hold_seconds: u64, +} + +fn config(variable: &str) -> Option { + let path = PathBuf::from(std::env::var_os(variable)?); + Some( + serde_json::from_slice(&std::fs::read(path).expect("read live Redis config")) + .expect("parse live Redis config"), + ) +} + +async fn roundtrip(config: LiveConfig) { + if rustls::crypto::CryptoProvider::get_default().is_none() { + let _ = rustls::crypto::ring::default_provider().install_default(); + } + let lookup = |name: &str| std::env::var(name).ok(); + let resolved = config + .redis + .resolve( + &SecretSnapshot::read(config.redis.secret_names(), &lookup), + &lookup, + ) + .expect("resolve live Redis config"); + let http_settings = HttpSettings::from_layers([HttpSettingsLayer::from_environment(&lookup)]); + let pool = HttpClientPool::new(Arc::new(PublicDnsResolver)); + let http = pool + .client( + &Resolution::from(&http_settings).config, + ClientVariant::Provider, + ) + .expect("build production HTTP client"); + let services = || HostServices { + events: None, + breaker: None, + http: http.clone(), + environment: Arc::new(|name: &str| std::env::var(name).ok()), + }; + let client = connect(resolved.clone(), services()).expect("build live Redis client"); + let marker = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let key = format!("litellm:redis-connect:live:{}:{marker}", std::process::id()); + let ttl = Duration::from_secs(config.hold_seconds + 120); + + client.ping().await.expect("authenticate and ping"); + client + .set(&key, "before-refresh", ttl) + .await + .expect("set before refresh"); + assert_eq!( + client.get::(&key).await.unwrap().as_deref(), + Some("before-refresh") + ); + + tokio::time::sleep(Duration::from_secs(config.hold_seconds)).await; + + assert_eq!( + client.get::(&key).await.unwrap().as_deref(), + Some("before-refresh") + ); + let fresh = connect(resolved, services()).expect("build a fresh live Redis client"); + assert_eq!( + fresh.get::(&key).await.unwrap().as_deref(), + Some("before-refresh") + ); + fresh + .set(&key, "after-refresh", ttl) + .await + .expect("set after refresh"); + assert_eq!( + client.get::(&key).await.unwrap().as_deref(), + Some("after-refresh") + ); + client.del(&[&key]).await.expect("delete live test key"); +}