diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 840f64cfede..0c1670e65c6 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2316,7 +2316,7 @@ async def get_key_object( if isinstance(cached_key_obj, dict): return UserAPIKeyAuth(**cached_key_obj) elif isinstance(cached_key_obj, UserAPIKeyAuth): - return cached_key_obj + return cached_key_obj.model_copy(deep=True) if check_cache_only: raise Exception( @@ -2366,7 +2366,7 @@ async def get_key_object( proxy_logging_obj=proxy_logging_obj, ) - return _response + return _response.model_copy(deep=True) @log_db_metrics diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9779a07b97b..ce93fa2c07a 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -303,6 +303,7 @@ async def user_api_key_auth_websocket(websocket: WebSocket): def update_valid_token_with_end_user_params( valid_token: UserAPIKeyAuth, end_user_params: dict ) -> UserAPIKeyAuth: + valid_token = valid_token.model_copy(deep=True) valid_token.end_user_id = end_user_params.get("end_user_id") # Only overwrite token fields when the DB-derived value is not None. # This prevents DB lookups (where the budget table has no value set) @@ -1220,11 +1221,8 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 raise e # update end-user params on valid token # These can change per request - it's important to update them here - valid_token.end_user_id = end_user_params.get("end_user_id") - valid_token.end_user_tpm_limit = end_user_params.get("end_user_tpm_limit") - valid_token.end_user_rpm_limit = end_user_params.get("end_user_rpm_limit") - valid_token.allowed_model_region = end_user_params.get( - "allowed_model_region" + valid_token = update_valid_token_with_end_user_params( + valid_token=valid_token, end_user_params=end_user_params ) # update key budget with temp budget increase valid_token = _update_key_budget_with_temp_budget_increase( diff --git a/tests/proxy_unit_tests/test_auth_checks_cache_copy.py b/tests/proxy_unit_tests/test_auth_checks_cache_copy.py new file mode 100644 index 00000000000..4e40c358d80 --- /dev/null +++ b/tests/proxy_unit_tests/test_auth_checks_cache_copy.py @@ -0,0 +1,33 @@ +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import get_key_object + + +@pytest.mark.asyncio +async def test_get_key_object_returns_copy_of_cached_user_api_key_auth(): + """ + In-memory cache returns object references. Auth callers must receive a copy so + request-scoped mutations do not poison the cached key object. + """ + cached_token = UserAPIKeyAuth( + token="hashed-token", + end_user_id=None, + end_user_tpm_limit=None, + ) + user_api_key_cache = MagicMock() + user_api_key_cache.async_get_cache = AsyncMock(return_value=cached_token) + + result = await get_key_object( + hashed_token="hashed-token", + prisma_client=MagicMock(), + user_api_key_cache=user_api_key_cache, + ) + + assert result is not cached_token + result.end_user_id = "request-user" + result.end_user_tpm_limit = 1 + assert cached_token.end_user_id is None + assert cached_token.end_user_tpm_limit is None diff --git a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py index 3af0cac6fd3..6bba59ae8b0 100644 --- a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py @@ -139,3 +139,34 @@ def test_update_valid_token_db_values_override_custom_auth_when_set(): # DB values should win assert result.end_user_tpm_limit == 500 assert result.end_user_model_max_budget == db_budget + + +def test_update_valid_token_does_not_mutate_original_token(): + """ + Request-scoped end-user limits must not mutate the cached UserAPIKeyAuth object. + """ + valid_token = UserAPIKeyAuth( + token="test_token", + end_user_id=None, + end_user_tpm_limit=None, + end_user_rpm_limit=None, + allowed_model_region=None, + ) + end_user_params = { + "end_user_id": "attacker-user", + "end_user_tpm_limit": 1, + "end_user_rpm_limit": 1, + "allowed_model_region": "eu", + } + + result = update_valid_token_with_end_user_params(valid_token, end_user_params) + + assert result is not valid_token + assert result.end_user_id == "attacker-user" + assert result.end_user_tpm_limit == 1 + assert result.end_user_rpm_limit == 1 + assert result.allowed_model_region == "eu" + assert valid_token.end_user_id is None + assert valid_token.end_user_tpm_limit is None + assert valid_token.end_user_rpm_limit is None + assert valid_token.allowed_model_region is None