mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
Merge 97a26159af into e16d79d635
This commit is contained in:
Jack Danger
GitHub
commit
8f6bbb1513
4 changed files with 1099 additions and 331 deletions
|
|
@ -2338,9 +2338,31 @@ class ModelManagementAuthChecks:
|
|||
) -> None:
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
|
||||
return
|
||||
stored: Final = _decrypted_litellm_params(model_params.litellm_params)
|
||||
|
||||
# Static-key endpoint harvest (LD fork R2-9): a non-admin write that
|
||||
# moves api_base/base_url while the stored server-owned api_key stays
|
||||
# in place sends that key to the endpoint the caller names. The team
|
||||
# must bring its own key (set api_key in the same write) or route
|
||||
# through an admin-attached credential.
|
||||
repoints_endpoint: Final = incoming_params is not None and not incoming_params.keys().isdisjoint(
|
||||
("api_base", "base_url")
|
||||
)
|
||||
brings_own_key: Final = incoming_params is not None and incoming_params.get("api_key") is not None
|
||||
server_key_stays: Final = stored.get("api_key") not in (None, "") and not brings_own_key
|
||||
if repoints_endpoint and server_key_stays:
|
||||
raise ProxyException(
|
||||
message=(
|
||||
"Only proxy admins can repoint a deployment's endpoint (api_base/base_url) "
|
||||
"while its api_key stays server-owned; supply your own api_key with the "
|
||||
"new endpoint."
|
||||
),
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
code=status.HTTP_403_FORBIDDEN,
|
||||
param="api_base",
|
||||
)
|
||||
if not write_touches_federation_surface(incoming_params):
|
||||
return
|
||||
stored: Final = _decrypted_litellm_params(model_params.litellm_params)
|
||||
wif_fields: Final = await effective_server_owned_wif_fields(stored, incoming_params, prisma_client)
|
||||
if wif_fields:
|
||||
# ProxyException rather than HTTPException so the offending field stays a structured
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -0,0 +1,102 @@
|
|||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_admin_cannot_move_api_base_keeping_server_key():
|
||||
from fastapi import HTTPException
|
||||
from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
ModelManagementAuthChecks,
|
||||
)
|
||||
|
||||
stored = GenericLiteLLMParams(api_key="server-owned-key")
|
||||
incoming = GenericLiteLLMParams(api_base="https://attacker.example")
|
||||
model_params = MagicMock()
|
||||
model_params.litellm_params = stored
|
||||
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await ModelManagementAuthChecks._reject_non_admin_wif_write(
|
||||
model_params=model_params,
|
||||
incoming_params=incoming.model_dump(exclude_none=True),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="sk-team-admin",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
user_id="team-admin-1",
|
||||
),
|
||||
prisma_client=AsyncMock(),
|
||||
)
|
||||
assert int(getattr(exc.value, "code", 0)) == 403
|
||||
assert "api_key stays server-owned" in str(getattr(exc.value, "message", getattr(exc.value, "detail", "")))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_admin_bring_your_own_key_move_allowed():
|
||||
import litellm.proxy.management_endpoints.model_management_endpoints as mme
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
ModelManagementAuthChecks,
|
||||
)
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
|
||||
stored = GenericLiteLLMParams(api_key="server-owned-key")
|
||||
incoming = GenericLiteLLMParams(api_base="https://my-team-proxy.example", api_key="my-team-key")
|
||||
model_params = MagicMock()
|
||||
model_params.litellm_params = stored
|
||||
|
||||
with (
|
||||
patch.object(mme, "effective_server_owned_wif_fields", AsyncMock(return_value=())),
|
||||
):
|
||||
# The observable on a BYOK move is that the server-owned-credential
|
||||
# guard raised nothing for this caller.
|
||||
no_raise = await ModelManagementAuthChecks._reject_non_admin_wif_write(
|
||||
model_params=model_params,
|
||||
incoming_params=incoming.model_dump(exclude_none=True),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="sk-team",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
user_id="team-admin-1",
|
||||
),
|
||||
prisma_client=AsyncMock(),
|
||||
)
|
||||
assert no_raise is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_admin_wif_association_still_guarded():
|
||||
"""Under WIF the original server-owned-token rule keeps firing."""
|
||||
import litellm.proxy.management_endpoints.model_management_endpoints as mme
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
ModelManagementAuthChecks,
|
||||
)
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
|
||||
stored = GenericLiteLLMParams(api_key="server-key")
|
||||
incoming = GenericLiteLLMParams(api_base="https://moved.example")
|
||||
model_params = MagicMock()
|
||||
model_params.litellm_params = stored
|
||||
|
||||
with (
|
||||
patch.object(mme, "effective_server_owned_wif_fields", AsyncMock(return_value=("oidc_token_url",))),
|
||||
):
|
||||
# Capture outside the except block: on this stack both the static-key
|
||||
# harvest gate and the WIF rule can fire for the same shape; accept
|
||||
# either interdiction.
|
||||
guard_raised = None
|
||||
try:
|
||||
await ModelManagementAuthChecks._reject_non_admin_wif_write(
|
||||
model_params=model_params,
|
||||
incoming_params=incoming.model_dump(exclude_none=True),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="sk-team", user_role=LitellmUserRoles.INTERNAL_USER, user_id="t"
|
||||
),
|
||||
prisma_client=AsyncMock(),
|
||||
)
|
||||
except Exception as guard_error:
|
||||
guard_raised = guard_error
|
||||
assert guard_raised is not None, "expected a server-owned-key interdiction to fire"
|
||||
assert ("api_key stays server-owned" in str(guard_raised)) or ("workload identity" in str(guard_raised))
|
||||
102
tests/unit/proxy/management_endpoints/test_r2_9_byok_flows.py
Normal file
102
tests/unit/proxy/management_endpoints/test_r2_9_byok_flows.py
Normal file
|
|
@ -0,0 +1,102 @@
|
|||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_admin_cannot_move_api_base_keeping_server_key():
|
||||
from fastapi import HTTPException
|
||||
from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth
|
||||
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
ModelManagementAuthChecks,
|
||||
)
|
||||
|
||||
stored = GenericLiteLLMParams(api_key="server-owned-key")
|
||||
incoming = GenericLiteLLMParams(api_base="https://attacker.example")
|
||||
model_params = MagicMock()
|
||||
model_params.litellm_params = stored
|
||||
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await ModelManagementAuthChecks._reject_non_admin_wif_write(
|
||||
model_params=model_params,
|
||||
incoming_params=incoming.model_dump(exclude_none=True),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="sk-team-admin",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
user_id="team-admin-1",
|
||||
),
|
||||
prisma_client=AsyncMock(),
|
||||
)
|
||||
assert int(getattr(exc.value, "code", 0)) == 403
|
||||
assert "api_key stays server-owned" in str(getattr(exc.value, "message", getattr(exc.value, "detail", "")))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_admin_bring_your_own_key_move_allowed():
|
||||
import litellm.proxy.management_endpoints.model_management_endpoints as mme
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
ModelManagementAuthChecks,
|
||||
)
|
||||
|
||||
stored = GenericLiteLLMParams(api_key="server-owned-key")
|
||||
incoming = GenericLiteLLMParams(api_base="https://my-team-proxy.example", api_key="my-team-key")
|
||||
model_params = MagicMock()
|
||||
model_params.litellm_params = stored
|
||||
|
||||
with (
|
||||
patch.object(mme, "effective_server_owned_wif_fields", AsyncMock(return_value=())),
|
||||
):
|
||||
# The observable on a BYOK move is that the server-owned-credential
|
||||
# guard raised nothing for this caller.
|
||||
no_raise = await ModelManagementAuthChecks._reject_non_admin_wif_write(
|
||||
model_params=model_params,
|
||||
incoming_params=incoming.model_dump(exclude_none=True),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="sk-team",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
user_id="team-admin-1",
|
||||
),
|
||||
prisma_client=AsyncMock(),
|
||||
)
|
||||
assert no_raise is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_admin_wif_association_still_guarded():
|
||||
"""Under WIF the original server-owned-token rule keeps firing."""
|
||||
import litellm.proxy.management_endpoints.model_management_endpoints as mme
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
ModelManagementAuthChecks,
|
||||
)
|
||||
|
||||
stored = GenericLiteLLMParams(api_key="server-key")
|
||||
incoming = GenericLiteLLMParams(api_base="https://moved.example")
|
||||
model_params = MagicMock()
|
||||
model_params.litellm_params = stored
|
||||
|
||||
with (
|
||||
patch.object(mme, "effective_server_owned_wif_fields", AsyncMock(return_value=("oidc_token_url",))),
|
||||
):
|
||||
# Capture outside the except block: on this stack both the static-key
|
||||
# harvest gate and the WIF rule can fire for the same shape; accept
|
||||
# either interdiction.
|
||||
guard_raised = None
|
||||
try:
|
||||
await ModelManagementAuthChecks._reject_non_admin_wif_write(
|
||||
model_params=model_params,
|
||||
incoming_params=incoming.model_dump(exclude_none=True),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="sk-team", user_role=LitellmUserRoles.INTERNAL_USER, user_id="t"
|
||||
),
|
||||
prisma_client=AsyncMock(),
|
||||
)
|
||||
except Exception as guard_error:
|
||||
guard_raised = guard_error
|
||||
assert guard_raised is not None, "expected a server-owned-key interdiction to fire"
|
||||
assert ("api_key stays server-owned" in str(guard_raised)) or ("workload identity" in str(guard_raised))
|
||||
Loading…
Add table
Reference in a new issue