mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
[Fix] Key Expiry Validation Against Team Max Duration
Virtual keys created or regenerated with a duration exceeding the team's team_member_key_duration limit were silently accepted. On creation, the backend silently capped the value giving the user a misleading success response. On regeneration, the user value was stored as-is, bypassing the team limit entirely. Add a duration check in _common_key_generation_helper (key creation) and a _validate_regenerate_key_duration_against_team helper called from _execute_virtual_key_regeneration (key regeneration). Both raise HTTP 400 when the requested duration exceeds the team maximum. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
yuneng-jiang
parent
4c1b15d685
commit
3ed3255c47
2 changed files with 343 additions and 2 deletions
|
|
@ -547,6 +547,27 @@ async def _common_key_generation_helper( # noqa: PLR0915
|
|||
},
|
||||
)
|
||||
|
||||
# Validate key duration against the team's max key duration
|
||||
if (
|
||||
team_table is not None
|
||||
and team_table.metadata is not None
|
||||
and team_table.metadata.get("team_member_key_duration")
|
||||
and data.duration is not None
|
||||
):
|
||||
team_max_duration = team_table.metadata["team_member_key_duration"]
|
||||
team_max_seconds = duration_in_seconds(duration=team_max_duration)
|
||||
if data.duration == "-1":
|
||||
user_key_duration: float = float("inf")
|
||||
else:
|
||||
user_key_duration = duration_in_seconds(duration=data.duration)
|
||||
if user_key_duration > team_max_seconds:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": f"Key duration exceeds team maximum. Requested: {data.duration}; Team maximum: {team_max_duration}"
|
||||
},
|
||||
)
|
||||
|
||||
# APPLY ENTERPRISE KEY MANAGEMENT PARAMS
|
||||
try:
|
||||
from litellm_enterprise.proxy.management_endpoints.key_management_endpoints import (
|
||||
|
|
@ -3372,6 +3393,55 @@ async def _insert_deprecated_key(
|
|||
"Failed to insert deprecated key for grace period: %s",
|
||||
deprecated_err,
|
||||
)
|
||||
async def _validate_regenerate_key_duration_against_team(
|
||||
data: Optional[RegenerateKeyRequest],
|
||||
key_in_db: LiteLLM_VerificationToken,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: DualCache,
|
||||
) -> None:
|
||||
"""Raise HTTP 400 if the requested duration exceeds the team's max key duration."""
|
||||
if data is None or not data.duration:
|
||||
return
|
||||
|
||||
team_id = getattr(key_in_db, "team_id", None)
|
||||
if not team_id:
|
||||
return
|
||||
|
||||
try:
|
||||
team_table = await get_team_object(
|
||||
team_id=team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
check_db_only=True,
|
||||
)
|
||||
except Exception:
|
||||
return
|
||||
|
||||
if (
|
||||
team_table is None
|
||||
or team_table.metadata is None
|
||||
or not team_table.metadata.get("team_member_key_duration")
|
||||
):
|
||||
return
|
||||
|
||||
team_max_duration = team_table.metadata["team_member_key_duration"]
|
||||
team_max_seconds = duration_in_seconds(duration=team_max_duration)
|
||||
|
||||
if data.duration == "-1":
|
||||
user_seconds: float = float("inf")
|
||||
else:
|
||||
user_seconds = duration_in_seconds(duration=data.duration)
|
||||
|
||||
if user_seconds > team_max_seconds:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": f"Key duration exceeds team maximum. Requested: {data.duration}; Team maximum: {team_max_duration}"
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def _execute_virtual_key_regeneration(
|
||||
*,
|
||||
prisma_client: PrismaClient,
|
||||
|
|
@ -3387,6 +3457,13 @@ async def _execute_virtual_key_regeneration(
|
|||
"""Generate new token, update DB, invalidate cache, and return response."""
|
||||
from litellm.proxy.proxy_server import hash_token
|
||||
|
||||
await _validate_regenerate_key_duration_against_team(
|
||||
data=data,
|
||||
key_in_db=key_in_db,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
new_token = get_new_token(data=data)
|
||||
new_token_hash = hash_token(new_token)
|
||||
new_token_key_name = f"sk-...{new_token[-4:]}"
|
||||
|
|
|
|||
|
|
@ -6436,7 +6436,7 @@ class TestValidateKeyAliasFormat:
|
|||
def test_validate_key_alias_format_invalid(self):
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import _validate_key_alias_format
|
||||
from litellm.proxy._types import ProxyException
|
||||
|
||||
|
||||
invalid_aliases = [
|
||||
"", # empty
|
||||
" ", # whitespace
|
||||
|
|
@ -6449,9 +6449,273 @@ class TestValidateKeyAliasFormat:
|
|||
" leading",
|
||||
"trailing ",
|
||||
]
|
||||
|
||||
|
||||
for alias in invalid_aliases:
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
_validate_key_alias_format(alias)
|
||||
assert str(exc.value.code) == "400"
|
||||
assert "Invalid key_alias format" in str(exc.value.message)
|
||||
|
||||
|
||||
class TestCommonKeyGenerationHelperTeamDurationValidation:
|
||||
"""Tests for team duration validation inside _common_key_generation_helper."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duration_exceeds_team_max_raises(self):
|
||||
"""Raises HTTP 400 when the requested duration exceeds the team's max."""
|
||||
team = MagicMock(spec=LiteLLM_TeamTableCachedObj)
|
||||
team.metadata = {"team_member_key_duration": "5d"}
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _common_key_generation_helper(
|
||||
data=GenerateKeyRequest(user_id="user-1", duration="10d"),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"
|
||||
),
|
||||
litellm_changed_by=None,
|
||||
team_table=team,
|
||||
)
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "Team maximum" in str(exc_info.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_never_expires_exceeds_team_max_raises(self):
|
||||
"""Duration '-1' (never expires) exceeds any finite team max."""
|
||||
team = MagicMock(spec=LiteLLM_TeamTableCachedObj)
|
||||
team.metadata = {"team_member_key_duration": "30d"}
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _common_key_generation_helper(
|
||||
data=GenerateKeyRequest(user_id="user-1", duration="-1"),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"
|
||||
),
|
||||
litellm_changed_by=None,
|
||||
team_table=team,
|
||||
)
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "Team maximum" in str(exc_info.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duration_within_team_max_passes(self):
|
||||
"""Duration within the team max proceeds past the validation block."""
|
||||
team = MagicMock(spec=LiteLLM_TeamTableCachedObj)
|
||||
team.metadata = {"team_member_key_duration": "5d"}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"key": "sk-test", "expires": None, "user_id": "user-1"},
|
||||
), patch("litellm.proxy.proxy_server.prisma_client"), patch(
|
||||
"litellm.proxy.proxy_server.llm_router"
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server.premium_user", False
|
||||
):
|
||||
await _common_key_generation_helper(
|
||||
data=GenerateKeyRequest(user_id="user-1", duration="3d"),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"
|
||||
),
|
||||
litellm_changed_by=None,
|
||||
team_table=team,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_duration_skips_team_validation(self):
|
||||
"""No duration provided — team validation is skipped."""
|
||||
team = MagicMock(spec=LiteLLM_TeamTableCachedObj)
|
||||
team.metadata = {"team_member_key_duration": "5d"}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"key": "sk-test", "expires": None, "user_id": "user-1"},
|
||||
), patch("litellm.proxy.proxy_server.prisma_client"), patch(
|
||||
"litellm.proxy.proxy_server.llm_router"
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server.premium_user", False
|
||||
):
|
||||
await _common_key_generation_helper(
|
||||
data=GenerateKeyRequest(user_id="user-1", duration=None),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"
|
||||
),
|
||||
litellm_changed_by=None,
|
||||
team_table=team,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_team_skips_team_validation(self):
|
||||
"""No team_table — team duration validation is skipped."""
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"key": "sk-test", "expires": None, "user_id": "user-1"},
|
||||
), patch("litellm.proxy.proxy_server.prisma_client"), patch(
|
||||
"litellm.proxy.proxy_server.llm_router"
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server.premium_user", False
|
||||
):
|
||||
await _common_key_generation_helper(
|
||||
data=GenerateKeyRequest(user_id="user-1", duration="100d"),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"
|
||||
),
|
||||
litellm_changed_by=None,
|
||||
team_table=None,
|
||||
)
|
||||
|
||||
|
||||
class TestValidateRegenerateKeyDurationAgainstTeam:
|
||||
"""Tests for _validate_regenerate_key_duration_against_team."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_data_skips_validation(self):
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_validate_regenerate_key_duration_against_team,
|
||||
)
|
||||
|
||||
mock_key = MagicMock(spec=LiteLLM_VerificationToken)
|
||||
mock_key.team_id = "team-123"
|
||||
mock_prisma = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
# Should not raise
|
||||
await _validate_regenerate_key_duration_against_team(
|
||||
data=None,
|
||||
key_in_db=mock_key,
|
||||
prisma_client=mock_prisma,
|
||||
user_api_key_cache=mock_cache,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_duration_skips_validation(self):
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_validate_regenerate_key_duration_against_team,
|
||||
)
|
||||
from litellm.proxy._types import RegenerateKeyRequest
|
||||
|
||||
mock_key = MagicMock(spec=LiteLLM_VerificationToken)
|
||||
mock_key.team_id = "team-123"
|
||||
data = RegenerateKeyRequest(duration=None)
|
||||
mock_prisma = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
# Should not raise
|
||||
await _validate_regenerate_key_duration_against_team(
|
||||
data=data,
|
||||
key_in_db=mock_key,
|
||||
prisma_client=mock_prisma,
|
||||
user_api_key_cache=mock_cache,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_team_id_skips_validation(self):
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_validate_regenerate_key_duration_against_team,
|
||||
)
|
||||
from litellm.proxy._types import RegenerateKeyRequest
|
||||
|
||||
mock_key = MagicMock(spec=LiteLLM_VerificationToken)
|
||||
mock_key.team_id = None
|
||||
data = RegenerateKeyRequest(duration="10d")
|
||||
mock_prisma = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
# Should not raise
|
||||
await _validate_regenerate_key_duration_against_team(
|
||||
data=data,
|
||||
key_in_db=mock_key,
|
||||
prisma_client=mock_prisma,
|
||||
user_api_key_cache=mock_cache,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duration_within_team_max_passes(self):
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_validate_regenerate_key_duration_against_team,
|
||||
)
|
||||
from litellm.proxy._types import RegenerateKeyRequest, LiteLLM_TeamTableCachedObj
|
||||
|
||||
mock_key = MagicMock(spec=LiteLLM_VerificationToken)
|
||||
mock_key.team_id = "team-123"
|
||||
data = RegenerateKeyRequest(duration="3d")
|
||||
|
||||
mock_team = MagicMock(spec=LiteLLM_TeamTableCachedObj)
|
||||
mock_team.metadata = {"team_member_key_duration": "5d"}
|
||||
|
||||
mock_prisma = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_team,
|
||||
):
|
||||
# Should not raise
|
||||
await _validate_regenerate_key_duration_against_team(
|
||||
data=data,
|
||||
key_in_db=mock_key,
|
||||
prisma_client=mock_prisma,
|
||||
user_api_key_cache=mock_cache,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duration_exceeds_team_max_raises(self):
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_validate_regenerate_key_duration_against_team,
|
||||
)
|
||||
from litellm.proxy._types import RegenerateKeyRequest, LiteLLM_TeamTableCachedObj
|
||||
|
||||
mock_key = MagicMock(spec=LiteLLM_VerificationToken)
|
||||
mock_key.team_id = "team-123"
|
||||
data = RegenerateKeyRequest(duration="10d")
|
||||
|
||||
mock_team = MagicMock(spec=LiteLLM_TeamTableCachedObj)
|
||||
mock_team.metadata = {"team_member_key_duration": "5d"}
|
||||
|
||||
mock_prisma = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_team,
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _validate_regenerate_key_duration_against_team(
|
||||
data=data,
|
||||
key_in_db=mock_key,
|
||||
prisma_client=mock_prisma,
|
||||
user_api_key_cache=mock_cache,
|
||||
)
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "Team maximum" in str(exc_info.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_never_expires_duration_exceeds_team_max_raises(self):
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_validate_regenerate_key_duration_against_team,
|
||||
)
|
||||
from litellm.proxy._types import RegenerateKeyRequest, LiteLLM_TeamTableCachedObj
|
||||
|
||||
mock_key = MagicMock(spec=LiteLLM_VerificationToken)
|
||||
mock_key.team_id = "team-123"
|
||||
data = RegenerateKeyRequest(duration="-1")
|
||||
|
||||
mock_team = MagicMock(spec=LiteLLM_TeamTableCachedObj)
|
||||
mock_team.metadata = {"team_member_key_duration": "30d"}
|
||||
|
||||
mock_prisma = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_team,
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _validate_regenerate_key_duration_against_team(
|
||||
data=data,
|
||||
key_in_db=mock_key,
|
||||
prisma_client=mock_prisma,
|
||||
user_api_key_cache=mock_cache,
|
||||
)
|
||||
assert exc_info.value.status_code == 400
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue