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 • committed by Yujong Lee
parent f85d956076
commit 56ffcf75d8
32 changed files with 1998 additions and 187 deletions

View file

@ -4440,17 +4440,31 @@ 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-redis-types",
"litellm-secret-value",
"rstest",
"rustls 0.23.42",
"serde",
"serde_json",
"strum",
"thiserror 2.0.19",
"time",
"tokio",
"tracing",
"url",
"veil",
]
[[package]]

View file

@ -39,6 +39,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,16 +1,21 @@
use std::time::Duration;
use std::{sync::Arc, time::Duration};
use litellm_circuit_breaker::{Breaker, Config as BreakerConfig, State};
use litellm_config::{
CacheSettings, Config, ConfiguredRedis, Expected, RedisError, RedisSettings, Resolution,
SecretSnapshot, Value, environment_has_target,
CacheSettings, Config, ConfiguredRedis, RedisError, RedisSettings, Resolution, 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 litellm_redis_types::ResolvedRedis;
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;
@ -59,13 +64,25 @@ 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 client = |params: Params, breaker: Option<BreakerConfig>| -> Result<RedisClient, Error> {
let breaker = litellm_redis_connect::breaker_config(lookup)?;
let events: Arc<dyn RedisEvents> = Arc::new(TracingEvents::new(lookup));
let client = |params: Params,
breaker: Option<litellm_circuit_breaker::Config>|
-> 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 {
@ -167,81 +184,9 @@ 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()
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
}
fn breaker_config(lookup: &dyn Lookup) -> Result<Option<BreakerConfig>, Error> {
let enabled = lookup
.get("REDIS_CIRCUIT_BREAKER_ENABLED")
.is_none_or(|value| value.eq_ignore_ascii_case("true"));
if !enabled {
return Ok(None);
}
let defaults = BreakerConfig::default();
Ok(Some(BreakerConfig {
failure_threshold: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD",
Expected::Integer,
|text| text.parse().ok(),
)?
.unwrap_or(defaults.failure_threshold),
recovery_timeout: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT",
Expected::Seconds,
parse_seconds,
)?
.unwrap_or(defaults.recovery_timeout),
timeout_min_duration: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION",
Expected::Seconds,
parse_seconds,
)?
.unwrap_or(defaults.timeout_min_duration),
}))
}
fn environment_value<T>(
lookup: &dyn Lookup,
name: &str,
expected: Expected,
parse: impl Fn(&str) -> Option<T>,
) -> Result<Option<T>, RedisError> {
lookup
.get(name)
.map(|value| {
parse(&value).ok_or_else(|| RedisError::Invalid {
setting: name.to_owned(),
expected,
})
})
.transpose()
}
#[cfg(test)]
mod tests {
use litellm_core_utils::settings::Lookup;
use std::{collections::HashMap, time::Duration};
use litellm_config::{Config, RedisError};
@ -250,9 +195,19 @@ mod tests {
use litellm_redis_types::RedisAuth;
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

@ -16,6 +16,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;
@ -49,13 +52,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,
};
@ -211,7 +234,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

@ -45,9 +45,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,
@ -65,11 +67,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))
}
@ -98,14 +108,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
@ -201,8 +215,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()),
@ -217,6 +237,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

