From db8ef44323d2a2b166ccca21cb6a3d99c25ee09c Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Sat, 25 Apr 2026 22:49:00 +0200 Subject: [PATCH] fix(key_management): enforce upperbound_key_generate_params on /key/regenerate (#26340) --- .../key_management_endpoints.py | 2 + .../test_key_management_endpoints.py | 324 ++++++++++++++++++ 2 files changed, 326 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 25cc48aa5f..a424f1558f 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3854,6 +3854,8 @@ async def _execute_virtual_key_regeneration( non_default_values = {} if data is not None: + # Enforce upperbound key params on regenerate (don't fill defaults) + _enforce_upperbound_key_params(data, fill_defaults=False) non_default_values = await prepare_key_update_data( data=data, existing_key_row=key_in_db ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 647be49e0d..f256ffd866 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -9046,6 +9046,330 @@ def test_enforce_upperbound_no_config_is_noop(): litellm.upperbound_key_generate_params = original +# --- Tests: _execute_virtual_key_regeneration enforces upperbound --- + + +def _make_regenerate_mock_prisma(): + """Mock prisma client shaped for _execute_virtual_key_regeneration.""" + + class DictLikeResult: + def __init__(self, data): + self._data = data + + def __iter__(self): + return iter(self._data.items()) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock( + return_value=DictLikeResult( + { + "token": "new-hashed-token", + "key_name": "sk-...ab12", + "user_id": "user-1", + } + ) + ) + mock_prisma_client.db.litellm_verificationtoken.create = AsyncMock( + return_value=None + ) + mock_prisma_client.jsonify_object = MagicMock(side_effect=lambda data: data) + return mock_prisma_client + + +def _make_regenerate_user_api_key_dict(): + return UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ) + + +def _make_regenerate_existing_key(): + return LiteLLM_VerificationToken( + token="abc123", + user_id="user-1", + models=["gpt-4"], + team_id=None, + max_budget=None, + tags=None, + ) + + +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(): + """Regenerate must reject durations exceeding upperbound_key_generate_params.duration.""" + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + from litellm.types.proxy.management_endpoints.ui_sso import ( + LiteLLM_UpperboundKeyGenerateParams, + ) + + original = litellm.upperbound_key_generate_params + try: + litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams( + duration="1h" + ) + existing_key = _make_regenerate_existing_key() + data = RegenerateKeyRequest(duration="2h") + user_api_key_dict = _make_regenerate_user_api_key_dict() + mock_prisma_client = _make_regenerate_mock_prisma() + + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + assert exc_info.value.status_code == 400 + assert "duration" in str(exc_info.value.detail) + # Rejected regenerate must not reach the DB update. + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 + finally: + litellm.upperbound_key_generate_params = original + + +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_allows_within_limit_duration(): + """Regenerate must accept durations within upperbound_key_generate_params.duration.""" + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + from litellm.types.proxy.management_endpoints.ui_sso import ( + LiteLLM_UpperboundKeyGenerateParams, + ) + + original = litellm.upperbound_key_generate_params + try: + litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams( + duration="1h" + ) + existing_key = _make_regenerate_existing_key() + data = RegenerateKeyRequest(duration="30m") + user_api_key_dict = _make_regenerate_user_api_key_dict() + mock_prisma_client = _make_regenerate_mock_prisma() + + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + finally: + litellm.upperbound_key_generate_params = original + + +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_rejects_over_limit_max_budget(): + """Regenerate must reject max_budget exceeding upperbound — proves the fix covers non-duration fields.""" + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + from litellm.types.proxy.management_endpoints.ui_sso import ( + LiteLLM_UpperboundKeyGenerateParams, + ) + + original = litellm.upperbound_key_generate_params + try: + litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams( + max_budget=10.0 + ) + existing_key = _make_regenerate_existing_key() + data = RegenerateKeyRequest(max_budget=500.0) + user_api_key_dict = _make_regenerate_user_api_key_dict() + mock_prisma_client = _make_regenerate_mock_prisma() + + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + assert exc_info.value.status_code == 400 + assert "max_budget" in str(exc_info.value.detail) + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 + finally: + litellm.upperbound_key_generate_params = original + + +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_skips_none_values(): + """Regenerate with data.duration=None must not raise, even when upperbound is set + (fill_defaults=False semantic — None means 'inherit from existing key').""" + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + from litellm.types.proxy.management_endpoints.ui_sso import ( + LiteLLM_UpperboundKeyGenerateParams, + ) + + original = litellm.upperbound_key_generate_params + try: + litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams( + duration="1h" + ) + existing_key = _make_regenerate_existing_key() + data = RegenerateKeyRequest() # all fields None + user_api_key_dict = _make_regenerate_user_api_key_dict() + mock_prisma_client = _make_regenerate_mock_prisma() + + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + finally: + litellm.upperbound_key_generate_params = original + + +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_no_upperbound_config_is_noop(): + """Regenerate with no upperbound config set must accept any duration.""" + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + + original = litellm.upperbound_key_generate_params + try: + litellm.upperbound_key_generate_params = None + existing_key = _make_regenerate_existing_key() + data = RegenerateKeyRequest(duration="30d") + user_api_key_dict = _make_regenerate_user_api_key_dict() + mock_prisma_client = _make_regenerate_mock_prisma() + + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + finally: + litellm.upperbound_key_generate_params = original + + class TestAllowedRoutesCallerPermission: """ Non-admins must not be able to set `allowed_routes` on a key. The field