From 8dcfa9fc3cb8a09b1f59469cf4e7fbe9bb352abf Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 4 Jun 2026 20:31:28 -0700 Subject: [PATCH] feat(proxy): auth_v2 slice 4 - policy administration API Adds endpoints to manage casbin policies and role assignments without raw DB inserts: add/remove a permission, add/remove a role assignment, and list rules under /auth/v2/policy/*. Requests go through a validated, normalized core (actions limited to read/write/delete/manage, effects to allow/deny, role and object strings normalized) so malformed rules can't reach the policy table. The policy surface governs itself: those routes are a "policy" resource in the route map, so only a role permitted to manage policy (the bootstrap proxy_admin role) can edit them, with an explicit proxy-admin check as defense in depth. Writes reset the policy snapshot cache so changes take effect immediately. Tests cover rule validation/normalization and the rule<->row conversion. --- litellm/proxy/auth/v2/management_endpoints.py | 159 ++++++++++++++++++ litellm/proxy/auth/v2/policy_admin.py | 54 ++++++ litellm/proxy/auth/v2/route_map.py | 7 + litellm/proxy/proxy_server.py | 3 + .../auth/v2/test_management_endpoints.py | 36 ++++ .../proxy/auth/v2/test_policy_admin.py | 68 ++++++++ 6 files changed, 327 insertions(+) create mode 100644 litellm/proxy/auth/v2/management_endpoints.py create mode 100644 litellm/proxy/auth/v2/policy_admin.py create mode 100644 tests/test_litellm/proxy/auth/v2/test_management_endpoints.py create mode 100644 tests/test_litellm/proxy/auth/v2/test_policy_admin.py diff --git a/litellm/proxy/auth/v2/management_endpoints.py b/litellm/proxy/auth/v2/management_endpoints.py new file mode 100644 index 00000000000..a13b2cd1572 --- /dev/null +++ b/litellm/proxy/auth/v2/management_endpoints.py @@ -0,0 +1,159 @@ +from typing import Any, Dict, List, Optional + +from fastapi import APIRouter, Depends, HTTPException, status +from pydantic import BaseModel + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + +from .policy_admin import ( + PolicyValidationError, + make_assignment_rule, + make_permission_rule, +) +from .policy_store import reset_cache + +router = APIRouter(tags=["auth_v2"]) + + +class PermissionRequest(BaseModel): + role: str + resource: str + action: str + effect: str = "allow" + domain: str = "*" + resource_id: Optional[str] = None + + +class AssignmentRequest(BaseModel): + subject_type: str + subject_id: str + role: str + + +def rule_to_row_data(rule: List[str]) -> Dict[str, str]: + """Convert a casbin rule list (``[ptype, v0, v1, ...]``) to a DB row dict.""" + data: Dict[str, str] = {"ptype": rule[0]} + for index, value in enumerate(rule[1:]): + data[f"v{index}"] = value + return data + + +def row_to_rule(row: Any) -> List[str]: + rule = [row.ptype] + for index in range(6): + value = getattr(row, f"v{index}", None) + if value is not None and value != "": + rule.append(value) + return rule + + +def _require_admin(user_api_key_dict: UserAPIKeyAuth) -> None: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="auth_v2 policy administration requires proxy admin", + ) + + +def _prisma() -> Any: + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="auth_v2 policy administration requires a connected database", + ) + return prisma_client + + +async def _add_rule(rule: List[str]) -> None: + await _prisma().db.litellm_casbinrule.create(data=rule_to_row_data(rule)) + reset_cache() + + +async def _remove_rule(rule: List[str]) -> int: + deleted = await _prisma().db.litellm_casbinrule.delete_many( + where=rule_to_row_data(rule) + ) + reset_cache() + return deleted + + +@router.post("/auth/v2/policy/permission/add") +async def add_permission( + body: PermissionRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + _require_admin(user_api_key_dict) + try: + rule = make_permission_rule( + role=body.role, + resource=body.resource, + action=body.action, + effect=body.effect, + domain=body.domain, + resource_id=body.resource_id, + ) + except PolicyValidationError as e: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) + await _add_rule(rule) + return {"added": rule} + + +@router.post("/auth/v2/policy/permission/remove") +async def remove_permission( + body: PermissionRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + _require_admin(user_api_key_dict) + try: + rule = make_permission_rule( + role=body.role, + resource=body.resource, + action=body.action, + effect=body.effect, + domain=body.domain, + resource_id=body.resource_id, + ) + except PolicyValidationError as e: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) + removed = await _remove_rule(rule) + return {"removed": removed, "rule": rule} + + +@router.post("/auth/v2/policy/assignment/add") +async def add_assignment( + body: AssignmentRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + _require_admin(user_api_key_dict) + try: + rule = make_assignment_rule(body.subject_type, body.subject_id, body.role) + except PolicyValidationError as e: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) + await _add_rule(rule) + return {"added": rule} + + +@router.post("/auth/v2/policy/assignment/remove") +async def remove_assignment( + body: AssignmentRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + _require_admin(user_api_key_dict) + try: + rule = make_assignment_rule(body.subject_type, body.subject_id, body.role) + except PolicyValidationError as e: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) + removed = await _remove_rule(rule) + return {"removed": removed, "rule": rule} + + +@router.get("/auth/v2/policy/list") +async def list_policies( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + _require_admin(user_api_key_dict) + rows = await _prisma().db.litellm_casbinrule.find_many() + return {"rules": [row_to_rule(row) for row in rows]} diff --git a/litellm/proxy/auth/v2/policy_admin.py b/litellm/proxy/auth/v2/policy_admin.py new file mode 100644 index 00000000000..f136a350e7e --- /dev/null +++ b/litellm/proxy/auth/v2/policy_admin.py @@ -0,0 +1,54 @@ +from typing import List, Optional + +VALID_ACTIONS = {"read", "write", "delete", "manage"} +VALID_EFFECTS = {"allow", "deny"} +VALID_SUBJECT_TYPES = {"user", "team"} + + +class PolicyValidationError(ValueError): + """Raised when a policy/assignment request is malformed.""" + + +def _role_token(role: str) -> str: + if not role or not role.strip(): + raise PolicyValidationError("role is required") + return role if role.startswith("role:") else f"role:{role}" + + +def normalize_object(resource: str, resource_id: Optional[str]) -> str: + if not resource or not resource.strip(): + raise PolicyValidationError("resource is required") + return f"{resource}:{resource_id}" if resource_id else f"{resource}:*" + + +def make_permission_rule( + role: str, + resource: str, + action: str, + effect: str = "allow", + domain: str = "*", + resource_id: Optional[str] = None, +) -> List[str]: + """Build a casbin ``p`` rule row from validated, normalized inputs.""" + if action not in VALID_ACTIONS: + raise PolicyValidationError( + f"action must be one of {sorted(VALID_ACTIONS)}, got '{action}'" + ) + if effect not in VALID_EFFECTS: + raise PolicyValidationError( + f"effect must be one of {sorted(VALID_EFFECTS)}, got '{effect}'" + ) + obj = normalize_object(resource, resource_id) + return ["p", _role_token(role), domain or "*", obj, action, effect] + + +def make_assignment_rule(subject_type: str, subject_id: str, role: str) -> List[str]: + """Build a casbin ``g`` rule row binding a user/team subject to a role.""" + if subject_type not in VALID_SUBJECT_TYPES: + raise PolicyValidationError( + f"subject_type must be one of {sorted(VALID_SUBJECT_TYPES)}, " + f"got '{subject_type}'" + ) + if not subject_id or not subject_id.strip(): + raise PolicyValidationError("subject_id is required") + return ["g", f"{subject_type}:{subject_id}", _role_token(role)] diff --git a/litellm/proxy/auth/v2/route_map.py b/litellm/proxy/auth/v2/route_map.py index fdc7788101a..6718a43d153 100644 --- a/litellm/proxy/auth/v2/route_map.py +++ b/litellm/proxy/auth/v2/route_map.py @@ -28,6 +28,13 @@ _GOVERNED: Dict[str, GovernedRoute] = { "/team/update": GovernedRoute("team", "write", _TEAM_ID_FIELDS), "/team/delete": GovernedRoute("team", "delete", _TEAM_ID_FIELDS), "/team/info": GovernedRoute("team", "read", _TEAM_ID_FIELDS), + # The policy-admin surface governs itself: only a role permitted to manage the + # "policy" resource (the bootstrap proxy_admin role does) may edit policies. + "/auth/v2/policy/permission/add": GovernedRoute("policy", "write"), + "/auth/v2/policy/permission/remove": GovernedRoute("policy", "delete"), + "/auth/v2/policy/assignment/add": GovernedRoute("policy", "write"), + "/auth/v2/policy/assignment/remove": GovernedRoute("policy", "delete"), + "/auth/v2/policy/list": GovernedRoute("policy", "read"), } diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7aed9ad894a..57e719be139 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15655,8 +15655,11 @@ async def get_routes(): # return {"token": token} +from litellm.proxy.auth.v2.management_endpoints import router as auth_v2_router + app.include_router(router) app.include_router(response_router) +app.include_router(auth_v2_router) app.include_router(public_endpoints_router) app.include_router(rerank_router) app.include_router(ocr_router) diff --git a/tests/test_litellm/proxy/auth/v2/test_management_endpoints.py b/tests/test_litellm/proxy/auth/v2/test_management_endpoints.py new file mode 100644 index 00000000000..8270a830d7d --- /dev/null +++ b/tests/test_litellm/proxy/auth/v2/test_management_endpoints.py @@ -0,0 +1,36 @@ +from litellm.proxy.auth.v2.management_endpoints import rule_to_row_data, row_to_rule + + +class _Row: + def __init__(self, ptype, *values): + self.ptype = ptype + for i in range(6): + setattr(self, f"v{i}", values[i] if i < len(values) else None) + + +def test_permission_rule_to_row(): + data = rule_to_row_data(["p", "role:x", "*", "model:*", "read", "allow"]) + assert data == { + "ptype": "p", + "v0": "role:x", + "v1": "*", + "v2": "model:*", + "v3": "read", + "v4": "allow", + } + + +def test_assignment_rule_to_row(): + data = rule_to_row_data(["g", "user:u1", "role:x"]) + assert data == {"ptype": "g", "v0": "user:u1", "v1": "role:x"} + + +def test_row_to_rule_trims_empty_columns(): + row = _Row("g", "user:u1", "role:x") + assert row_to_rule(row) == ["g", "user:u1", "role:x"] + + +def test_round_trip_permission_rule(): + rule = ["p", "role:x", "team:eng", "model:gpt-4o", "write", "allow"] + row = _Row(*([rule[0]] + rule[1:])) + assert row_to_rule(row) == rule diff --git a/tests/test_litellm/proxy/auth/v2/test_policy_admin.py b/tests/test_litellm/proxy/auth/v2/test_policy_admin.py new file mode 100644 index 00000000000..f66c61d3f78 --- /dev/null +++ b/tests/test_litellm/proxy/auth/v2/test_policy_admin.py @@ -0,0 +1,68 @@ +import pytest + +from litellm.proxy.auth.v2.policy_admin import ( + PolicyValidationError, + make_assignment_rule, + make_permission_rule, + normalize_object, +) + + +def test_permission_rule_shape(): + rule = make_permission_rule("model_reader", "model", "read") + assert rule == ["p", "role:model_reader", "*", "model:*", "read", "allow"] + + +def test_permission_rule_with_resource_id_and_domain(): + rule = make_permission_rule( + "gpt_owner", "model", "write", domain="team:eng", resource_id="gpt-4o" + ) + assert rule == ["p", "role:gpt_owner", "team:eng", "model:gpt-4o", "write", "allow"] + + +def test_role_prefix_is_idempotent(): + assert make_permission_rule("role:x", "model", "read")[1] == "role:x" + + +def test_deny_effect_is_supported(): + assert make_permission_rule("r", "model", "read", effect="deny")[-1] == "deny" + + +def test_invalid_action_is_rejected(): + with pytest.raises(PolicyValidationError): + make_permission_rule("r", "model", "execute") + + +def test_invalid_effect_is_rejected(): + with pytest.raises(PolicyValidationError): + make_permission_rule("r", "model", "read", effect="maybe") + + +def test_empty_resource_is_rejected(): + with pytest.raises(PolicyValidationError): + make_permission_rule("r", "", "read") + + +def test_empty_role_is_rejected(): + with pytest.raises(PolicyValidationError): + make_permission_rule(" ", "model", "read") + + +def test_normalize_object_wildcard_default(): + assert normalize_object("team", None) == "team:*" + assert normalize_object("team", "eng") == "team:eng" + + +def test_assignment_rule_for_user_and_team(): + assert make_assignment_rule("user", "u1", "admin") == ["g", "user:u1", "role:admin"] + assert make_assignment_rule("team", "eng", "role:x") == ["g", "team:eng", "role:x"] + + +def test_assignment_rejects_unknown_subject_type(): + with pytest.raises(PolicyValidationError): + make_assignment_rule("org", "o1", "admin") + + +def test_assignment_requires_subject_id(): + with pytest.raises(PolicyValidationError): + make_assignment_rule("user", "", "admin")