@ -293,6 +293,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]
@ -50,7 +116,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));
@ -69,3 +71,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_redis_types::ResolvedRedis` into a `RedisClient`, the same way for every host. It depends on `litellm-redis-types`, never on `litellm-config`, so a host that resolves Redis some other way connects the same way. `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_redis_types::ResolvedRedis` into a `RedisClient`, the same way for every host. It depends on `litellm-redis-types`, never on `litellm-config`, so a host that resolves Redis some other way connects the same way. `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. The expiry deadline lives only in the published `Latest::Ready`, so expiring it leaves `Ready` and cannot fire again. Never keep a copy of the deadline in the loop
- 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,16 +6,33 @@ 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-redis-types.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-config.workspace = true
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,18 +1,48 @@
use std::time::Duration;
use litellm_circuit_breaker::Breaker;
use std::sync::Arc;
use litellm_circuit_breaker::{Breaker, Config as BreakerConfig};
use litellm_core_utils::settings::Lookup;
use litellm_redis_types::{RedisAuth, ResolvedRedis};
use litellm_redis_client::{
Credentials, DEFAULT_TIMEOUT, Endpoint, Namespace, RedisClient, RedisConfig, Timeouts,
Credentials, CredentialsProvider, DEFAULT_TIMEOUT, Endpoint, Namespace, RedisClient,
RedisConfig, RedisEvents, Timeouts,
};
use litellm_redis_types::{RedisAuth, ResolvedRedis};
use crate::Error;
use crate::{
Error, RefreshingCredentials,
sources::{AzureAd, ElastiCacheIam, GcpIam},
};
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,
}
// litellm/constants.py REDIS_CLUSTER_HEALTH_CHECK_INTERVAL, the default _redis.py sets for cluster clients only
const CLUSTER_HEALTH_CHECK_INTERVAL: Duration = Duration::from_secs(25);
pub fn connect(resolved: ResolvedRedis, breaker: Option<Breaker>) -> Result<RedisClient, Error> {
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);
}
}),
)
});
let health_check = match &resolved.endpoint {
Endpoint::Cluster(_) => resolved
.health_check_interval
@ -22,30 +52,83 @@ pub fn connect(resolved: ResolvedRedis, breaker: Option<Breaker>) -> Result<Redi
.filter(|interval| !interval.is_zero());
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,
health_check,
})?)
}
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()),
)),
})
}
/// The breaker `RedisCache` builds from `REDIS_CIRCUIT_BREAKER_*`, or none when it is disabled.
pub fn breaker_config(lookup: &dyn Lookup) -> Result<Option<BreakerConfig>, Error> {
let enabled = lookup
.get("REDIS_CIRCUIT_BREAKER_ENABLED")
.is_none_or(|value| value.eq_ignore_ascii_case("true"));
if !enabled {
return Ok(None);
}
let defaults = BreakerConfig::default();
Ok(Some(BreakerConfig {
failure_threshold: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD",
|text| text.parse().ok(),
)?
.unwrap_or(defaults.failure_threshold),
recovery_timeout: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT",
parse_seconds,
)?
.unwrap_or(defaults.recovery_timeout),
timeout_min_duration: environment_value(
lookup,
"REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION",
parse_seconds,
)?
.unwrap_or(defaults.timeout_min_duration),
}))
}
fn parse_seconds(text: &str) -> Option<Duration> {
text.parse::<f64>()
.ok()
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
}
fn environment_value<T>(
lookup: &dyn Lookup,
name: &'static str,
parse: impl Fn(&str) -> Option<T>,
) -> Result<Option<T>, Error> {
lookup
.get(name)
.map(|value| parse(&value).ok_or(Error::InvalidBreakerSetting { setting: name }))
.transpose()
}

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),
@ -6,8 +6,96 @@ pub enum Error {
"Redis setting credential_provider names a Python object, which this host cannot supply"
)]
NoHostCredentialProvider,
#[error("invalid Redis circuit breaker setting {setting}")]
InvalidBreakerSetting { setting: &'static str },
#[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, breaker_config, connect};
pub use error::{CredentialStage, Error};
pub use events::TracingEvents;
pub(crate) use refresh::{CredentialSource, Expiring, RefreshingCredentials, SourceFuture};

View file

@ -0,0 +1,747 @@
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,
expires_at: Option<Instant>,
},
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
)
}
fn expires_at(&self) -> Option<Instant> {
match self {
Self::Ready { expires_at, .. } => *expires_at,
Self::Pending | Self::Unavailable(_) => None,
}
}
}
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,
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>,
) -> Option<T> {
tokio::pin!(future);
loop {
let expires_at = sender.borrow().expires_at();
let expiration = async {
match expires_at {
Some(deadline) => tokio::time::sleep_until(deadline).await,
None => pending().await,
}
};
tokio::select! {
biased;
() = sender.closed() => return None,
() = expiration => {
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;
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).await else {
return;
};
let wait = match result {
Ok(Fresh { latest, refresh_in }) => {
retry = FIRST_RETRY;
sender.send_if_modified(|current| {
let changed = !current.is_same_token(&latest);
*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)
.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, SentinelEndpoint, Timeouts, TlsOptions,
};
use litellm_redis_testing::{
ClientConnection, RedisDeployment, RedisSentinel, SERVICE_NAME, 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 expiration_is_published_once_and_later_failures_keep_their_cause() {
let replies = std::iter::once(token("first", Duration::from_secs(8)))
.chain(std::iter::repeat_n(Err("token endpoint down"), 16))
.collect();
let provider = RefreshingCredentials::new(Scripted {
replies: Mutex::new(replies),
calls: Arc::new(AtomicUsize::new(0)),
});
let mut stream = provider.subscribe();
assert_eq!(stream.next().await.unwrap().unwrap().password(), "first");
let failures: Vec<String> = stream
.take(5)
.map(|item| item.unwrap_err().to_string())
.collect()
.await;
assert!(failures[0].contains("expired"), "{failures:?}");
assert!(
failures[1..]
.iter()
.all(|failure| failure.contains("caller identity") && !failure.contains("expired")),
"{failures:?}"
);
}
#[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,
health_check: 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)));
}
/// Cloud credentials rotate while Sentinel fails over: the refreshed user must be what the
/// client presents to the master Sentinel promotes, not the one it connected with first.
#[rstest]
#[tokio::test]
async fn a_refreshed_user_reaches_the_master_sentinel_promotes() {
let sentinel = RedisSentinel::start().await.unwrap();
sentinel
.cli_on_data_nodes(&[
"ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all",
])
.await
.unwrap();
sentinel
.cli_on_data_nodes(&[
"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 (host, port) = sentinel.sentinel();
let client = RedisClient::new(RedisConfig {
endpoint: Endpoint::Sentinel(SentinelEndpoint {
sentinels: vec![ClusterNode {
host: host.to_owned(),
port,
}],
service_name: SERVICE_NAME.to_owned(),
sentinel_password: None,
tls: None,
db: 0,
}),
credentials: Credentials::Provider(CredentialsProvider(Arc::new(provider))),
namespace: Namespace::default(),
timeouts: Timeouts::new(Duration::from_secs(1), None),
breaker: None,
events: None,
health_check: None,
})
.unwrap();
client
.set("before", "1", Duration::from_secs(60))
.await
.unwrap();
let master = sentinel.master_port().await.unwrap();
let app_users = |connections: &[ClientConnection]| -> BTreeSet<String> {
connections
.iter()
.filter(|connection| connection.user != "default")
.map(|connection| connection.user.clone())
.collect()
};
eventually(async || {
let connections = sentinel.clients_on(master).await.ok()?;
(app_users(&connections) == BTreeSet::from(["bob".to_owned()])).then_some(())
})
.await;
sentinel
.cli_on_data_nodes(&["ACL", "DELUSER", "alice"])
.await
.unwrap();
let promoted = sentinel.failover().await.unwrap();
eventually(async || client.set("after", "2", Duration::from_secs(60)).await.ok()).await;
let on_promoted = sentinel.clients_on(promoted).await.unwrap();
assert_eq!(promoted, sentinel.replica_port());
assert_eq!(
app_users(&on_promoted),
BTreeSet::from(["bob".to_owned()]),
"{on_promoted:?}"
);
assert_eq!(
client.get::<String>("after").await.unwrap().as_deref(),
Some("2")
);
}
}

View file

@ -0,0 +1,188 @@
use std::time::{Duration, SystemTime};
use litellm_auth_aws::{AwsAuthService, SigningScope, presign_url};
use litellm_redis_types::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_redis_types::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_redis_types::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, 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 litellm_redis_types::ResolvedRedis;
use rstest::rstest;
@ -20,6 +22,15 @@ fn resolved(kwargs: serde_json::Value, environment: &[(&str, &str)]) -> Resolved
.redis
}
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()
}
@ -41,7 +52,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:?}");
}
@ -55,20 +69,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;
@ -86,7 +111,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
@ -110,24 +135,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");
}

View file

@ -7,7 +7,8 @@ use testcontainers_modules::testcontainers::{
};
use crate::{
Error,
ClientConnection, Error,
deployment::parse_client_list,
exec::{cli, run},
ports::{HOST, start_on_free_ports, wait_for_pong},
standalone::REDIS_TAG,
@ -71,6 +72,21 @@ impl RedisSentinel {
(HOST, self.ports[2])
}
/// Runs one redis-cli command on the primary and the replica, which keep their own ACLs.
pub async fn cli_on_data_nodes(&self, args: &[&str]) -> Result<(), Error> {
for port in [self.ports[0], self.ports[1]] {
run(&self.container, cli(port, args)).await?;
}
Ok(())
}
pub async fn clients_on(&self, port: u16) -> Result<Vec<ClientConnection>, Error> {
Ok(parse_client_list(
port,
&run(&self.container, cli(port, &["CLIENT", "LIST"])).await?,
))
}
pub fn primary_port(&self) -> u16 {
self.ports[0]
}