mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
feat(proxy): restrict aws_session_tags on model management to proxy admins
Team admins can create and edit team models through /model/new, PUT /model/update
and PATCH /model/{id}/update. The proxy forwards aws_session_tags to STS under its
own identity, so a team admin could pick tags that unlock aws:PrincipalTag gated
resources. Only proxy admins may now set or change aws_session_tags there; an
unchanged tag set still passes so team admins can edit other fields
This commit is contained in:
ryan-crabbe-berri
parent
d909c101ab
commit
b1fb478e37
2 changed files with 223 additions and 0 deletions
|
|
@ -111,6 +111,7 @@ from litellm.router_utils.auto_router_model_naming import (
|
|||
validate_strategy_router_model_write,
|
||||
)
|
||||
from litellm.router_utils.auto_router_tuning_baseline import is_mutable_tuned_candidate, tuning_quota_violation
|
||||
from litellm.types.llms.bedrock import AwsSessionTag
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
AutoRouterClassifierDefaultPromptResponse,
|
||||
UpdateUsefulLinksRequest,
|
||||
|
|
@ -897,6 +898,12 @@ async def patch_model(
|
|||
existing_litellm_params=db_model.litellm_params,
|
||||
)
|
||||
|
||||
ModelManagementAuthChecks.can_user_set_aws_session_tags(
|
||||
litellm_params=patch_data.litellm_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
existing_litellm_params=db_model.litellm_params,
|
||||
)
|
||||
|
||||
_raise_on_strategy_router_write_violation(
|
||||
incoming_params=patch_data.litellm_params,
|
||||
existing_params=db_model.litellm_params,
|
||||
|
|
@ -1650,6 +1657,10 @@ async def _update_existing_team_model_assignment(
|
|||
# No team_model_add/delete calls required; public name is already registered
|
||||
|
||||
|
||||
def _canonical_session_tags(tags: Sequence[AwsSessionTag]) -> tuple[tuple[str, str], ...]:
|
||||
return tuple(sorted((tag["Key"], tag["Value"]) for tag in tags))
|
||||
|
||||
|
||||
class ModelManagementAuthChecks:
|
||||
"""
|
||||
Common auth checks for model management endpoints
|
||||
|
|
@ -1704,6 +1715,28 @@ class ModelManagementAuthChecks:
|
|||
param="litellm_credential_name",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def can_user_set_aws_session_tags(
|
||||
litellm_params: GenericLiteLLMParams | None,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
existing_litellm_params: GenericLiteLLMParams | None = None,
|
||||
) -> Literal[True]:
|
||||
if litellm_params is None or litellm_params.aws_session_tags is None:
|
||||
return True
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
|
||||
return True
|
||||
existing_tags: Final = existing_litellm_params.aws_session_tags if existing_litellm_params is not None else None
|
||||
if existing_tags is not None and _canonical_session_tags(existing_tags) == _canonical_session_tags(
|
||||
litellm_params.aws_session_tags
|
||||
):
|
||||
return True
|
||||
raise ProxyException(
|
||||
message=f"Only a proxy admin can set aws_session_tags on a model. Your role={user_api_key_dict.user_role}.",
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
code=status.HTTP_403_FORBIDDEN,
|
||||
param="aws_session_tags",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def allow_team_model_action(
|
||||
model_params: Deployment | updateDeployment,
|
||||
|
|
@ -2037,6 +2070,11 @@ async def add_new_model(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
ModelManagementAuthChecks.can_user_set_aws_session_tags(
|
||||
litellm_params=model_params.litellm_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
_raise_on_strategy_router_write_violation(
|
||||
incoming_params=model_params.litellm_params,
|
||||
existing_params=None,
|
||||
|
|
@ -2221,6 +2259,12 @@ async def update_model(
|
|||
existing_litellm_params=deployment.litellm_params,
|
||||
)
|
||||
|
||||
ModelManagementAuthChecks.can_user_set_aws_session_tags(
|
||||
litellm_params=model_params.litellm_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
existing_litellm_params=deployment.litellm_params,
|
||||
)
|
||||
|
||||
_raise_on_strategy_router_write_violation(
|
||||
incoming_params=model_params.litellm_params,
|
||||
existing_params=deployment.litellm_params,
|
||||
|
|
|
|||
|
|
@ -392,6 +392,185 @@ class TestModelManagementAuthChecks:
|
|||
assert exc_info.value.code == "403"
|
||||
mock_update.assert_not_awaited()
|
||||
|
||||
def test_can_user_set_aws_session_tags_admin_success(self):
|
||||
result = ModelManagementAuthChecks.can_user_set_aws_session_tags(
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="bedrock/test_model", aws_session_tags=[{"Key": "team", "Value": "genai"}]
|
||||
),
|
||||
user_api_key_dict=self.admin_user,
|
||||
)
|
||||
assert result is True
|
||||
|
||||
def test_can_user_set_aws_session_tags_without_tags_allows_any_role(self):
|
||||
result = ModelManagementAuthChecks.can_user_set_aws_session_tags(
|
||||
litellm_params=LiteLLM_Params(model="bedrock/test_model", aws_role_name="arn:aws:iam::123:role/x"),
|
||||
user_api_key_dict=self.team_admin_user,
|
||||
)
|
||||
assert result is True
|
||||
|
||||
def test_can_user_set_aws_session_tags_team_admin_fails(self):
|
||||
with pytest.raises(Exception, match="Only a proxy admin can set aws_session_tags") as exc_info:
|
||||
ModelManagementAuthChecks.can_user_set_aws_session_tags(
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="bedrock/test_model", aws_session_tags=[{"Key": "team", "Value": "genai"}]
|
||||
),
|
||||
user_api_key_dict=self.team_admin_user,
|
||||
)
|
||||
assert exc_info.value.code == "403"
|
||||
assert exc_info.value.param == "aws_session_tags"
|
||||
|
||||
def test_can_user_set_aws_session_tags_unchanged_existing_allows_any_role(self):
|
||||
result = ModelManagementAuthChecks.can_user_set_aws_session_tags(
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="bedrock/test_model",
|
||||
aws_session_tags=[{"Key": "team", "Value": "genai"}, {"Key": "env", "Value": "prod"}],
|
||||
),
|
||||
user_api_key_dict=self.team_admin_user,
|
||||
existing_litellm_params=LiteLLM_Params(
|
||||
model="bedrock/test_model",
|
||||
aws_session_tags=[{"Key": "env", "Value": "prod"}, {"Key": "team", "Value": "genai"}],
|
||||
),
|
||||
)
|
||||
assert result is True
|
||||
|
||||
def test_can_user_set_aws_session_tags_changed_value_fails_for_team_admin(self):
|
||||
with pytest.raises(Exception, match="Only a proxy admin can set aws_session_tags") as exc_info:
|
||||
ModelManagementAuthChecks.can_user_set_aws_session_tags(
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="bedrock/test_model", aws_session_tags=[{"Key": "team", "Value": "platform"}]
|
||||
),
|
||||
user_api_key_dict=self.team_admin_user,
|
||||
existing_litellm_params=LiteLLM_Params(
|
||||
model="bedrock/test_model", aws_session_tags=[{"Key": "team", "Value": "genai"}]
|
||||
),
|
||||
)
|
||||
assert exc_info.value.code == "403"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_new_model_rejects_aws_session_tags_for_non_admin(self):
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
add_new_model,
|
||||
)
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch( # test-quality-ok: prior auth check needs a live DB; only the session tag check is under test
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
|
||||
new=AsyncMock(return_value=None),
|
||||
),
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await add_new_model(
|
||||
model_params=Deployment(
|
||||
model_name="tagged-bedrock",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="bedrock/anthropic.claude-opus-4-6-v1:0",
|
||||
aws_role_name="arn:aws:iam::123456789012:role/team-role",
|
||||
aws_session_tags=[{"Key": "team", "Value": "genai"}],
|
||||
),
|
||||
model_info={"id": "session-tags-create-test", "team_id": "test_team"},
|
||||
),
|
||||
user_api_key_dict=self.team_admin_user,
|
||||
)
|
||||
assert exc_info.value.code == "403"
|
||||
assert exc_info.value.param == "aws_session_tags"
|
||||
mock_prisma.db.litellm_proxymodeltable.create.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_patch_model_rejects_aws_session_tags_for_non_admin(self):
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
patch_model,
|
||||
)
|
||||
from litellm.types.router import updateLiteLLMParams
|
||||
|
||||
model_id = "session-tags-patch-test"
|
||||
db_model = Deployment(
|
||||
model_name="tagged-bedrock",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="bedrock/anthropic.claude-opus-4-6-v1:0",
|
||||
aws_role_name="arn:aws:iam::123456789012:role/team-role",
|
||||
),
|
||||
model_info={"id": model_id, "team_id": "test_team"},
|
||||
)
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server.llm_router", MagicMock()), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch( # test-quality-ok: stubs the DB row fetch; only the session tag check is under test
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.get_db_model",
|
||||
new=AsyncMock(return_value=db_model),
|
||||
),
|
||||
patch( # test-quality-ok: prior auth check needs a live DB; only the session tag check is under test
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
|
||||
new=AsyncMock(return_value=None),
|
||||
),
|
||||
patch( # test-quality-ok: asserts the DB write is never reached on rejection
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints._update_team_model_in_db",
|
||||
new=AsyncMock(),
|
||||
) as mock_update,
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await patch_model(
|
||||
model_id=model_id,
|
||||
patch_data=updateDeployment(
|
||||
litellm_params=updateLiteLLMParams(aws_session_tags=[{"Key": "team", "Value": "genai"}])
|
||||
),
|
||||
user_api_key_dict=self.team_admin_user,
|
||||
)
|
||||
assert exc_info.value.code == "403"
|
||||
assert exc_info.value.param == "aws_session_tags"
|
||||
mock_update.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_model_rejects_aws_session_tags_for_non_admin(self):
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
update_model,
|
||||
)
|
||||
from litellm.types.router import updateLiteLLMParams
|
||||
|
||||
model_id = "session-tags-put-test"
|
||||
existing = Deployment(
|
||||
model_name="tagged-bedrock",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="bedrock/anthropic.claude-opus-4-6-v1:0",
|
||||
aws_role_name="arn:aws:iam::123456789012:role/team-role",
|
||||
),
|
||||
model_info={"id": model_id, "team_id": "test_team"},
|
||||
)
|
||||
existing_row = MagicMock()
|
||||
existing_row.model_dump.return_value = existing.model_dump()
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row)
|
||||
mock_prisma.db.litellm_proxymodeltable.update = AsyncMock()
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
|
||||
patch( # test-quality-ok: prior auth check needs a live DB; only the session tag check is under test
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
|
||||
new=AsyncMock(return_value=None),
|
||||
),
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await update_model(
|
||||
model_params=updateDeployment(
|
||||
litellm_params=updateLiteLLMParams(aws_session_tags=[{"Key": "team", "Value": "genai"}]),
|
||||
model_info=ModelInfo(id=model_id),
|
||||
),
|
||||
user_api_key_dict=self.team_admin_user,
|
||||
)
|
||||
assert exc_info.value.code == "403"
|
||||
assert exc_info.value.param == "aws_session_tags"
|
||||
mock_prisma.db.litellm_proxymodeltable.update.assert_not_awaited()
|
||||
|
||||
def test_can_user_attach_credential_internal_user_fails(self):
|
||||
with pytest.raises(Exception, match="Only a proxy admin can attach a stored credential") as exc_info:
|
||||
ModelManagementAuthChecks.can_user_attach_credential(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue