feat(redis): rotating cloud credentials for Redis
Some checks are pending
LiteLLM Rust DB / db-queries (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run

This commit is contained in:
Devin AI 2026-11-02 00:17:32 +00:00
parent 359ae381e7
commit b0ea7e5527
34 changed files with 2007 additions and 132 deletions

View file

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

View file

@ -268,16 +268,35 @@ impl AwsAuthService {
) -> Result<Credentials, Error> {
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<Credentials, Error> {
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<Credentials, Error> {
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)
}

View file

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

View file

@ -84,12 +84,36 @@ impl AzureAuthService {
Self { native }
}
pub async fn identity_token(
&self,
identity: &AzureIdentity,
scope: &str,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Option<ResolvedCredential>, 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<String> + Sync),
) -> Result<Option<Sourced<ResolvedCredential>>, 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<String> + Sync),
) -> Result<Option<Sourced<ResolvedCredential>>, 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<String>,
) -> Result<AzureCredentialPlan, Error> {
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<String>,
@ -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() {

View file

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

View file

@ -74,15 +74,23 @@ pub fn install_crypto_provider() {
}
}
pub async fn build_redis(config: &Config) -> Result<RedisClients, Error> {
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<RedisClients, Error> {
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(

View file

@ -47,7 +47,8 @@ async fn main() -> Result<(), Box<dyn Error>> {
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())?;

View file

@ -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<Plan, Error>
})
}
pub async fn connect(config: &Config, lookup: &dyn Lookup) -> Result<RedisClients, Error> {
pub(crate) async fn connect(
config: &Config,
lookup: &dyn Lookup,
resources: &Resources,
) -> Result<RedisClients, Error> {
let plan = plan(config, lookup)?;
let breaker = breaker_config(lookup)?;
let events: Arc<dyn RedisEvents> = Arc::new(TracingEvents::new(lookup));
let client = |params: Params, breaker: Option<BreakerConfig>| -> Result<RedisClient, Error> {
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<RedisClient> {
}
}
fn log_transition(transition: &litellm_circuit_breaker::Transition<State>) {
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<Duration> {
text.parse::<f64>()
.ok()
@ -237,6 +236,7 @@ fn environment_value<T>(
#[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<super::RedisClients, Error> {
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 {

View file

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

View file

@ -15,6 +15,7 @@ redis = { workspace = true, features = [
"connection-manager",
"sentinel",
"tls-rustls-insecure",
"token-based-authentication",
"tokio-rustls-comp",
] }
pem = "4.0.0"

View file

@ -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<SecretValue>,
pub password: Option<SecretValue>,
pub enum Credentials {
Static {
username: Option<SecretValue>,
password: Option<SecretValue>,
},
Provider(CredentialsProvider),
}
#[derive(Clone)]
pub struct CredentialsProvider(pub Arc<dyn StreamingCredentialsProvider>);
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<Box<dyn Stream<Item = RedisResult<BasicAuth>> + 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()),

View file

@ -44,9 +44,11 @@ impl Connection {
timeouts: Timeouts,
) -> Result<Self, ConnectError> {
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<Self, ConnectError> {
fn standalone(
info: ConnectionInfo,
credentials: &Credentials,
timeouts: Timeouts,
) -> Result<Self, ConnectError> {
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),

View file

@ -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<State>);
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]

View file

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

View file

@ -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<State>) {}
}
#[rstest]

View file

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

View file

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

View file

@ -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<ClientConnection> {
deployment.clients().await.unwrap()
}
pub struct Rotating(watch::Sender<BasicAuth>);
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<Box<dyn Stream<Item = Result<BasicAuth, RedisError>> + 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)))
},
))
}
}

View file

@ -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 <token>`, because redis-rs always sends a username with a credentials provider. On Redis 6+ that is the same as Python's `AUTH <token>`
- 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

View file

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

View file

@ -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<Breaker>) -> Result<RedisClient, Error> {
pub type SharedEnvironment = Arc<dyn Lookup + Send + Sync>;
pub struct HostServices {
pub events: Option<Arc<dyn RedisEvents>>,
pub breaker: Option<BreakerConfig>,
pub http: litellm_http::Client,
pub environment: SharedEnvironment,
}
pub fn connect(resolved: ResolvedRedis, host: HostServices) -> Result<RedisClient, Error> {
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<Credentials, Error> {
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<Credentials, Error> {
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()),
)),
})
}

View file

@ -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<u16>,
#[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<E: std::error::Error + Send + Sync + 'static>(
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<u16> {
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<u16>) -> String {
status.map_or_else(String::new, |status| format!(" (HTTP {status})"))
}

View file

@ -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::<f64>().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<State>) {
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<Instant>, u64)>,
}
impl TimeoutLogThrottle {
fn new(interval: Duration) -> Self {
Self {
interval,
state: Mutex::new((None, 0)),
}
}
fn admit(&self, now: Instant) -> Option<u64> {
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)]);
}
}

View file

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

View file

@ -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<String>,
pub password: SecretValue,
pub expires_at: Option<SystemTime>,
}
pub type SourceFuture<'a> = Pin<Box<dyn Future<Output = Result<Expiring, Error>> + Send + 'a>>;
pub trait CredentialSource: Send + Sync + 'static {
fn fetch(&self) -> SourceFuture<'_>;
}
#[derive(Clone)]
enum Latest {
Pending,
Ready {
username: String,
password: SecretValue,
},
Unavailable(Arc<Error>),
}
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<dyn FnOnce() + Send>;
#[derive(Clone)]
pub struct RefreshingCredentials {
latest: watch::Receiver<Latest>,
start: Arc<Mutex<Option<Start>>>,
}
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<Box<dyn Stream<Item = Result<BasicAuth, RedisError>> + 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<Instant>,
refresh_in: Duration,
}
fn fresh(credentials: Expiring, now: SystemTime) -> Result<Fresh, Error> {
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<T>(
future: impl Future<Output = T>,
sender: &watch::Sender<Latest>,
expires_at: Option<Instant>,
) -> Option<T> {
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<Latest>) {
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<SystemTime>, 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<Duration>,
#[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<Vec<Result<Expiring, &'static str>>>,
calls: Arc<AtomicUsize>,
}
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<Expiring, &'static str> {
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<Option<Expiring>>,
started: Arc<tokio::sync::Notify>,
cancelled: Arc<AtomicUsize>,
}
struct Cancellation(Arc<AtomicUsize>);
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::<std::io::Error>(&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::<litellm_http::Error>(&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<Expiring, &'static str> {
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::<String>("k").await.unwrap().as_deref(),
Some("1")
);
server
.cli_on_every_node(&["CLIENT", "KILL", "USER", "bob"])
.await
.unwrap();
let value = eventually(async || client.get::<String>("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)));
}
}

View file

@ -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<Self, Error> {
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);
}
}

View file

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

View file

@ -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::<std::io::Error>(&error).map(std::io::Error::kind),
Some(std::io::ErrorKind::NotFound)
);
}
}

View file

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

View file

@ -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::<Vec<_>>(),
"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
));
}

View file

@ -0,0 +1,17 @@
use rstest::{fixture, rstest};
use super::{LiveConfig, config, roundtrip};
#[fixture]
fn live_config() -> Option<LiveConfig> {
config("LITELLM_REDIS_AWS_LIVE_CONFIG")
}
#[rstest]
#[tokio::test]
async fn credentials_authenticate_existing_and_fresh_connections(live_config: Option<LiveConfig>) {
let Some(config) = live_config else {
return;
};
roundtrip(config).await;
}

View file

@ -0,0 +1,17 @@
use rstest::{fixture, rstest};
use super::{LiveConfig, config, roundtrip};
#[fixture]
fn live_config() -> Option<LiveConfig> {
config("LITELLM_REDIS_AZURE_LIVE_CONFIG")
}
#[rstest]
#[tokio::test]
async fn credentials_authenticate_existing_and_fresh_connections(live_config: Option<LiveConfig>) {
let Some(config) = live_config else {
return;
};
roundtrip(config).await;
}

View file

@ -0,0 +1,17 @@
use rstest::{fixture, rstest};
use super::{LiveConfig, config, roundtrip};
#[fixture]
fn live_config() -> Option<LiveConfig> {
config("LITELLM_REDIS_GCP_LIVE_CONFIG")
}
#[rstest]
#[tokio::test]
async fn credentials_authenticate_existing_and_fresh_connections(live_config: Option<LiveConfig>) {
let Some(config) = live_config else {
return;
};
roundtrip(config).await;
}

View file

@ -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<LiveConfig> {
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::<String>(&key).await.unwrap().as_deref(),
Some("before-refresh")
);
tokio::time::sleep(Duration::from_secs(config.hold_seconds)).await;
assert_eq!(
client.get::<String>(&key).await.unwrap().as_deref(),
Some("before-refresh")
);
let fresh = connect(resolved, services()).expect("build a fresh live Redis client");
assert_eq!(
fresh.get::<String>(&key).await.unwrap().as_deref(),
Some("before-refresh")
);
fresh
.set(&key, "after-refresh", ttl)
.await
.expect("set after refresh");
assert_eq!(
client.get::<String>(&key).await.unwrap().as_deref(),
Some("after-refresh")
);
client.del(&[&key]).await.expect("delete live test key");
}