From a286106f4fa6a3c76a58cab869931cafc9592273 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 22:17:11 +0000 Subject: [PATCH] fix(rust): isolate explicit s3 keys from env tokens and treat 403 misses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/cache-s3/src/auth.rs | 75 ++++++++++++++++++++- litellm-rust/crates/cache-s3/src/cache.rs | 10 +-- litellm-rust/crates/cache-s3/tests/cache.rs | 8 +++ litellm/rust_bridge/_native.pyi | 2 + tests/test_litellm_rust/support/s3_stub.py | 8 ++- tests/test_litellm_rust/test_cache.py | 30 ++++++--- 6 files changed, 113 insertions(+), 20 deletions(-) diff --git a/litellm-rust/crates/cache-s3/src/auth.rs b/litellm-rust/crates/cache-s3/src/auth.rs index 8f7eb66059d..b7ca722cea3 100644 --- a/litellm-rust/crates/cache-s3/src/auth.rs +++ b/litellm-rust/crates/cache-s3/src/auth.rs @@ -1,14 +1,22 @@ -use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future}; +use aws_credential_types::{ + Credentials as AwsCredentials, + provider::{ProvideCredentials, error::CredentialsError, future}, +}; use litellm_auth_aws::{AwsAuthConfig, resolve_credentials}; #[derive(Clone)] pub(crate) struct Credentials { config: AwsAuthConfig, + env: fn(&str) -> Option, } impl Credentials { pub(crate) fn new(config: AwsAuthConfig) -> Self { - Self { config } + Self::with_env(config, |name| std::env::var(name).ok()) + } + + pub(crate) fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option) -> Self { + Self { config, env } } } @@ -18,7 +26,19 @@ impl ProvideCredentials for Credentials { Self: 'a, { future::ProvideCredentials::new(async { - resolve_credentials(self.config.clone(), &|name| std::env::var(name).ok()) + if let (Some(access_key_id), Some(secret_access_key)) = ( + self.config.access_key_id.clone(), + self.config.secret_access_key.clone(), + ) { + return Ok(AwsCredentials::new( + access_key_id, + secret_access_key, + self.config.session_token.clone(), + None, + "litellm-s3-cache", + )); + } + resolve_credentials(self.config.clone(), &self.env) .await .map_err(|_| CredentialsError::provider_error("S3 cache authentication failed")) }) @@ -30,3 +50,52 @@ impl std::fmt::Debug for Credentials { f.debug_struct("Credentials").finish_non_exhaustive() } } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn explicit_keys_ignore_an_ambient_session_token() { + let provider = Credentials::with_env( + AwsAuthConfig { + access_key_id: Some("key".to_string()), + secret_access_key: Some("secret".to_string()), + region_name: Some("us-east-1".to_string()), + ..Default::default() + }, + |name| (name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string()), + ); + let credentials = provider.provide_credentials().await.unwrap(); + assert_eq!(credentials.access_key_id(), "key"); + assert_eq!(credentials.secret_access_key(), "secret"); + assert_eq!(credentials.session_token(), None); + } + + #[tokio::test] + async fn explicit_keys_keep_their_session_token() { + let provider = Credentials::new(AwsAuthConfig { + access_key_id: Some("key".to_string()), + secret_access_key: Some("secret".to_string()), + session_token: Some("t".to_string()), + region_name: Some("us-east-1".to_string()), + ..Default::default() + }); + let credentials = provider.provide_credentials().await.unwrap(); + assert_eq!(credentials.session_token(), Some("t")); + } + + #[tokio::test] + async fn environment_keys_resolve_with_their_session_token() { + let provider = Credentials::with_env(AwsAuthConfig::default(), |name| match name { + "AWS_ACCESS_KEY_ID" => Some("env-key".to_string()), + "AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()), + "AWS_SESSION_TOKEN" => Some("env-token".to_string()), + _ => None, + }); + let credentials = provider.provide_credentials().await.unwrap(); + assert_eq!(credentials.access_key_id(), "env-key"); + assert_eq!(credentials.secret_access_key(), "env-secret"); + assert_eq!(credentials.session_token(), Some("env-token")); + } +} diff --git a/litellm-rust/crates/cache-s3/src/cache.rs b/litellm-rust/crates/cache-s3/src/cache.rs index 78e5f39dbcc..9c791f42c3b 100644 --- a/litellm-rust/crates/cache-s3/src/cache.rs +++ b/litellm-rust/crates/cache-s3/src/cache.rs @@ -134,11 +134,13 @@ impl S3Cache { Ok(output) => output, Err(error) => { if let SdkError::ServiceError(service) = &error { + let status = error + .raw_response() + .map(|response| response.status().as_u16()); let not_found = service.err().is_no_such_key() - || error - .raw_response() - .map(|response| response.status().as_u16()) - == Some(404); + || service.err().meta().code() == Some("AccessDenied") + || status == Some(404) + || status == Some(403); if not_found { return Ok(None); } diff --git a/litellm-rust/crates/cache-s3/tests/cache.rs b/litellm-rust/crates/cache-s3/tests/cache.rs index 6ca3eb9a758..9a71656286b 100644 --- a/litellm-rust/crates/cache-s3/tests/cache.rs +++ b/litellm-rust/crates/cache-s3/tests/cache.rs @@ -125,6 +125,13 @@ async fn get_hit_miss_expired_and_invalid_entries() { ) .mount(&server) .await; + Mock::given(method("GET")) + .and(path("/cache-bucket/team/denied")) + .respond_with( + ResponseTemplate::new(403).set_body_string("AccessDenied"), + ) + .mount(&server) + .await; Mock::given(method("GET")) .and(path("/cache-bucket/team/expired")) .respond_with( @@ -147,6 +154,7 @@ async fn get_hit_miss_expired_and_invalid_entries() { Some(json!({"answer": 3})) ); assert_eq!(cache.get_cache("missing", &context).unwrap(), None); + assert_eq!(cache.get_cache("denied", &context).unwrap(), None); assert_eq!(cache.get_cache("expired", &context).unwrap(), None); assert_eq!( cache.get_cache("malformed", &context), diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 6ad0e6be26e..3f747dadc56 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -165,6 +165,8 @@ class _CacheTestHandle: startup_nodes: Sequence[tuple[str, int]] | None = None, ) -> _CacheTestHandle: ... @staticmethod + def azure_blob(account_url: str, container: str) -> _CacheTestHandle: ... + @staticmethod def s3( bucket: str, *, diff --git a/tests/test_litellm_rust/support/s3_stub.py b/tests/test_litellm_rust/support/s3_stub.py index 18863ad8083..7e5d068b975 100644 --- a/tests/test_litellm_rust/support/s3_stub.py +++ b/tests/test_litellm_rust/support/s3_stub.py @@ -51,7 +51,7 @@ class S3Stub: def do_PUT(self) -> None: body: Final = self._read_body() headers: Final = {name: self.headers[name] for name in _STORED_HEADERS if name in self.headers} - stub._objects[self._key()] = S3Object(body=body, headers=headers) + stub._objects = {**stub._objects, self._key(): S3Object(body=body, headers=headers)} self.send_response(200) self.send_header("ETag", '"stub"') self.send_header("Content-Length", "0") @@ -83,7 +83,9 @@ class S3Stub: if send_body: self.wfile.write(entry.body) - def log_message(self, format: str, *args: object) -> None: # noqa: A002 + def log_message( + self, format: str, *args: object + ) -> None: # BaseHTTPRequestHandler.log_message names this parameter format pass self._server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) @@ -100,7 +102,7 @@ class S3Stub: return self._objects def put_object(self, key: str, body: bytes, headers: dict[str, str] | None = None) -> None: - self._objects[key] = S3Object(body=body, headers=headers or {}) + self._objects = {**self._objects, key: S3Object(body=body, headers=headers or {})} def expires(self, key: str) -> object: header: Final = self._objects[key].headers.get("expires") diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index cfa8e95703c..2c98f0ca6d3 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -11,6 +11,7 @@ from collections.abc import Generator from datetime import datetime from types import SimpleNamespace from typing import Final, Protocol, cast +from unittest.mock import Mock from urllib.parse import urlparse import boto3 @@ -24,8 +25,8 @@ import litellm from litellm.caching.azure_blob_cache import AzureBlobCache from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache from litellm.caching.in_memory_cache import InMemoryCache -from litellm.caching.s3_cache import S3Cache from litellm.caching.redis_cluster_cache import RedisClusterCache +from litellm.caching.s3_cache import S3Cache from litellm.rust_bridge import _native from litellm.types.caching import LiteLLMCacheType from tests.test_litellm_rust.support.isolation import rebound @@ -416,15 +417,20 @@ def test_azure_blob_facade_serves_natively_and_python_reads_the_same_blobs(azure assert handle.backend == "azure-blob" account_url: Final = backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}") with pytest.raises(TypeError, match="containers must match"): - _native._CacheTestHandle.azure_blob(account_url, f"{backend.container_client.container_name}-other")._bind_facade( - azure_blob_facade - ) + _native._CacheTestHandle.azure_blob( + account_url, f"{backend.container_client.container_name}-other" + )._bind_facade(azure_blob_facade) handle._bind_facade(azure_blob_facade) resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=azure_blob_facade)) native: Final = resolver.resolve() assert native.kind == "native" - response: Final = {"choices": [{"text": "caf\u00e9 \u2603"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} + response: Final = { + "choices": [{"text": "caf\u00e9 \u2603"}], + "usage": {"total_tokens": 3}, + "flag": True, + "empty": None, + } native.store({**request("sync"), "ttl_seconds": 0.001}, response) native.store(request("sync"), {"choices": [{"text": "second"}]}) time.sleep(0.01) @@ -479,7 +485,9 @@ async def test_azure_blob_native_async_writes_overwrite_batch_and_flush_like_pyt await binding.async_store({**request("async"), "ttl_seconds": 0.001}, {"value": 2}) time.sleep(0.01) assert await binding.async_lookup(request("async")) == {"value": 2} - assert await backend.async_get_cache("async") == json.loads(backend.container_client.download_blob("async").readall()) + assert await backend.async_get_cache("async") == json.loads( + backend.container_client.download_blob("async").readall() + ) assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="async") == {"value": 2} await binding.async_store_batch([request("first"), request("second")], [{"value": 3}, {"value": 4}]) @@ -624,11 +632,11 @@ def test_s3_facade_binds_only_exact_configuration_and_falls_back_on_mutation(s3_ binding: Final = resolver.resolve() assert binding.kind == "native" - calls: Final = [] - facade.cache.s3_client.meta.events.register("before-call.s3.*", lambda **_kwargs: calls.append(1)) + handler: Final = Mock() + facade.cache.s3_client.meta.events.register("before-call.s3.*", handler) binding.store(request("native"), {"answer": 1}) assert binding.lookup(request("native")) == {"answer": 1} - assert calls == [] + assert handler.call_count == 0 assert "team/native" in s3_stub.objects with rebound(facade.cache, "bucket_name", "other"): @@ -756,7 +764,9 @@ async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_n await binding.async_flush() - remaining: Final = tuple(sorted(key for node in client.get_primaries() for key in client.keys("parity:*", target_nodes=node))) + remaining: Final = tuple( + sorted(key for node in client.get_primaries() for key in client.keys("parity:*", target_nodes=node)) + ) assert remaining == (), remaining assert client.get("unscoped") == b"stays" client.delete("unscoped")