From 16f74fed15f2504fd5d9e428855efb812aea284c Mon Sep 17 00:00:00 2001 From: harish-berri Date: Wed, 29 Apr 2026 18:56:13 +0000 Subject: [PATCH] refactor(auth): ensure cached UserAPIKeyAuth objects are returned as deep copies to prevent mutation Currently the in memory cache uses a hash map which returns references to elements and not deep copies. Since we dont want to enforce this at the interface level, allowing copy creation at call sites necessary. This minimizes performane overhead. - Updated `get_key_object` to return a deep copy of cached UserAPIKeyAuth objects. - Modified `update_valid_token_with_end_user_params` to create a deep copy of the valid token before applying updates. - Added unit tests to verify that original cached tokens remain unmodified after updates. --- litellm/proxy/auth/auth_checks.py | 4 +-- litellm/proxy/auth/user_api_key_auth.py | 8 ++--- .../test_auth_checks_cache_copy.py | 33 +++++++++++++++++++ .../auth/test_custom_auth_end_user_budget.py | 31 +++++++++++++++++ 4 files changed, 69 insertions(+), 7 deletions(-) create mode 100644 tests/proxy_unit_tests/test_auth_checks_cache_copy.py 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