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

fix(redis): accept every truthy flag and sign serverless ElastiCache caches

Browse source 
The ElastiCache IAM gate read `aws_iam_auth` and `ssl` with a helper that
only accepted the literal string "true", while the kwarg coercion that runs
later accepts "true", "1" and "yes". Type coercion happens after the gate, so
`REDIS_AWS_IAM_AUTH=1` silently skipped IAM auth and `REDIS_SSL=1` made the
"requires TLS" check fail closed on a connection that was in fact TLS. Both
helpers now share `_str_to_bool`.

AWS signs serverless cache tokens with an extra `ResourceType=ServerlessCache`
query parameter, so tokens minted for a serverless cache were rejected. Adds an
`aws_iam_serverless` setting (`REDIS_AWS_IAM_SERVERLESS`) that puts the
parameter into the signed URL, and lowercases the cache name because
ElastiCache lowercases it at creation time.
This commit is contained in:
eugene-yao-zocdoc 2026-09-09 14:43:58 -04:00
parent bf73c49d73
commit c4fc20bcf9
8 changed files with 192 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
litellm/_redis.py
View file

@ -39,6 +39,14 @@ from ._logging import verbose_logger
AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default"
_AWS_IAM_KWARG_NAMES: Final = (
"aws_iam_auth",
"aws_iam_user_name",
"aws_iam_cache_name",
"aws_iam_region",
"aws_iam_serverless",
)
def _unwrapped_init_args(cls: type) -> frozenset[str]:
"""Every parameter on a single class's own ``__init__``, decorator-unwrapped.
@ -76,10 +84,7 @@ def _get_redis_kwargs():
"azure_client_id",
"azure_tenant_id",
"azure_client_secret",
"aws_iam_auth",
"aws_iam_user_name",
"aws_iam_cache_name",
"aws_iam_region",
*_AWS_IAM_KWARG_NAMES,
}
available_args: Final = {x for x in _unwrapped_init_args(redis.Redis) if x not in exclude_args} | include_args
@ -275,22 +280,25 @@ def _redis_kwargs_from_environment():
return return_dict
def _is_true(value: object | None) -> bool:
return value is True or (isinstance(value, str) and value.lower() == "true")
def _coerces_to_true(value: object | None) -> bool:
return _str_to_bool(value) if isinstance(value, str) else bool(value)
def _uses_tls(redis_kwargs: Mapping[str, object]) -> bool:
if redis_kwargs.get("startup_nodes") is not None:
return _is_true(redis_kwargs.get("ssl"))
return _coerces_to_true(redis_kwargs.get("ssl"))
url: Final = redis_kwargs.get("url")
return urlsplit(url).scheme.lower() == "rediss" if isinstance(url, str) else _is_true(redis_kwargs.get("ssl"))
if isinstance(url, str):
return urlsplit(url).scheme.lower() == "rediss"
return _coerces_to_true(redis_kwargs.get("ssl"))
def _build_elasticache_iam_provider(
user_name: object | None,
cache_name: object | None,
region: object | None,
) -> ElastiCacheIAMCredentialProvider:
def _build_elasticache_iam_provider(redis_kwargs: Mapping[str, object]) -> ElastiCacheIAMCredentialProvider:
user_name: Final = redis_kwargs.get("aws_iam_user_name")
cache_name: Final = redis_kwargs.get("aws_iam_cache_name")
region: Final = (
redis_kwargs.get("aws_iam_region") or get_secret_str("AWS_REGION") or get_secret_str("AWS_DEFAULT_REGION")
)
required_settings: Final = (
("aws_iam_user_name", user_name),
("aws_iam_cache_name", cache_name),
@ -304,6 +312,7 @@ def _build_elasticache_iam_provider(
user_name=str(user_name),
cache_name=str(cache_name),
region=str(region),
is_serverless=_coerces_to_true(redis_kwargs.get("aws_iam_serverless")),
)
@ -577,7 +586,7 @@ def _get_redis_client_logic(**env_overrides):
_azure_redis_ad_token: Final = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN")
_azure_ad_enabled: Final = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true"
_aws_iam_enabled: Final = _is_true(redis_kwargs.get("aws_iam_auth"))
_aws_iam_enabled: Final = _coerces_to_true(redis_kwargs.get("aws_iam_auth"))
if _azure_ad_enabled and _gcp_service_account is not None:
verbose_logger.warning(
@ -619,13 +628,7 @@ def _get_redis_client_logic(**env_overrides):
if not _uses_tls(redis_kwargs):
raise ValueError("AWS ElastiCache IAM Redis authentication requires TLS")
verbose_logger.debug("Setting up AWS ElastiCache IAM authentication for Redis.")
redis_kwargs["credential_provider"] = _build_elasticache_iam_provider(
user_name=redis_kwargs.get("aws_iam_user_name"),
cache_name=redis_kwargs.get("aws_iam_cache_name"),
region=redis_kwargs.get("aws_iam_region")
or get_secret_str("AWS_REGION")
or get_secret_str("AWS_DEFAULT_REGION"),
)
redis_kwargs["credential_provider"] = _build_elasticache_iam_provider(redis_kwargs)
redis_kwargs.pop("gcp_service_account", None)
redis_kwargs.pop("gcp_ssl_ca_certs", None)
@ -635,10 +638,8 @@ def _get_redis_client_logic(**env_overrides):
redis_kwargs.pop("azure_client_id", None)
redis_kwargs.pop("azure_tenant_id", None)
redis_kwargs.pop("azure_client_secret", None)
redis_kwargs.pop("aws_iam_auth", None)
redis_kwargs.pop("aws_iam_user_name", None)
redis_kwargs.pop("aws_iam_cache_name", None)
redis_kwargs.pop("aws_iam_region", None)
for aws_iam_key in _AWS_IAM_KWARG_NAMES:
redis_kwargs.pop(aws_iam_key, None)
if redis_kwargs.get("credential_provider") is not None:
redis_kwargs.pop("redis_connect_func", None)

17
litellm/_redis_credential_provider.py
View file

@ -5,7 +5,7 @@ import threading
import time
from collections.abc import Callable
from typing import TYPE_CHECKING, Final, Protocol
from urllib.parse import quote
from urllib.parse import urlencode
from redis.credentials import CredentialProvider
@ -126,6 +126,7 @@ class GCPIAMCredentialProvider(CredentialProvider):
_ELASTICACHE_SERVICE_NAME: Final = "elasticache"
_ELASTICACHE_TOKEN_TTL_SECONDS: Final = 900
_ELASTICACHE_SERVERLESS_RESOURCE_TYPE: Final = "ServerlessCache"
class ElastiCacheIAMCredentialProvider(CredentialProvider):
@ -134,12 +135,14 @@ class ElastiCacheIAMCredentialProvider(CredentialProvider):
user_name: str,
cache_name: str,
region: str,
is_serverless: bool = False,
credentials_resolver: Callable[[], Credentials | None] | None = None,
token_lifetime_seconds: int = _ELASTICACHE_TOKEN_TTL_SECONDS,
) -> None:
self._user_name = user_name
self._cache_name = cache_name
self._cache_name = cache_name.lower()
self._region = region
self._is_serverless = is_serverless
self._credentials_resolver = credentials_resolver or self._resolve_credentials
self._credentials: Credentials | None = None
self._token_lifetime_seconds = token_lifetime_seconds
@ -171,10 +174,14 @@ class ElastiCacheIAMCredentialProvider(CredentialProvider):
"botocore is required for ElastiCache IAM Redis authentication. Install it with: pip install boto3"
) from e
request: Final = AWSRequest(
method="GET",
url=(f"https://{self._cache_name}/?Action=connect&User={quote(self._user_name, safe='')}"),
query: Final = urlencode(
(
("Action", "connect"),
("User", self._user_name),
*((("ResourceType", _ELASTICACHE_SERVERLESS_RESOURCE_TYPE),) if self._is_serverless else ()),
)
)
request: Final = AWSRequest(method="GET", url=f"https://{self._cache_name}/?{query}")
SigV4QueryAuth(
frozen_credentials,
_ELASTICACHE_SERVICE_NAME,

3
litellm/proxy/_types.py
View file

@ -2431,6 +2431,9 @@ class CoordinationRedisParams(LiteLLMPydanticObjectBase):
aws_iam_user_name: str | None = Field(None, description="AWS ElastiCache IAM user name")
aws_iam_cache_name: str | None = Field(None, description="AWS ElastiCache cache name")
aws_iam_region: str | None = Field(None, description="AWS region for ElastiCache IAM authentication")
aws_iam_serverless: bool | str | None = Field(
None, description="the ElastiCache cache is serverless rather than a self-designed cluster"
)
def has_connection_target(self) -> bool:
return any(value is not None for value in (self.host, self.url, self.startup_nodes, self.sentinel_nodes))

9
litellm/types/management_endpoints/cache_settings_endpoints.py
View file

@ -273,4 +273,13 @@ CACHE_SETTINGS_FIELDS: Final[list[CacheSettingsField]] = [
ui_field_name="AWS IAM Region",
redis_type=None,
),
CacheSettingsField(
field_name="aws_iam_serverless",
field_type="Boolean",
field_value=None,
field_description="The ElastiCache cache is serverless rather than a self-designed cluster",
field_default=False,
ui_field_name="AWS IAM Serverless Cache",
redis_type=None,
),
]

8
litellm/types/management_endpoints/coordination_redis_endpoints.py
View file

@ -131,4 +131,12 @@ COORDINATION_REDIS_SETTINGS_FIELDS: Final[list[CoordinationRedisSettingsField]]
ui_field_name="AWS IAM Region",
section="connection",
),
CoordinationRedisSettingsField(
field_name="aws_iam_serverless",
field_type="Boolean",
field_description="The ElastiCache cache is serverless rather than a self-designed cluster",
field_default=False,
ui_field_name="AWS IAM Serverless Cache",
section="connection",
),
]

82
tests/test_litellm/test_redis.py
View file

@ -113,15 +113,25 @@ def test_credential_provider_is_not_environment_derived():
assert "credential_provider" not in mapping.values()
_AWS_IAM_SETTINGS = {
"aws_iam_auth",
"aws_iam_user_name",
"aws_iam_cache_name",
"aws_iam_region",
"aws_iam_serverless",
}
def test_aws_iam_settings_are_environment_derived():
allowed = _get_redis_kwargs()
mapping = _get_redis_env_kwarg_mapping()
assert {"aws_iam_auth", "aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"} <= allowed
assert _AWS_IAM_SETTINGS <= allowed
assert mapping["REDIS_AWS_IAM_AUTH"] == "aws_iam_auth"
assert mapping["REDIS_AWS_IAM_USER_NAME"] == "aws_iam_user_name"
assert mapping["REDIS_AWS_IAM_CACHE_NAME"] == "aws_iam_cache_name"
assert mapping["REDIS_AWS_IAM_REGION"] == "aws_iam_region"
assert mapping["REDIS_AWS_IAM_SERVERLESS"] == "aws_iam_serverless"
def test_sync_direct_preserves_credential_provider_identity(clean_redis_environment):
@ -319,12 +329,15 @@ def test_aws_iam_environment_settings_install_provider(clean_redis_environment,
monkeypatch.setenv("REDIS_AWS_IAM_USER_NAME", "iam-user")
monkeypatch.setenv("REDIS_AWS_IAM_CACHE_NAME", "cache.example.com")
monkeypatch.setenv("REDIS_AWS_IAM_REGION", "us-east-1")
monkeypatch.setenv("REDIS_SSL", "true")
monkeypatch.setenv("REDIS_AWS_IAM_SERVERLESS", "1")
monkeypatch.setenv("REDIS_SSL", "1")
redis_kwargs = _get_redis_client_logic(host="cache.example.com", port=6379)
assert isinstance(redis_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider)
assert not {"aws_iam_auth", "aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"} & redis_kwargs.keys()
provider = redis_kwargs["credential_provider"]
assert isinstance(provider, ElastiCacheIAMCredentialProvider)
assert provider._is_serverless is True
assert not _AWS_IAM_SETTINGS & redis_kwargs.keys()
@pytest.mark.parametrize(
@ -332,6 +345,9 @@ def test_aws_iam_environment_settings_install_provider(clean_redis_environment,
[
pytest.param({"host": "cache.example.com", "port": 6379}, id="host_without_ssl"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": False}, id="host_ssl_false"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "false"}, id="host_ssl_false_string"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "0"}, id="host_ssl_zero_string"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "no"}, id="host_ssl_no_string"),
pytest.param({"url": "redis://cache.example.com:6379", "ssl": True}, id="plaintext_url"),
pytest.param(
{"startup_nodes": [{"host": "cache.example.com", "port": 6379}]},
@ -368,6 +384,13 @@ def test_aws_iam_auth_rejects_non_tls_connections(clean_redis_environment, trans
"transport",
[
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": True}, id="host"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "true"}, id="host_ssl_true_string"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "1"}, id="host_ssl_one_string"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "yes"}, id="host_ssl_yes_string"),
pytest.param(
{"startup_nodes": [{"host": "cache.example.com", "port": 6379}], "ssl": "1"},
id="cluster_ssl_one_string",
),
pytest.param({"url": "rediss://cache.example.com:6379"}, id="url"),
pytest.param(
{
@ -413,7 +436,7 @@ def test_aws_iam_settings_are_removed_for_url_and_static_credentials(clean_redis
assert redis_kwargs["url"] == "rediss://cache.example.com:6380"
assert "username" not in redis_kwargs
assert "password" not in redis_kwargs
assert not {"aws_iam_auth", "aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"} & redis_kwargs.keys()
assert not _AWS_IAM_SETTINGS & redis_kwargs.keys()
@pytest.mark.parametrize("missing", ["aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"])
@ -499,7 +522,7 @@ def test_aws_iam_settings_map_to_distinct_provider_fields(clean_redis_environmen
assert provider._region == "iam-region-value"
@pytest.mark.parametrize("aws_iam_auth", [False, "false"])
@pytest.mark.parametrize("aws_iam_auth", [None, False, "", "false", "0", "no"])
def test_aws_iam_auth_disabled_does_not_install_provider(clean_redis_environment, aws_iam_auth):
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
@ -511,7 +534,52 @@ def test_aws_iam_auth_disabled_does_not_install_provider(clean_redis_environment
)
assert "credential_provider" not in redis_kwargs
assert not {"aws_iam_auth", "aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"} & redis_kwargs.keys()
assert not _AWS_IAM_SETTINGS & redis_kwargs.keys()
@pytest.mark.parametrize("aws_iam_auth", [True, "true", "True", "TRUE", "1", "yes"])
def test_aws_iam_auth_enabled_by_any_truthy_flag(clean_redis_environment, aws_iam_auth):
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
ssl=True,
aws_iam_auth=aws_iam_auth,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache-name",
aws_iam_region="us-east-1",
)
assert isinstance(redis_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider)
@pytest.mark.parametrize(
"aws_iam_serverless, expected",
[
pytest.param(None, False, id="unset"),
pytest.param(False, False, id="bool_false"),
pytest.param("false", False, id="string_false"),
pytest.param("0", False, id="string_zero"),
pytest.param(True, True, id="bool_true"),
pytest.param("true", True, id="string_true"),
pytest.param("1", True, id="string_one"),
],
)
def test_aws_iam_serverless_flag_reaches_the_provider(clean_redis_environment, aws_iam_serverless, expected):
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
ssl=True,
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache-name",
aws_iam_region="us-east-1",
aws_iam_serverless=aws_iam_serverless,
)
provider = redis_kwargs["credential_provider"]
assert isinstance(provider, ElastiCacheIAMCredentialProvider)
assert provider._is_serverless is expected
assert "aws_iam_serverless" not in redis_kwargs
def test_explicit_provider_wins_over_aws_iam(clean_redis_environment):

54
tests/test_litellm/test_redis_credential_provider.py
View file

@ -175,3 +175,57 @@ def test_elasticache_provider_recovers_after_a_failed_resolution():
assert user_name == "iam-user"
assert token
assert resolver.calls == 2
@pytest.mark.parametrize(
"provider_kwargs, expected_resource_type",
[
pytest.param({}, None, id="default_is_self_designed"),
pytest.param({"is_serverless": False}, None, id="self_designed"),
pytest.param({"is_serverless": True}, ["ServerlessCache"], id="serverless"),
],
)
def test_elasticache_provider_signs_resource_type_only_for_serverless(provider_kwargs, expected_resource_type):
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="cache-name",
region="us-east-1",
credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")),
**provider_kwargs,
)
_, token = provider.get_credentials()
query = parse_qs(urlsplit("https://" + token).query)
assert query.get("ResourceType") == expected_resource_type
assert query["X-Amz-Signature"]
def test_elasticache_provider_lowercases_the_cache_name():
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="Mixed-Case-Cache",
region="us-east-1",
credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")),
)
_, token = provider.get_credentials()
assert urlsplit("https://" + token).netloc == "mixed-case-cache"
def test_elasticache_provider_encodes_reserved_characters_in_the_user_name():
user_name = "iam user/with+reserved&chars"
provider = ElastiCacheIAMCredentialProvider(
user_name=user_name,
cache_name="cache-name",
region="us-east-1",
credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")),
)
returned_user_name, token = provider.get_credentials()
query = parse_qs(urlsplit("https://" + token).query)
assert returned_user_name == user_name
assert query["User"] == [user_name]
assert query["Action"] == ["connect"]

5
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -26372,6 +26372,11 @@ export interface components {
* @description AWS region for ElastiCache IAM authentication
*/
aws_iam_region?: string | null;
/**
* Aws Iam Serverless
* @description the ElastiCache cache is serverless rather than a self-designed cluster
*/
aws_iam_serverless?: boolean | string | null;
/**
* Aws Iam User Name
* @description AWS ElastiCache IAM user name