scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 84

fix: ensure auto-rotation updates existing AWS secret instead of creating new one (#19455)

Browse source
This commit is contained in:
Harshit Jain 2026-01-21 08:00:36 +05:30 • committed by GitHub
parent 20323feecc
commit b36e704e06
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 207 additions and 41 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

101
litellm/proxy/common_utils/key_rotation_manager.py
View file

@ -26,97 +26,119 @@ class KeyRotationManager:
"""
Manages automated key rotation based on individual key rotation schedules.
"""
def __init__(self, prisma_client: PrismaClient):
self.prisma_client = prisma_client
async def process_rotations(self):
"""
Main entry point - find and rotate keys that are due for rotation
"""
try:
verbose_proxy_logger.info("Starting scheduled key rotation check...")
# Find keys that are due for rotation
keys_to_rotate = await self._find_keys_needing_rotation()
if not keys_to_rotate:
verbose_proxy_logger.debug("No keys are due for rotation at this time")
return
verbose_proxy_logger.info(f"Found {len(keys_to_rotate)} keys due for rotation")
verbose_proxy_logger.info(
f"Found {len(keys_to_rotate)} keys due for rotation"
)
# Rotate each key
for key in keys_to_rotate:
try:
await self._rotate_key(key)
key_identifier = key.key_name or (key.token[:8] + "..." if key.token else "unknown")
verbose_proxy_logger.info(f"Successfully rotated key: {key_identifier}")
key_identifier = key.key_name or (
key.token[:8] + "..." if key.token else "unknown"
)
verbose_proxy_logger.info(
f"Successfully rotated key: {key_identifier}"
)
except Exception as e:
key_identifier = key.key_name or (key.token[:8] + "..." if key.token else "unknown")
verbose_proxy_logger.error(f"Failed to rotate key {key_identifier}: {e}")
key_identifier = key.key_name or (
key.token[:8] + "..." if key.token else "unknown"
)
verbose_proxy_logger.error(
f"Failed to rotate key {key_identifier}: {e}"
)
except Exception as e:
verbose_proxy_logger.error(f"Key rotation process failed: {e}")
async def _find_keys_needing_rotation(self) -> List[LiteLLM_VerificationToken]:
"""
Find keys that are due for rotation based on their key_rotation_at timestamp.
Logic:
- Key has auto_rotate = true
- key_rotation_at is null (needs initial setup) OR key_rotation_at <= now
"""
now = datetime.now(timezone.utc)
keys_with_rotation = await self.prisma_client.db.litellm_verificationtoken.find_many(
where={
"auto_rotate": True, # Only keys marked for auto rotation
"OR": [
{"key_rotation_at": None}, # Keys that need initial rotation time setup
{"key_rotation_at": {"lte": now}} # Keys where rotation time has passed
]
}
keys_with_rotation = (
await self.prisma_client.db.litellm_verificationtoken.find_many(
where={
"auto_rotate": True, # Only keys marked for auto rotation
"OR": [
{
"key_rotation_at": None
}, # Keys that need initial rotation time setup
{
"key_rotation_at": {"lte": now}
}, # Keys where rotation time has passed
],
}
)
)
return keys_with_rotation
def _should_rotate_key(self, key: LiteLLM_VerificationToken, now: datetime) -> bool:
"""
Determine if a key should be rotated based on key_rotation_at timestamp.
"""
if not key.rotation_interval:
return False
# If key_rotation_at is not set, rotate immediately (and set it)
if key.key_rotation_at is None:
return True
# Check if the rotation time has passed
return now >= key.key_rotation_at
async def _rotate_key(self, key: LiteLLM_VerificationToken):
"""
Rotate a single key using existing regenerate_key_fn and call the rotation hook
"""
# Create regenerate request
# Create regenerate request
regenerate_request = RegenerateKeyRequest(
key=key.token or ""
key=key.token or "",
key_alias=key.key_alias, # Pass key alias to ensure correct secret is updated in AWS Secrets Manager
)
# Create a system user for key rotation
from litellm.proxy._types import UserAPIKeyAuth
system_user = UserAPIKeyAuth.get_litellm_internal_jobs_user_api_key_auth()
# Use existing regenerate key function
response = await regenerate_key_fn(
data=regenerate_request,
user_api_key_dict=system_user,
litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
)
# Update the NEW key with rotation info (regenerate_key_fn creates a new token)
if isinstance(response, GenerateKeyResponse) and response.token_id and key.rotation_interval:
if (
isinstance(response, GenerateKeyResponse)
and response.token_id
and key.rotation_interval
):
# Calculate next rotation time using helper function
now = datetime.now(timezone.utc)
next_rotation_time = _calculate_key_rotation_time(key.rotation_interval)
@ -125,10 +147,10 @@ class KeyRotationManager:
data={
"rotation_count": (key.rotation_count or 0) + 1,
"last_rotation_at": now,
"key_rotation_at": next_rotation_time
}
"key_rotation_at": next_rotation_time,
},
)
# Call the existing rotation hook for notifications, audit logs, etc.
if isinstance(response, GenerateKeyResponse):
await KeyManagementEventHooks.async_key_rotated_hook(
@ -136,6 +158,5 @@ class KeyRotationManager:
existing_key_row=key,
response=response,
user_api_key_dict=system_user,
litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
)

3
litellm/proxy/hooks/key_management_event_hooks.py
View file

@ -152,7 +152,8 @@ class KeyManagementEventHooks:
)
await KeyManagementEventHooks._rotate_virtual_key_in_secret_manager(
current_secret_name=initial_secret_name,
new_secret_name=data.key_alias
new_secret_name=response.key_alias
or data.key_alias
or f"virtual-key-{response.token_id}",
new_secret_value=response.key,
)

144
tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py Normal file
View file

@ -0,0 +1,144 @@
"""
Regression test for AWS Secrets Manager Auto-Rotation Bug Fix
This test verifies that KeyRotationManager correctly passes key_alias
when calling regenerate_key_fn, ensuring the secret is rotated at the
correct location in AWS Secrets Manager.
Bug Fixed: Key alias was not passed during auto-rotation, causing
secrets to be created at a new location instead of updating in-place.
"""
import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
sys.path.insert(0, os.path.abspath("../../../.."))
from litellm.proxy._types import (
GenerateKeyResponse,
LiteLLM_VerificationToken,
RegenerateKeyRequest,
)
from litellm.proxy.common_utils.key_rotation_manager import KeyRotationManager
class TestKeyRotationManagerPassesKeyAlias:
"""
Regression tests to ensure KeyRotationManager passes key_alias
to regenerate_key_fn during auto-rotation.
"""
@pytest.mark.asyncio
async def test_rotate_key_passes_key_alias_to_regenerate_request(self):
"""
Verify that _rotate_key includes key_alias in the RegenerateKeyRequest.
This is the core fix: previously, key_alias was NOT passed, causing
the secret manager hook to use a generated name instead of the alias.
"""
# Create a mock key with an alias
test_alias = "tenant1/my-important-key"
test_token = "sk-test-token-hash-12345"
mock_key = MagicMock(spec=LiteLLM_VerificationToken)
mock_key.token = test_token
mock_key.key_alias = test_alias
mock_key.key_name = "sk-...1234"
mock_key.rotation_interval = "30d"
mock_key.rotation_count = 0
# Create mock prisma client
mock_prisma = MagicMock()
mock_prisma.db.litellm_verificationtoken.update = AsyncMock(
return_value=mock_key
)
# Create mock response
mock_response = GenerateKeyResponse(
key="sk-new-key-value",
token_id="new-token-hash",
key_alias=test_alias,
)
# Capture the RegenerateKeyRequest passed to regenerate_key_fn
captured_request = None
async def capture_regenerate_key_fn(
data, user_api_key_dict, litellm_changed_by
):
nonlocal captured_request
captured_request = data
return mock_response
# Patch regenerate_key_fn to capture the request
with patch(
"litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn",
side_effect=capture_regenerate_key_fn,
):
with patch(
"litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook",
new_callable=AsyncMock,
):
rotation_manager = KeyRotationManager(mock_prisma)
await rotation_manager._rotate_key(mock_key)
# CRITICAL ASSERTION: key_alias must be passed
assert captured_request is not None, "regenerate_key_fn should have been called"
assert isinstance(captured_request, RegenerateKeyRequest)
assert captured_request.key == test_token, "Token should be passed correctly"
assert captured_request.key_alias == test_alias, (
f"key_alias should be '{test_alias}' but was '{captured_request.key_alias}'. "
"This is the bug we fixed - key_alias was not being passed!"
)
@pytest.mark.asyncio
async def test_rotate_key_passes_none_alias_when_key_has_no_alias(self):
"""
Verify that _rotate_key handles keys without an alias gracefully.
"""
test_token = "sk-test-token-hash-67890"
mock_key = MagicMock(spec=LiteLLM_VerificationToken)
mock_key.token = test_token
mock_key.key_alias = None # No alias set
mock_key.key_name = "sk-...5678"
mock_key.rotation_interval = "30d"
mock_key.rotation_count = 0
mock_prisma = MagicMock()
mock_prisma.db.litellm_verificationtoken.update = AsyncMock(
return_value=mock_key
)
mock_response = GenerateKeyResponse(
key="sk-new-key-value",
token_id="new-token-hash",
)
captured_request = None
async def capture_regenerate_key_fn(
data, user_api_key_dict, litellm_changed_by
):
nonlocal captured_request
captured_request = data
return mock_response
with patch(
"litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn",
side_effect=capture_regenerate_key_fn,
):
with patch(
"litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook",
new_callable=AsyncMock,
):
rotation_manager = KeyRotationManager(mock_prisma)
await rotation_manager._rotate_key(mock_key)
assert captured_request is not None
assert captured_request.key == test_token
assert (
captured_request.key_alias is None
), "key_alias should be None for keys without alias"