mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-17 10:17:46 +00:00
fix(key_management): enforce upperbound_key_generate_params on /key/regenerate (#26340)
This commit is contained in:
michelligabriele
@@ -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
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user