scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

fix(proxy): enforce custom_key_update policy on /key/regenerate

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
oliver 2026-09-11 07:41:13 +00:00
parent 362033cb2a
commit e205be80df
2 changed files with 184 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

47
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -294,6 +294,35 @@ def _custom_key_update_hook(
return hooks.user_custom_key_update
async def _enforce_custom_key_update_policy(
hook: Callable[..., Awaitable[Mapping[str, object]]] | None,
data: UpdateKeyRequest,
) -> None:
if hook is None:
return
if not inspect.iscoroutinefunction(hook):
raise ValueError("user_custom_key_update must be a coroutine")
result: Final = await hook(data)
if not result.get("decision", True):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=result.get("message", "Authentication Failed - Custom Auth Rule"),
)
def _regenerate_request_as_update_request(key: str, data: RegenerateKeyRequest) -> UpdateKeyRequest | None:
changed_fields: Final = MappingProxyType(
{
field: value
for field, value in data.model_dump(exclude_unset=True).items()
if field in UpdateKeyRequest.model_fields and field != "key"
}
)
if not changed_fields:
return None
return UpdateKeyRequest(key=key, **changed_fields)
class _LegacyDumpable(Protocol):
def dict(self) -> Mapping[str, object]: ...
@ -3066,19 +3095,7 @@ async def update_key_fn(
user_api_key_cache=user_api_key_cache,
)
# Custom key update hook
custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = _custom_key_update_hook(
proxy_server
)
if custom_key_update_hook is not None:
if inspect.iscoroutinefunction(custom_key_update_hook):
result: Final = await custom_key_update_hook(data)
else:
raise ValueError("user_custom_key_update must be a coroutine")
decision: Final = result.get("decision", True)
message: Final = result.get("message", "Authentication Failed - Custom Auth Rule")
if not decision:
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=message)
await _enforce_custom_key_update_policy(hook=_custom_key_update_hook(proxy_server), data=data)
# Enforce upperbound key params on update (don't fill defaults)
_enforce_upperbound_key_params(data, fill_defaults=False)
@ -5069,6 +5086,7 @@ async def _execute_virtual_key_regeneration(
proxy_logging_obj: ProxyLogging,
) -> GenerateKeyResponse:
"""Generate new token, update DB, invalidate cache, and return response."""
from litellm.proxy import proxy_server
from litellm.proxy.proxy_server import hash_token
# Mirror the /key/update ownership rebind guard. See helper docstring.
@ -5116,6 +5134,9 @@ async def _execute_virtual_key_regeneration(
non_default_values = {}
if data is not None:
update_request: Final = _regenerate_request_as_update_request(key=hashed_api_key, data=data)
if update_request is not None:
await _enforce_custom_key_update_policy(hook=_custom_key_update_hook(proxy_server), data=update_request)
# 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)

150
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -24,18 +24,21 @@ from litellm.proxy._types import (
LitellmUserRoles,
Member,
ProxyException,
RegenerateKeyRequest,
ResetSpendRequest,
UpdateKeyRequest,
)
from litellm.proxy.auth.auth_checks import _delete_cache_key_object, _project_cache_key
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy.management_endpoints.key_management_endpoints import (
_check_org_key_limits,
_check_project_key_limits,
_check_team_key_limits,
_common_key_generation_helper,
_enforce_upperbound_key_params,
_execute_virtual_key_regeneration,
_get_and_validate_existing_key,
_list_key_helper,
_persist_deleted_verification_tokens,
@ -12011,6 +12014,153 @@ async def test_execute_virtual_key_regeneration_allows_within_limit_duration(mon
assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1
@pytest.mark.asyncio
async def test_execute_virtual_key_regeneration_rejects_when_custom_key_update_hook_denies():
existing_key = _make_regenerate_existing_key()
data = RegenerateKeyRequest(duration="3000d")
mock_prisma_client = _make_regenerate_mock_prisma()
received_data: list[UpdateKeyRequest] = []
async def hook(data: UpdateKeyRequest) -> dict[str, object]:
received_data.append(data)
if data.duration and duration_in_seconds(data.duration) > duration_in_seconds("7d"):
return {"decision": False, "message": "duration must be <= 7d"}
return {"decision": True}
with (
patch( # test-quality-ok: deterministic token setup for policy rejection
"litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
new_callable=AsyncMock,
return_value="sk-newtoken1234ab12",
),
patch( # test-quality-ok: grace-period path is outside policy rejection
"litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
new_callable=AsyncMock,
) as insert_deprecated_key,
patch( # test-quality-ok: cache eviction is outside policy rejection
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
new_callable=AsyncMock,
),
patch( # test-quality-ok: rotation callback is outside policy rejection
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
new_callable=AsyncMock,
),
patch("litellm.proxy.proxy_server.user_custom_key_update", hook), # test-quality-ok: inject policy hook
):
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=_make_regenerate_user_api_key_dict(),
litellm_changed_by=None,
user_api_key_cache=MagicMock(),
proxy_logging_obj=MagicMock(),
)
assert exc_info.value.status_code == 403
assert exc_info.value.detail == "duration must be <= 7d"
insert_deprecated_key.assert_not_awaited()
assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0
assert len(received_data) == 1
assert received_data[0].key == "abc123"
assert received_data[0].duration == "3000d"
@pytest.mark.asyncio
async def test_execute_virtual_key_regeneration_allows_when_custom_key_update_hook_approves():
existing_key = _make_regenerate_existing_key()
data = RegenerateKeyRequest(duration="5d")
mock_prisma_client = _make_regenerate_mock_prisma()
received_data: list[UpdateKeyRequest] = []
async def hook(data: UpdateKeyRequest) -> dict[str, object]:
received_data.append(data)
if data.duration and duration_in_seconds(data.duration) > duration_in_seconds("7d"):
return {"decision": False, "message": "duration must be <= 7d"}
return {"decision": True}
with (
patch( # test-quality-ok: deterministic token setup for policy approval
"litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
new_callable=AsyncMock,
return_value="sk-newtoken1234ab12",
),
patch( # test-quality-ok: grace-period path is outside policy approval
"litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
new_callable=AsyncMock,
),
patch( # test-quality-ok: cache eviction is outside policy approval
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
new_callable=AsyncMock,
),
patch( # test-quality-ok: rotation callback is outside policy approval
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
new_callable=AsyncMock,
),
patch("litellm.proxy.proxy_server.user_custom_key_update", hook), # test-quality-ok: inject policy hook
):
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=_make_regenerate_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
assert len(received_data) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("data", [None, RegenerateKeyRequest()])
async def test_execute_virtual_key_regeneration_skips_custom_key_update_hook_without_changes(data):
mock_prisma_client = _make_regenerate_mock_prisma()
async def hook(data: UpdateKeyRequest) -> dict[str, object]:
raise AssertionError(f"custom key update hook called with {data}")
with (
patch( # test-quality-ok: deterministic token setup for unchanged request
"litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
new_callable=AsyncMock,
return_value="sk-newtoken1234ab12",
),
patch( # test-quality-ok: grace-period path is outside unchanged request
"litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
new_callable=AsyncMock,
),
patch( # test-quality-ok: cache eviction is outside unchanged request
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
new_callable=AsyncMock,
),
patch( # test-quality-ok: rotation callback is outside unchanged request
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
new_callable=AsyncMock,
),
patch("litellm.proxy.proxy_server.user_custom_key_update", hook), # test-quality-ok: inject policy hook
):
await _execute_virtual_key_regeneration(
prisma_client=mock_prisma_client,
key_in_db=_make_regenerate_existing_key(),
hashed_api_key="abc123",
key="abc123",
data=data,
user_api_key_dict=_make_regenerate_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
@pytest.mark.asyncio
async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new_token():
"""