mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge 47ea8f7cb4 into e9b1c11098
This commit is contained in:
commit
f282597bf6
5 changed files with 84 additions and 56 deletions
|
|
@ -6493,8 +6493,7 @@ async def regenerate_key_fn(
|
|||
- key: str (path parameter) - The key to regenerate
|
||||
- data: Optional[RegenerateKeyRequest] - Request body containing optional parameters to update
|
||||
- key: Optional[str] - The key to regenerate.
|
||||
- new_master_key: Optional[str] - The new master key to use, if key is the master key.
|
||||
- new_key: Optional[str] - The new key to use, if key is not the master key. Must start with 'sk-' and be at least 16 characters long. If both set, new_master_key will be used.
|
||||
- new_key: Optional[str] - The new key to use. Must start with 'sk-' and be at least 16 characters long.
|
||||
- key_alias: Optional[str] - User-friendly key alias
|
||||
- user_id: Optional[str] - User ID associated with key
|
||||
- team_id: Optional[str] - Team ID associated with key
|
||||
|
|
@ -6574,21 +6573,21 @@ async def regenerate_key_fn(
|
|||
allow_safe_presets=True,
|
||||
)
|
||||
|
||||
# Premium-gate bypass for master-key rotation must verify the
|
||||
# caller actually holds the master key, not just that the request
|
||||
# body has a ``new_master_key`` field. A presence-only check let
|
||||
# any non-premium caller skip the enterprise gate by sending any
|
||||
# value in that field.
|
||||
regenerate_target_key: Final = data.key if data and data.key else key
|
||||
is_master_key_regeneration: Final = (
|
||||
data is not None
|
||||
and data.new_master_key is not None
|
||||
and is_master_key(api_key=regenerate_target_key, _master_key=master_key)
|
||||
)
|
||||
if master_key is not None and is_master_key(api_key=regenerate_target_key, _master_key=master_key):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail={
|
||||
"error": (
|
||||
"Rotating the master key through /key/regenerate is no longer supported. "
|
||||
"Set the new value in LITELLM_MASTER_KEY or general_settings.master_key and restart every "
|
||||
"proxy instance. If LITELLM_SALT_KEY is not set, first set it to your current master key. "
|
||||
"See https://docs.litellm.ai/docs/proxy/master_key_rotations"
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
if (
|
||||
premium_user is not True and not is_master_key_regeneration
|
||||
): # allow master key regeneration for non-premium users
|
||||
if premium_user is not True:
|
||||
raise ValueError(
|
||||
f"Regenerating Virtual Keys is an Enterprise feature, {CommonProxyErrors.not_premium_user.value}"
|
||||
)
|
||||
|
|
@ -6611,27 +6610,6 @@ async def regenerate_key_fn(
|
|||
detail={"error": "DB not connected. prisma_client is None"},
|
||||
)
|
||||
|
||||
_is_master_key_valid: Final = is_master_key(api_key=key, _master_key=master_key)
|
||||
|
||||
if master_key is not None and data and _is_master_key_valid:
|
||||
if data.new_master_key is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail={"error": "New master key is required."},
|
||||
)
|
||||
await rotate_master_key(
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
current_master_key=master_key,
|
||||
new_master_key=data.new_master_key,
|
||||
)
|
||||
return GenerateKeyResponse(
|
||||
key=data.new_master_key,
|
||||
token=data.new_master_key,
|
||||
key_name=data.new_master_key,
|
||||
expires=None,
|
||||
)
|
||||
|
||||
if "sk" not in key:
|
||||
hashed_api_key = key
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -156,5 +156,4 @@ class RegenerateKeyRequest(GenerateKeyRequest):
|
|||
duration: str | None = None
|
||||
spend: float | None = None
|
||||
metadata: dict[Any, Any] | None = None
|
||||
new_master_key: str | None = None
|
||||
grace_period: str | None = None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke
|
||||
|
|
|
|||
|
|
@ -1,7 +1,12 @@
|
|||
import json
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, Actor
|
||||
from tests.integration._support.actors import create_scratch_key, owned_settings_client
|
||||
from tests.integration._support.client import Gateway, object_value
|
||||
from tests.integration._support.wire import Reply, Request, wire_server
|
||||
|
||||
pytestmark = pytest.mark.asyncio(loop_scope="session")
|
||||
|
||||
|
|
@ -40,6 +45,56 @@ async def _info(gateway_client, cleartext: str):
|
|||
)
|
||||
|
||||
|
||||
async def test_master_key_regeneration_is_rejected_without_rotating_model(gateway: Gateway):
|
||||
master_key_regeneration_error: Final = (
|
||||
"Rotating the master key through /key/regenerate is no longer supported. "
|
||||
"Set the new value in LITELLM_MASTER_KEY or general_settings.master_key and restart every proxy instance. "
|
||||
"If LITELLM_SALT_KEY is not set, first set it to your current master key. "
|
||||
"See https://docs.litellm.ai/docs/proxy/master_key_rotations"
|
||||
)
|
||||
|
||||
def upstream(_: Request) -> Reply:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": "master-key-regeneration-response",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "available"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
with gateway.scenario() as scenario, wire_server(upstream) as provider:
|
||||
model: Final = scenario.model(api_base=f"{provider.url}/v1")
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/key/regenerate",
|
||||
{"key": gateway.key, "new_master_key": "sk-new-master-key"},
|
||||
)
|
||||
assert response.status_code == 400, response.text
|
||||
assert response.json()["error"]["message"] == str({"error": master_key_regeneration_error})
|
||||
|
||||
model_info: Final = gateway.request("GET", "/model/info")
|
||||
assert model_info.status_code == 200, model_info.text
|
||||
listed: Final = model_info.json()["data"]
|
||||
assert isinstance(listed, list)
|
||||
assert any(object_value(entry)["model_name"] == model for entry in listed)
|
||||
served: Final = gateway.chat(model, text="master key regeneration")
|
||||
assert object_value(served["choices"][0])["message"] == {
|
||||
"role": "assistant",
|
||||
"content": "available",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"actor,target_shape,expected_status",
|
||||
[(a, t, s) for (_id, a, t, s) in _SCENARIOS],
|
||||
|
|
|
|||
|
|
@ -16184,16 +16184,13 @@ async def test_regenerate_user_id_rebind_guard(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_regenerate_premium_gate_requires_actual_master_key():
|
||||
# ``regenerate_key_fn``'s decorator wraps the underlying ValueError
|
||||
# into a ProxyException with empty ``message``. The exception type
|
||||
# alone confirms the premium gate fired.
|
||||
async def test_regenerate_premium_gate_rejects_non_master_key():
|
||||
from litellm.proxy._types import RegenerateKeyRequest
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
regenerate_key_fn,
|
||||
)
|
||||
|
||||
data = RegenerateKeyRequest(key="sk-not-master", new_master_key="anything")
|
||||
data = RegenerateKeyRequest(key="sk-not-master")
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.premium_user", False),
|
||||
|
|
@ -16209,25 +16206,22 @@ async def test_regenerate_premium_gate_requires_actual_master_key():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_regenerate_premium_gate_allows_actual_master_key_holder():
|
||||
async def test_regenerate_rejects_master_key():
|
||||
from litellm.proxy._types import RegenerateKeyRequest
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
regenerate_key_fn,
|
||||
)
|
||||
|
||||
master = "sk-the-real-master-key"
|
||||
data = RegenerateKeyRequest(key=master, new_master_key="sk-new-master")
|
||||
data = RegenerateKeyRequest(key=master)
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.premium_user", False),
|
||||
patch("litellm.proxy.proxy_server.master_key", master),
|
||||
patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.rotate_master_key",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
pytest.raises(ProxyException) as exc_info,
|
||||
):
|
||||
result = await regenerate_key_fn(
|
||||
await regenerate_key_fn(
|
||||
key=master,
|
||||
data=data,
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
|
|
@ -16237,7 +16231,13 @@ async def test_regenerate_premium_gate_allows_actual_master_key_holder():
|
|||
),
|
||||
)
|
||||
|
||||
assert result.token == "sk-new-master"
|
||||
assert exc_info.value.code == "400"
|
||||
assert exc_info.value.message == (
|
||||
"{'error': 'Rotating the master key through /key/regenerate is no longer supported. "
|
||||
"Set the new value in LITELLM_MASTER_KEY or general_settings.master_key and restart every proxy instance. "
|
||||
"If LITELLM_SALT_KEY is not set, first set it to your current master key. "
|
||||
"See https://docs.litellm.ai/docs/proxy/master_key_rotations'}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
8
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
8
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -8892,8 +8892,7 @@ export interface paths {
|
|||
* - key: str (path parameter) - The key to regenerate
|
||||
* - data: Optional[RegenerateKeyRequest] - Request body containing optional parameters to update
|
||||
* - key: Optional[str] - The key to regenerate.
|
||||
* - new_master_key: Optional[str] - The new master key to use, if key is the master key.
|
||||
* - new_key: Optional[str] - The new key to use, if key is not the master key. Must start with 'sk-' and be at least 16 characters long. If both set, new_master_key will be used.
|
||||
* - new_key: Optional[str] - The new key to use. Must start with 'sk-' and be at least 16 characters long.
|
||||
* - key_alias: Optional[str] - User-friendly key alias
|
||||
* - user_id: Optional[str] - User ID associated with key
|
||||
* - team_id: Optional[str] - Team ID associated with key
|
||||
|
|
@ -9191,8 +9190,7 @@ export interface paths {
|
|||
* - key: str (path parameter) - The key to regenerate
|
||||
* - data: Optional[RegenerateKeyRequest] - Request body containing optional parameters to update
|
||||
* - key: Optional[str] - The key to regenerate.
|
||||
* - new_master_key: Optional[str] - The new master key to use, if key is the master key.
|
||||
* - new_key: Optional[str] - The new key to use, if key is not the master key. Must start with 'sk-' and be at least 16 characters long. If both set, new_master_key will be used.
|
||||
* - new_key: Optional[str] - The new key to use. Must start with 'sk-' and be at least 16 characters long.
|
||||
* - key_alias: Optional[str] - User-friendly key alias
|
||||
* - user_id: Optional[str] - User ID associated with key
|
||||
* - team_id: Optional[str] - Team ID associated with key
|
||||
|
|
@ -43568,8 +43566,6 @@ export interface components {
|
|||
models: unknown[] | null;
|
||||
/** New Key */
|
||||
new_key?: string | null;
|
||||
/** New Master Key */
|
||||
new_master_key?: string | null;
|
||||
object_permission?: components["schemas"]["LiteLLM_ObjectPermissionBase"] | null;
|
||||
/** Organization Id */
|
||||
organization_id?: string | null;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue