From ccb4b06cd6131083c13bcdecc00468200d6fe688 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 21 Aug 2026 03:24:22 +0000 Subject: [PATCH] feat(rbac): add custom RBAC roles with route allow-lists and inheritance Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../migration.sql | 13 + .../litellm_proxy_extras/schema.prisma | 11 + litellm/proxy/_lazy_features.py | 5 + litellm/proxy/_types.py | 90 +++++-- litellm/proxy/auth/auth_checks.py | 12 +- litellm/proxy/auth/custom_rbac.py | 211 +++++++++++++++ litellm/proxy/auth/route_checks.py | 35 ++- litellm/proxy/auth/user_api_key_auth.py | 3 +- .../custom_rbac_role_endpoints.py | 212 +++++++++++++++ .../internal_user_endpoints.py | 5 + .../team_metadata_validation.py | 4 +- litellm/proxy/schema.prisma | 11 + litellm/repositories/prisma_args.py | 12 + litellm/repositories/table_repositories.py | 4 + litellm/types/custom_rbac.py | 54 ++++ schema.prisma | 11 + .../proxy/auth/test_custom_rbac.py | 250 ++++++++++++++++++ .../test_custom_rbac_role_endpoints.py | 194 ++++++++++++++ 18 files changed, 1103 insertions(+), 34 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260820030000_add_custom_rbac_role_table/migration.sql create mode 100644 litellm/proxy/auth/custom_rbac.py create mode 100644 litellm/proxy/management_endpoints/custom_rbac_role_endpoints.py create mode 100644 litellm/repositories/prisma_args.py create mode 100644 litellm/types/custom_rbac.py create mode 100644 tests/test_litellm/proxy/auth/test_custom_rbac.py create mode 100644 tests/test_litellm/proxy/management_endpoints/test_custom_rbac_role_endpoints.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260820030000_add_custom_rbac_role_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260820030000_add_custom_rbac_role_table/migration.sql new file mode 100644 index 00000000000..d2c7503baa7 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260820030000_add_custom_rbac_role_table/migration.sql @@ -0,0 +1,13 @@ +-- CreateTable +CREATE TABLE "LiteLLM_CustomRBACRoleTable" ( + "role_name" TEXT NOT NULL, + "description" TEXT, + "allowed_routes" TEXT[], + "inherits" TEXT[], + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "created_by" TEXT, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_by" TEXT, + + CONSTRAINT "LiteLLM_CustomRBACRoleTable_pkey" PRIMARY KEY ("role_name") +); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index d9959677116..71b06bcd5c2 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -599,6 +599,17 @@ model LiteLLM_TagTable { updated_at DateTime @default(now()) @updatedAt @map("updated_at") } +model LiteLLM_CustomRBACRoleTable { + role_name String @id + description String? + allowed_routes String[] + inherits String[] + created_at DateTime @default(now()) @map("created_at") + created_by String? + updated_at DateTime @default(now()) @updatedAt @map("updated_at") + updated_by String? +} + // store proxy config.yaml model LiteLLM_Config { param_name String @id diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index fdd15a89aa5..99214463d1e 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -256,6 +256,11 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( module_path="litellm.proxy.management_endpoints.access_group_endpoints", path_prefixes=("/access_group", "/v1/access_group", "/v1/unified_access_group"), ), + LazyFeature( + name="custom_rbac_roles", + module_path="litellm.proxy.management_endpoints.custom_rbac_role_endpoints", + path_prefixes=("/custom_role",), + ), ) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 00cbd13cfdc..7050c7c68ca 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -206,6 +206,54 @@ class LitellmUserRoles(str, enum.Enum): ] +ASSIGNABLE_BUILTIN_USER_ROLES: Final = frozenset( + ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + ) +) + + +def user_role_name(value: "LitellmUserRoles | str | None") -> str | None: + """The wire name of a built-in or custom RBAC role""" + return value.value if isinstance(value, LitellmUserRoles) else value + + +def parse_stored_user_role(value: str) -> "LitellmUserRoles | str": + """A persisted user_role as its built-in enum member, or verbatim when it is a custom RBAC role""" + try: + return LitellmUserRoles(value) + except ValueError: + return value + + +def coerce_assignable_user_role( + value: "LitellmUserRoles | str | None", +) -> "LitellmUserRoles | str | None": + """Normalize a user_role from a management request. + + Built-in roles are returned as enum members and are rejected unless they can be assigned + to a user directly. Any other string is kept verbatim as a custom RBAC role name; whether + such a role actually exists is checked against the configured roles by the endpoint + """ + if value is None: + return None + + try: + role: Final = LitellmUserRoles(value) + except ValueError: + return str(value) + + if role not in ASSIGNABLE_BUILTIN_USER_ROLES: + raise ValueError( + f"user_role={role.value} cannot be assigned to a user. " + f"Allowed built-in roles: {sorted(_role.value for _role in ASSIGNABLE_BUILTIN_USER_ROLES)}" + ) + return role + + class LitellmTableNames(str, enum.Enum): """ Enum for Table Names used by LiteLLM @@ -1674,34 +1722,23 @@ class NewUserRequest(GenerateRequestBase): max_budget: float | None = None user_email: str | None = None user_alias: str | None = None - user_role: ( - Literal[ - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, - ] - | None - ) = None + user_role: LitellmUserRoles | str | None = None teams: list[str] | list[NewUserRequestTeam] | None = None auto_create_key: bool = True # flag used for returning a key as part of the /user/new response send_invite_email: bool | None = None sso_user_id: str | None = None organizations: list[str] | None = None + @field_validator("user_role") + @classmethod + def _check_user_role(cls, value: LitellmUserRoles | str | None) -> LitellmUserRoles | str | None: + return coerce_assignable_user_role(value) + class NewUserResponse(GenerateKeyResponse): max_budget: float | None = None user_email: str | None = None - user_role: ( - Literal[ - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, - ] - | None - ) = None + user_role: LitellmUserRoles | str | None = None teams: list | None = None user_alias: str | None = None model_max_budget: dict | None = None @@ -1714,17 +1751,14 @@ class UpdateUserRequestNoUserIDorEmail(GenerateRequestBase): # shared with Bulk spend: float | None = None metadata: dict | None = None user_alias: str | None = None - user_role: ( - Literal[ - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, - ] - | None - ) = None + user_role: LitellmUserRoles | str | None = None max_budget: float | None = None + @field_validator("user_role") + @classmethod + def _check_user_role(cls, value: LitellmUserRoles | str | None) -> LitellmUserRoles | str | None: + return coerce_assignable_user_role(value) + class UpdateUserRequest(UpdateUserRequestNoUserIDorEmail): # Note: the defaults of all Params here MUST BE NONE @@ -2795,7 +2829,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob """ api_key: str | None = None - user_role: LitellmUserRoles | None = None + user_role: LitellmUserRoles | str | None = None allowed_model_region: AllowedModelRegion | None = None parent_otel_span: Span | None = None rpm_limit_per_model: dict[str, int] | None = None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 12d6b44a648..79eb7591ced 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -63,11 +63,16 @@ from litellm.proxy._types import ( RoleBasedPermissions, SpecialModelNames, UserAPIKeyAuth, + parse_stored_user_role, ) from litellm.proxy.auth.budget_throttle import ( budget_throttle_percentage, should_throttle_budget_exceeded, ) +from litellm.proxy.auth.custom_rbac import ( + CustomRBACEngine, + get_active_custom_rbac_engine, +) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec @@ -1046,6 +1051,7 @@ async def common_checks( request_data=request_body_for_route_check, valid_token=valid_token, user_obj=user_object, + custom_rbac_engine=await get_active_custom_rbac_engine(), ) # 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store @@ -1091,6 +1097,7 @@ def _is_api_route_allowed( request_data: dict, valid_token: UserAPIKeyAuth | None, user_obj: LiteLLM_UserTable | None = None, + custom_rbac_engine: CustomRBACEngine | None = None, ) -> bool: """ - Route b/w api token check and normal token check @@ -1108,6 +1115,7 @@ def _is_api_route_allowed( request=request, request_data=request_data, valid_token=valid_token, + custom_rbac_engine=custom_rbac_engine, ) return True @@ -3007,7 +3015,7 @@ class ExperimentalUIJWTToken: team_id="litellm-dashboard", models=user_info.models, max_parallel_requests=None, - user_role=LitellmUserRoles(user_info.user_role), + user_role=parse_stored_user_role(user_info.user_role), ) return encrypt_value_helper(valid_token.model_dump_json(exclude_none=True)) @@ -3075,7 +3083,7 @@ class ExperimentalUIJWTToken: team_model_aliases=dict(team_model_aliases) if team_model_aliases is not None else None, models=[] if _team_id is not None else user_info.models, max_parallel_requests=None, - user_role=LitellmUserRoles(user_info.user_role), + user_role=parse_stored_user_role(user_info.user_role), is_session_token=True, ) diff --git a/litellm/proxy/auth/custom_rbac.py b/litellm/proxy/auth/custom_rbac.py new file mode 100644 index 00000000000..b8450a6d68b --- /dev/null +++ b/litellm/proxy/auth/custom_rbac.py @@ -0,0 +1,211 @@ +""" +Custom RBAC roles. + +Proxy admins define named roles whose permissions are an allow-list of routes, either in +``general_settings.custom_rbac_roles`` or through the ``/custom_role`` endpoints. Users +assigned such a role are governed entirely by it: any route the role does not grant is +denied, so the built-in role permissions never widen a custom role. +""" + +import time +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Protocol + +from fastapi import HTTPException +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLMRoutes, LitellmUserRoles +from litellm.types.custom_rbac import CustomRBACRole, CustomRBACRoleResponse + +CUSTOM_RBAC_ROLES_CONFIG_KEY: Final = "custom_rbac_roles" +ALL_ROUTES_WILDCARD: Final = "*" +_ROUTE_GROUP_NAMES: Final = frozenset(LiteLLMRoutes.__members__) +_BUILTIN_ROLE_NAMES: Final = frozenset(role.value for role in LitellmUserRoles) +_ENGINE_CACHE_TTL_SECONDS: Final = 30.0 +_CONFIG_ROLES_ADAPTER: Final = TypeAdapter(tuple[CustomRBACRole, ...]) +_ORDER_BY_ROLE_NAME: Final[Mapping[str, object]] = MappingProxyType({"role_name": "asc"}) + + +class _RoleRecord(Protocol): + def dict(self) -> Mapping[str, object]: ... + + +class _CustomRoleTable(Protocol): + async def find_many(self, order: Mapping[str, object]) -> Sequence[_RoleRecord]: ... + + +@dataclass(frozen=True, slots=True) +class CustomRBACEngine: + """Resolved route permissions per custom role, with inheritance already flattened.""" + + effective_routes: Mapping[str, frozenset[str]] + + def is_governed_role(self, role_name: str | None) -> bool: + return role_name is not None and role_name in self.effective_routes + + def is_route_allowed(self, role_name: str, route: str) -> bool: + return any( + _permission_grants_route(permission=permission, route=route) + for permission in self.effective_routes.get(role_name, frozenset()) + ) + + +def _permission_grants_route(permission: str, route: str) -> bool: + from litellm.proxy.auth.route_checks import RouteChecks + + if permission == ALL_ROUTES_WILDCARD: + return True + if permission in _ROUTE_GROUP_NAMES: + return RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes[permission].value) + return RouteChecks.check_route_access(route=route, allowed_routes=(permission,)) + + +def _resolve_routes( + role_name: str, + roles_by_name: Mapping[str, CustomRBACRole], + visited: frozenset[str], +) -> frozenset[str]: + role: Final = roles_by_name.get(role_name) + if role is None or role_name in visited: + return frozenset() + return frozenset(role.allowed_routes).union( + *( + _resolve_routes(role_name=parent, roles_by_name=roles_by_name, visited=visited | frozenset((role_name,))) + for parent in role.inherits + ), + frozenset(), + ) + + +def build_custom_rbac_engine(roles: Sequence[CustomRBACRole]) -> CustomRBACEngine: + roles_by_name: Final = MappingProxyType({role.role_name: role for role in roles}) + return CustomRBACEngine( + effective_routes=MappingProxyType( + { + role_name: _resolve_routes(role_name=role_name, roles_by_name=roles_by_name, visited=frozenset()) + for role_name in roles_by_name + } + ) + ) + + +def is_reserved_role_name(role_name: str) -> bool: + return role_name in _BUILTIN_ROLE_NAMES + + +def validate_role_permissions(allowed_routes: Sequence[str]) -> tuple[str, ...]: + """The entries that are neither ``*``, a route group name, nor a route path.""" + return tuple( + permission + for permission in allowed_routes + if permission != ALL_ROUTES_WILDCARD and permission not in _ROUTE_GROUP_NAMES and not permission.startswith("/") + ) + + +def get_config_custom_rbac_roles() -> tuple[CustomRBACRole, ...]: + from litellm.proxy.proxy_server import general_settings + + configured: Final = general_settings.get(CUSTOM_RBAC_ROLES_CONFIG_KEY) + if not configured: + return () + try: + return _CONFIG_ROLES_ADAPTER.validate_python(configured) + except ValidationError as exc: + verbose_proxy_logger.error("Invalid general_settings.%s: %s", CUSTOM_RBAC_ROLES_CONFIG_KEY, exc) + return () + + +async def get_db_custom_rbac_roles(table: _CustomRoleTable) -> tuple[CustomRBACRoleResponse, ...]: + records: Final = await table.find_many(order=_ORDER_BY_ROLE_NAME) + return tuple(CustomRBACRoleResponse.model_validate(record.dict()) for record in records) + + +class _EngineCache: + def __init__(self, ttl_seconds: float) -> None: + self._ttl_seconds = ttl_seconds + self._engine: CustomRBACEngine | None = None + self._expires_at: float = 0.0 + + def get_fresh(self) -> CustomRBACEngine | None: + if self._engine is None or time.monotonic() >= self._expires_at: + return None + return self._engine + + def get_stale(self) -> CustomRBACEngine | None: + return self._engine + + def set(self, engine: CustomRBACEngine) -> None: + self._engine = engine + self._expires_at = time.monotonic() + self._ttl_seconds + + def clear(self) -> None: + self._engine = None + self._expires_at = 0.0 + + +_ENGINE_CACHE: Final = _EngineCache(ttl_seconds=_ENGINE_CACHE_TTL_SECONDS) + + +def invalidate_custom_rbac_engine_cache() -> None: + _ENGINE_CACHE.clear() + + +def _custom_role_table() -> _CustomRoleTable | None: + from litellm.proxy.proxy_server import prisma_client + from litellm.repositories.table_repositories import CustomRBACRoleRepository + + if prisma_client is None: + return None + return CustomRBACRoleRepository(prisma_client=prisma_client).table + + +async def validate_assigned_user_role(user_role: LitellmUserRoles | str | None) -> None: + """Reject a user_role that is neither a built-in role nor a currently defined custom role.""" + if user_role is None or isinstance(user_role, LitellmUserRoles): + return + + engine: Final = await get_active_custom_rbac_engine() + if engine is not None and engine.is_governed_role(user_role): + return + + raise HTTPException( + status_code=400, + detail=f"user_role={user_role} is not a built-in role and no custom RBAC role with that name is defined", + ) + + +async def get_active_custom_rbac_engine() -> CustomRBACEngine | None: + """The engine for the currently configured roles, or None when no custom role exists. + + A DB read failure reuses the last known policy so a transient outage cannot silently + downgrade a governed role to the built-in role permissions. + """ + cached: Final = _ENGINE_CACHE.get_fresh() + if cached is not None: + return cached + + table: Final = _custom_role_table() + try: + db_roles: Final = () if table is None else await get_db_custom_rbac_roles(table=table) + except Exception as exc: # noqa: BLE001 # any DB failure must keep the last known policy, not drop it + verbose_proxy_logger.exception("Failed to load custom RBAC roles from the DB: %s", exc) + return _ENGINE_CACHE.get_stale() + + roles: Final = get_config_custom_rbac_roles() + tuple( + CustomRBACRole( + role_name=role.role_name, + description=role.description, + allowed_routes=role.allowed_routes, + inherits=role.inherits, + ) + for role in db_roles + ) + if not roles: + return None + + engine: Final = build_custom_rbac_engine(roles=roles) + _ENGINE_CACHE.set(engine) + return engine diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index cea21ca088b..df60ffe8b4d 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -1,6 +1,6 @@ import re from collections.abc import Sequence -from typing import Final +from typing import TYPE_CHECKING, Final from fastapi import HTTPException, Request, status @@ -16,6 +16,9 @@ from litellm.proxy._types import ( from .auth_checks_organization import _user_is_org_admin +if TYPE_CHECKING: + from litellm.proxy.auth.custom_rbac import CustomRBACEngine + # Management write routes denied to PROXY_ADMIN_VIEW_ONLY. Adding a new write # endpoint to a management router REQUIRES adding it here too — the surrounding # check falls through to "allow" if the route is not matched, which previously @@ -239,6 +242,28 @@ class RouteChecks: f"Only proxy admin can be used to generate, delete, update info for new keys/users/teams. Route={route}. Your role={user_role}. Your user_id={masked_user_id}" ) + @staticmethod + def custom_rbac_route_allowed( + user_obj: LiteLLM_UserTable | None, + route: str, + custom_rbac_engine: "CustomRBACEngine | None", + ) -> bool: + """Whether the caller's custom RBAC role grants ``route``. + + False when the caller's role is not governed by a custom role, so the built-in role + checks still apply. Governed roles are default-deny: a route the role does not grant + raises 403 instead of falling through to the built-in permissions + """ + raw_role: Final = user_obj.user_role if user_obj is not None else None + if custom_rbac_engine is None or not custom_rbac_engine.is_governed_role(raw_role) or raw_role is None: + return False + if custom_rbac_engine.is_route_allowed(role_name=raw_role, route=route): + return True + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"user_role={raw_role} is not allowed to access route={route}", + ) + @staticmethod def non_proxy_admin_allowed_routes_check( user_obj: LiteLLM_UserTable | None, @@ -247,6 +272,7 @@ class RouteChecks: request: Request, valid_token: UserAPIKeyAuth, request_data: dict, + custom_rbac_engine: "CustomRBACEngine | None" = None, ): """ Checks if Non Proxy Admin User is allowed to access the route @@ -257,6 +283,13 @@ class RouteChecks: route=route, ) + if RouteChecks.custom_rbac_route_allowed( + user_obj=user_obj, + route=route, + custom_rbac_engine=custom_rbac_engine, + ): + return + if RouteChecks.is_auth_enforced_pass_through_route( route=route, method=RouteChecks._get_request_method(request=request), diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 99592d44f9b..25bcd7d8161 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -32,6 +32,7 @@ from litellm.integrations.otel.runtime import phase_span, seed_request_identity from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.proxy._types import * +from litellm.proxy._types import parse_stored_user_role from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, _cache_key_object, @@ -1379,7 +1380,7 @@ async def _user_api_key_auth_builder( team_rpm_limit=(team_object.rpm_limit if team_object is not None else None), team_models=(team_object.models if team_object is not None else []), user_role=( - LitellmUserRoles(user_object.user_role) + parse_stored_user_role(user_object.user_role) if user_object is not None and user_object.user_role is not None else LitellmUserRoles.INTERNAL_USER ), diff --git a/litellm/proxy/management_endpoints/custom_rbac_role_endpoints.py b/litellm/proxy/management_endpoints/custom_rbac_role_endpoints.py new file mode 100644 index 00000000000..8c91bea2a69 --- /dev/null +++ b/litellm/proxy/management_endpoints/custom_rbac_role_endpoints.py @@ -0,0 +1,212 @@ +"""CRUD endpoints for custom RBAC roles. Proxy admin only.""" + +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Annotated, Final, Protocol + +from fastapi import APIRouter, Depends, HTTPException, status + +from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.custom_rbac import ( + get_config_custom_rbac_roles, + invalidate_custom_rbac_engine_cache, + is_reserved_role_name, + validate_role_permissions, +) +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.utils import get_prisma_client_or_throw +from litellm.repositories.prisma_args import prisma_args, prisma_str_list +from litellm.repositories.table_repositories import CustomRBACRoleRepository +from litellm.types.custom_rbac import ( + CustomRBACRoleCreateRequest, + CustomRBACRoleDeleteRequest, + CustomRBACRoleDeleteResponse, + CustomRBACRoleListResponse, + CustomRBACRoleResponse, + CustomRBACRoleUpdateRequest, +) + +router: Final = APIRouter(tags=["custom rbac role management"]) # mutable-ok: APIRouter concatenates its tags list + +_ORDER_BY_ROLE_NAME: Final[Mapping[str, object]] = MappingProxyType({"role_name": "asc"}) + + +class _RoleRecord(Protocol): + def dict(self) -> Mapping[str, object]: ... + + +class _CustomRoleTable(Protocol): + async def find_unique(self, where: Mapping[str, object]) -> _RoleRecord | None: ... + + async def find_many(self, order: Mapping[str, object]) -> Sequence[_RoleRecord]: ... + + async def create(self, data: Mapping[str, object]) -> _RoleRecord: ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _RoleRecord: ... + + async def delete(self, where: Mapping[str, object]) -> object: ... + + +def _require_proxy_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=CommonProxyErrors.not_allowed_access.value, + ) + + +def _role_table() -> _CustomRoleTable: + prisma_client: Final = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) + return CustomRBACRoleRepository(prisma_client).table + + +def _where_role(role_name: str) -> dict[str, object]: # mutable-ok: prisma requires a plain dict + return prisma_args(MappingProxyType({"role_name": role_name})) + + +def _to_response(record: _RoleRecord) -> CustomRBACRoleResponse: + return CustomRBACRoleResponse.model_validate(record.dict()) + + +def _reject_reserved_or_invalid(role_name: str, allowed_routes: Sequence[str]) -> None: + if is_reserved_role_name(role_name): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"role_name={role_name} is a built-in LiteLLM role and cannot be redefined", + ) + if any(role.role_name == role_name for role in get_config_custom_rbac_roles()): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"role_name={role_name} is defined in general_settings.custom_rbac_roles, edit the config instead", + ) + invalid: Final = validate_role_permissions(allowed_routes=allowed_routes) + if invalid: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"allowed_routes entries must be '*', a LiteLLMRoutes group, or a route path. Invalid: {invalid}", + ) + + +async def _reject_unknown_inherits( + role_name: str, + inherits: Sequence[str], + table: _CustomRoleTable, +) -> None: + if not inherits: + return + known: Final = ( + frozenset(role.role_name for role in get_config_custom_rbac_roles()) + | frozenset(str(record.dict()["role_name"]) for record in await table.find_many(order=_ORDER_BY_ROLE_NAME)) + | frozenset((role_name,)) + ) + unknown: Final = tuple(parent for parent in inherits if parent not in known) + if unknown: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"inherits references roles that do not exist: {unknown}", + ) + + +@router.post("/custom_role/new", response_model=CustomRBACRoleResponse, status_code=status.HTTP_201_CREATED) +async def new_custom_role( + data: CustomRBACRoleCreateRequest, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> CustomRBACRoleResponse: + _require_proxy_admin(user_api_key_dict) + _reject_reserved_or_invalid(role_name=data.role_name, allowed_routes=data.allowed_routes) + + table: Final = _role_table() + if await table.find_unique(where=_where_role(data.role_name)) is not None: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"role_name={data.role_name} already exists", + ) + await _reject_unknown_inherits(role_name=data.role_name, inherits=data.inherits, table=table) + + record: Final = await table.create( + data=prisma_args( + MappingProxyType( + { + "role_name": data.role_name, + "description": data.description, + "allowed_routes": prisma_str_list(data.allowed_routes), + "inherits": prisma_str_list(data.inherits), + "created_by": user_api_key_dict.user_id, + "updated_by": user_api_key_dict.user_id, + } + ) + ) + ) + invalidate_custom_rbac_engine_cache() + return _to_response(record) + + +@router.post("/custom_role/update", response_model=CustomRBACRoleResponse) +async def update_custom_role( + data: CustomRBACRoleUpdateRequest, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> CustomRBACRoleResponse: + _require_proxy_admin(user_api_key_dict) + _reject_reserved_or_invalid(role_name=data.role_name, allowed_routes=data.allowed_routes or ()) + + table: Final = _role_table() + if await table.find_unique(where=_where_role(data.role_name)) is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"role_name={data.role_name} not found", + ) + await _reject_unknown_inherits(role_name=data.role_name, inherits=data.inherits or (), table=table) + + changes: Final = MappingProxyType( + { + key: value + for key, value in ( + ("description", data.description), + ("allowed_routes", None if data.allowed_routes is None else prisma_str_list(data.allowed_routes)), + ("inherits", None if data.inherits is None else prisma_str_list(data.inherits)), + ("updated_by", user_api_key_dict.user_id), + ) + if value is not None + } + ) + record: Final = await table.update(where=_where_role(data.role_name), data=prisma_args(changes)) + invalidate_custom_rbac_engine_cache() + return _to_response(record) + + +@router.get("/custom_role/list", response_model=CustomRBACRoleListResponse) +async def list_custom_roles( + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> CustomRBACRoleListResponse: + _require_proxy_admin(user_api_key_dict) + + config_roles: Final = tuple( + CustomRBACRoleResponse( + role_name=role.role_name, + description=role.description, + allowed_routes=role.allowed_routes, + inherits=role.inherits, + source="config", + ) + for role in get_config_custom_rbac_roles() + ) + records: Final = await _role_table().find_many(order=_ORDER_BY_ROLE_NAME) + return CustomRBACRoleListResponse(roles=config_roles + tuple(_to_response(record) for record in records)) + + +@router.post("/custom_role/delete", response_model=CustomRBACRoleDeleteResponse) +async def delete_custom_role( + data: CustomRBACRoleDeleteRequest, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> CustomRBACRoleDeleteResponse: + _require_proxy_admin(user_api_key_dict) + + table: Final = _role_table() + if await table.find_unique(where=_where_role(data.role_name)) is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"role_name={data.role_name} not found", + ) + await table.delete(where=_where_role(data.role_name)) + invalidate_custom_rbac_engine_cache() + return CustomRBACRoleDeleteResponse(role_name=data.role_name) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 9c725c54d08..19dc493bc9e 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -27,6 +27,7 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import get_team_object, get_user_object +from litellm.proxy.auth.custom_rbac import validate_assigned_user_role from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.user_api_key_cache import ( object_permission_cache_key, @@ -512,6 +513,8 @@ async def new_user( if prisma_client is None: raise HTTPException(status_code=400, detail=CommonProxyErrors.db_not_connected_error.value) + await validate_assigned_user_role(data.user_role) + if prisma_client is None: raise HTTPException( status_code=500, @@ -1578,6 +1581,8 @@ async def user_update( try: verbose_proxy_logger.debug("/user/update: Received data = %s", data) + await validate_assigned_user_role(data.user_role) + response: Final = await _update_single_user_helper( user_request=data, user_api_key_dict=user_api_key_dict, diff --git a/litellm/proxy/management_helpers/team_metadata_validation.py b/litellm/proxy/management_helpers/team_metadata_validation.py index 7bc66c240c7..fd91b111363 100644 --- a/litellm/proxy/management_helpers/team_metadata_validation.py +++ b/litellm/proxy/management_helpers/team_metadata_validation.py @@ -18,7 +18,7 @@ from typing import Final, Literal, Protocol from fastapi import HTTPException, status from pydantic import BaseModel, JsonValue, TypeAdapter -from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth, user_role_name from litellm.types.proxy.management_endpoints.team_endpoints import ( TeamMetadataFieldSchema, ) @@ -181,7 +181,7 @@ async def validate_team_metadata_if_configured( requester=TeamMetadataRequester( user_id=user_api_key_dict.user_id, user_email=user_api_key_dict.user_email, - user_role=user_api_key_dict.user_role.value if user_api_key_dict.user_role is not None else None, + user_role=user_role_name(user_api_key_dict.user_role), ), ) await run_team_metadata_validation( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index d9959677116..71b06bcd5c2 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -599,6 +599,17 @@ model LiteLLM_TagTable { updated_at DateTime @default(now()) @updatedAt @map("updated_at") } +model LiteLLM_CustomRBACRoleTable { + role_name String @id + description String? + allowed_routes String[] + inherits String[] + created_at DateTime @default(now()) @map("created_at") + created_by String? + updated_at DateTime @default(now()) @updatedAt @map("updated_at") + updated_by String? +} + // store proxy config.yaml model LiteLLM_Config { param_name String @id diff --git a/litellm/repositories/prisma_args.py b/litellm/repositories/prisma_args.py new file mode 100644 index 00000000000..c00a93856eb --- /dev/null +++ b/litellm/repositories/prisma_args.py @@ -0,0 +1,12 @@ +"""Prisma's query builder serializes plain ``dict``/``list`` values only, so read-only +mappings and sequences are converted here instead of at every call site.""" + +from collections.abc import Mapping, Sequence + + +def prisma_args(fields: Mapping[str, object]) -> dict[str, object]: # mutable-ok: prisma requires a plain dict + return dict(fields) # mutable-ok: prisma's query builder rejects read-only mappings + + +def prisma_str_list(values: Sequence[str]) -> list[str]: # mutable-ok: prisma requires a plain list + return list(values) # mutable-ok: prisma serializes String[] columns from plain lists diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 131f4d377ef..54867de578b 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -174,6 +174,10 @@ class SkillsRepository(PrismaTableRepository): table_name = "litellm_skillstable" +class CustomRBACRoleRepository(PrismaTableRepository): + table_name = "litellm_customrbacroletable" + + class CacheConfigRepository(PrismaTableRepository): table_name = "litellm_cacheconfig" diff --git a/litellm/types/custom_rbac.py b/litellm/types/custom_rbac.py new file mode 100644 index 00000000000..64ecec52cfb --- /dev/null +++ b/litellm/types/custom_rbac.py @@ -0,0 +1,54 @@ +from datetime import datetime +from typing import Literal, TypeAlias + +from pydantic import BaseModel, Field + +CustomRBACRoleSource: TypeAlias = Literal["config", "db"] + + +class CustomRBACRole(BaseModel): + """A custom RBAC role: a named allow-list of routes, optionally inheriting other custom roles. + + Each entry in ``allowed_routes`` is either ``"*"``, a ``LiteLLMRoutes`` group name + (e.g. ``llm_api_routes``), an exact route, or a wildcard pattern (e.g. ``/team/*``) + """ + + role_name: str + description: str | None = None + allowed_routes: tuple[str, ...] = () + inherits: tuple[str, ...] = () + + +class CustomRBACRoleCreateRequest(BaseModel): + role_name: str = Field(min_length=1) + description: str | None = None + allowed_routes: tuple[str, ...] = () + inherits: tuple[str, ...] = () + + +class CustomRBACRoleUpdateRequest(BaseModel): + role_name: str = Field(min_length=1) + description: str | None = None + allowed_routes: tuple[str, ...] | None = None + inherits: tuple[str, ...] | None = None + + +class CustomRBACRoleDeleteRequest(BaseModel): + role_name: str = Field(min_length=1) + + +class CustomRBACRoleResponse(CustomRBACRole): + source: CustomRBACRoleSource = "db" + created_at: datetime | None = None + updated_at: datetime | None = None + created_by: str | None = None + updated_by: str | None = None + + +class CustomRBACRoleListResponse(BaseModel): + roles: tuple[CustomRBACRoleResponse, ...] + + +class CustomRBACRoleDeleteResponse(BaseModel): + role_name: str + status: Literal["deleted"] = "deleted" diff --git a/schema.prisma b/schema.prisma index d9959677116..71b06bcd5c2 100644 --- a/schema.prisma +++ b/schema.prisma @@ -599,6 +599,17 @@ model LiteLLM_TagTable { updated_at DateTime @default(now()) @updatedAt @map("updated_at") } +model LiteLLM_CustomRBACRoleTable { + role_name String @id + description String? + allowed_routes String[] + inherits String[] + created_at DateTime @default(now()) @map("created_at") + created_by String? + updated_at DateTime @default(now()) @updatedAt @map("updated_at") + updated_by String? +} + // store proxy config.yaml model LiteLLM_Config { param_name String @id diff --git a/tests/test_litellm/proxy/auth/test_custom_rbac.py b/tests/test_litellm/proxy/auth/test_custom_rbac.py new file mode 100644 index 00000000000..ea43dca7db6 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_custom_rbac.py @@ -0,0 +1,250 @@ +import time +from unittest.mock import patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import ( + LiteLLM_UserTable, + LitellmUserRoles, + NewUserRequest, + parse_stored_user_role, +) +from litellm.proxy.auth.custom_rbac import ( + build_custom_rbac_engine, + get_active_custom_rbac_engine, + get_config_custom_rbac_roles, + invalidate_custom_rbac_engine_cache, + is_reserved_role_name, + validate_assigned_user_role, + validate_role_permissions, +) +from litellm.proxy.auth.route_checks import RouteChecks +from litellm.types.custom_rbac import CustomRBACRole + + +def _user(role: str) -> LiteLLM_UserTable: + return LiteLLM_UserTable(user_id="u1", user_role=role, max_budget=None, spend=0.0) + + +class TestCustomRBACEngine: + def test_exact_route_permission(self): + engine = build_custom_rbac_engine( + roles=(CustomRBACRole(role_name="viewer", allowed_routes=("/key/info",)),) + ) + assert engine.is_governed_role("viewer") is True + assert engine.is_route_allowed(role_name="viewer", route="/key/info") is True + assert engine.is_route_allowed(role_name="viewer", route="/key/generate") is False + + def test_route_group_permission(self): + engine = build_custom_rbac_engine( + roles=(CustomRBACRole(role_name="llm-only", allowed_routes=("llm_api_routes",)),) + ) + assert engine.is_route_allowed(role_name="llm-only", route="/chat/completions") is True + assert engine.is_route_allowed(role_name="llm-only", route="/team/new") is False + + def test_wildcard_and_all_routes(self): + engine = build_custom_rbac_engine( + roles=( + CustomRBACRole(role_name="team-mgr", allowed_routes=("/team/*",)), + CustomRBACRole(role_name="superuser", allowed_routes=("*",)), + ) + ) + assert engine.is_route_allowed(role_name="team-mgr", route="/team/new") is True + assert engine.is_route_allowed(role_name="team-mgr", route="/key/generate") is False + assert engine.is_route_allowed(role_name="superuser", route="/key/generate") is True + + def test_inheritance_is_transitive(self): + engine = build_custom_rbac_engine( + roles=( + CustomRBACRole(role_name="base", allowed_routes=("/key/info",)), + CustomRBACRole(role_name="mid", allowed_routes=("/team/info",), inherits=("base",)), + CustomRBACRole(role_name="top", allowed_routes=("/user/info",), inherits=("mid",)), + ) + ) + assert engine.is_route_allowed(role_name="top", route="/key/info") is True + assert engine.is_route_allowed(role_name="top", route="/team/info") is True + assert engine.is_route_allowed(role_name="top", route="/user/info") is True + assert engine.is_route_allowed(role_name="base", route="/user/info") is False + + def test_cyclic_inheritance_terminates(self): + engine = build_custom_rbac_engine( + roles=( + CustomRBACRole(role_name="a", allowed_routes=("/key/info",), inherits=("b",)), + CustomRBACRole(role_name="b", allowed_routes=("/team/info",), inherits=("a",)), + ) + ) + assert engine.is_route_allowed(role_name="a", route="/team/info") is True + assert engine.is_route_allowed(role_name="b", route="/key/info") is True + assert engine.is_route_allowed(role_name="a", route="/user/info") is False + + def test_unknown_role_is_not_governed(self): + engine = build_custom_rbac_engine(roles=(CustomRBACRole(role_name="viewer"),)) + assert engine.is_governed_role("internal_user") is False + assert engine.is_governed_role(None) is False + + +class TestRouteEnforcement: + def test_governed_role_denies_ungranted_route(self): + engine = build_custom_rbac_engine( + roles=(CustomRBACRole(role_name="viewer", allowed_routes=("/key/info",)),) + ) + assert ( + RouteChecks.custom_rbac_route_allowed( + user_obj=_user("viewer"), route="/key/info", custom_rbac_engine=engine + ) + is True + ) + with pytest.raises(HTTPException) as exc: + RouteChecks.custom_rbac_route_allowed( + user_obj=_user("viewer"), route="/key/generate", custom_rbac_engine=engine + ) + assert exc.value.status_code == 403 + + def test_builtin_role_falls_through_to_builtin_checks(self): + engine = build_custom_rbac_engine( + roles=(CustomRBACRole(role_name="viewer", allowed_routes=("/key/info",)),) + ) + assert ( + RouteChecks.custom_rbac_route_allowed( + user_obj=_user(LitellmUserRoles.INTERNAL_USER.value), + route="/key/generate", + custom_rbac_engine=engine, + ) + is False + ) + + def test_no_engine_falls_through(self): + assert ( + RouteChecks.custom_rbac_route_allowed( + user_obj=_user("viewer"), route="/key/generate", custom_rbac_engine=None + ) + is False + ) + + +class TestRoleValidation: + def test_reserved_role_names(self): + assert is_reserved_role_name("internal_user") is True + assert is_reserved_role_name("data-scientist") is False + + def test_invalid_permissions_are_reported(self): + assert validate_role_permissions(allowed_routes=("*", "llm_api_routes", "/key/info")) == () + assert validate_role_permissions(allowed_routes=("key/info", "nonsense")) == ("key/info", "nonsense") + + def test_custom_role_string_survives_request_validation(self): + assert NewUserRequest(user_role="data-scientist").user_role == "data-scientist" + assert NewUserRequest(user_role="internal_user").user_role is LitellmUserRoles.INTERNAL_USER + + def test_non_assignable_builtin_role_is_rejected(self): + with pytest.raises(ValueError): + NewUserRequest(user_role="team") + + def test_stored_role_parsing_keeps_custom_names(self): + assert parse_stored_user_role("internal_user") is LitellmUserRoles.INTERNAL_USER + assert parse_stored_user_role("data-scientist") == "data-scientist" + + +class _FakeRecord: + def __init__(self, role_name: str, allowed_routes: tuple[str, ...]) -> None: + self._role_name = role_name + self._allowed_routes = allowed_routes + + def dict(self) -> dict[str, object]: + return { + "role_name": self._role_name, + "description": None, + "allowed_routes": list(self._allowed_routes), + "inherits": [], + } + + +class _FakeTable: + def __init__(self, records: tuple[_FakeRecord, ...], fail: bool = False) -> None: + self._records = records + self._fail = fail + self.reads = 0 + + async def find_many(self, order): + self.reads += 1 + if self._fail: + raise RuntimeError("db down") + return self._records + + +_CONFIG_KEY = "custom_rbac_roles" +_TABLE_PATH = "litellm.proxy.auth.custom_rbac._custom_role_table" + + +class TestEngineLoading: + def setup_method(self): + invalidate_custom_rbac_engine_cache() + + def teardown_method(self): + invalidate_custom_rbac_engine_cache() + + def test_config_roles_are_parsed(self): + settings = {_CONFIG_KEY: [{"role_name": "cfg", "allowed_routes": ["/key/info"]}]} + with patch("litellm.proxy.proxy_server.general_settings", settings): + assert get_config_custom_rbac_roles() == ( + CustomRBACRole(role_name="cfg", allowed_routes=("/key/info",)), + ) + + def test_malformed_config_roles_are_ignored(self): + with patch("litellm.proxy.proxy_server.general_settings", {_CONFIG_KEY: [{"allowed_routes": 5}]}): + assert get_config_custom_rbac_roles() == () + + @pytest.mark.asyncio + async def test_config_and_db_roles_are_both_active_and_cached(self): + table = _FakeTable(records=(_FakeRecord("db-role", ("/team/info",)),)) + settings = {_CONFIG_KEY: [{"role_name": "cfg-role", "allowed_routes": ["/key/info"]}]} + with ( + patch("litellm.proxy.proxy_server.general_settings", settings), + patch(_TABLE_PATH, return_value=table), + ): + engine = await get_active_custom_rbac_engine() + cached = await get_active_custom_rbac_engine() + + assert engine is not None + assert engine.is_route_allowed(role_name="cfg-role", route="/key/info") is True + assert engine.is_route_allowed(role_name="db-role", route="/team/info") is True + assert engine.is_route_allowed(role_name="db-role", route="/key/generate") is False + assert cached is engine + assert table.reads == 1 + + @pytest.mark.asyncio + async def test_db_failure_reuses_last_known_policy(self): + healthy = _FakeTable(records=(_FakeRecord("db-role", ("/team/info",)),)) + with patch("litellm.proxy.proxy_server.general_settings", {}), patch(_TABLE_PATH, return_value=healthy): + engine = await get_active_custom_rbac_engine() + + broken = _FakeTable(records=(), fail=True) + with ( + patch("litellm.proxy.proxy_server.general_settings", {}), + patch(_TABLE_PATH, return_value=broken), + patch("time.monotonic", return_value=time.monotonic() + 3600), + ): + after_failure = await get_active_custom_rbac_engine() + + assert after_failure is engine + + @pytest.mark.asyncio + async def test_no_roles_means_no_engine(self): + with ( + patch("litellm.proxy.proxy_server.general_settings", {}), + patch(_TABLE_PATH, return_value=_FakeTable(records=())), + ): + assert await get_active_custom_rbac_engine() is None + + @pytest.mark.asyncio + async def test_assigning_undefined_custom_role_is_rejected(self): + table = _FakeTable(records=(_FakeRecord("db-role", ("/team/info",)),)) + with ( + patch("litellm.proxy.proxy_server.general_settings", {}), + patch(_TABLE_PATH, return_value=table), + ): + await validate_assigned_user_role("db-role") + await validate_assigned_user_role(LitellmUserRoles.INTERNAL_USER) + with pytest.raises(HTTPException) as exc: + await validate_assigned_user_role("ghost-role") + assert exc.value.status_code == 400 diff --git a/tests/test_litellm/proxy/management_endpoints/test_custom_rbac_role_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_custom_rbac_role_endpoints.py new file mode 100644 index 00000000000..566f0327b82 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_custom_rbac_role_endpoints.py @@ -0,0 +1,194 @@ +from unittest.mock import patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.management_endpoints.custom_rbac_role_endpoints import ( + delete_custom_role, + new_custom_role, + update_custom_role, +) +from litellm.types.custom_rbac import ( + CustomRBACRoleCreateRequest, + CustomRBACRoleDeleteRequest, + CustomRBACRoleUpdateRequest, +) + +_TABLE_PATH = "litellm.proxy.management_endpoints.custom_rbac_role_endpoints._role_table" + + +class _FakeRecord: + def __init__(self, values: dict): + self._values = values + + def dict(self) -> dict: + return self._values + + +class _FakeTable: + def __init__(self, rows: dict[str, dict] | None = None): + self.rows = dict(rows or {}) + self.deleted: list[str] = [] + + async def find_unique(self, where): + row = self.rows.get(where["role_name"]) + return None if row is None else _FakeRecord(row) + + async def find_many(self, order): + return [_FakeRecord(row) for row in self.rows.values()] + + async def create(self, data): + self.rows[data["role_name"]] = dict(data) + return _FakeRecord(dict(data)) + + async def update(self, where, data): + merged = {**self.rows[where["role_name"]], **data} + self.rows[where["role_name"]] = merged + return _FakeRecord(merged) + + async def delete(self, where): + self.deleted.append(where["role_name"]) + return self.rows.pop(where["role_name"]) + + +def _admin() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + +def _internal_user() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER) + + +@pytest.mark.asyncio +async def test_create_persists_role_and_returns_it(): + table = _FakeTable() + with patch(_TABLE_PATH, return_value=table): + response = await new_custom_role( + data=CustomRBACRoleCreateRequest( + role_name="data-scientist", + allowed_routes=("llm_api_routes", "/key/info"), + ), + user_api_key_dict=_admin(), + ) + assert response.role_name == "data-scientist" + assert table.rows["data-scientist"]["allowed_routes"] == ["llm_api_routes", "/key/info"] + + +@pytest.mark.asyncio +async def test_non_admin_cannot_create_role(): + table = _FakeTable() + with patch(_TABLE_PATH, return_value=table): + with pytest.raises(HTTPException) as exc: + await new_custom_role( + data=CustomRBACRoleCreateRequest(role_name="data-scientist"), + user_api_key_dict=_internal_user(), + ) + assert exc.value.status_code == 403 + assert table.rows == {} + + +@pytest.mark.asyncio +async def test_builtin_role_name_is_rejected(): + table = _FakeTable() + with patch(_TABLE_PATH, return_value=table): + with pytest.raises(HTTPException) as exc: + await new_custom_role( + data=CustomRBACRoleCreateRequest(role_name="internal_user"), + user_api_key_dict=_admin(), + ) + assert exc.value.status_code == 400 + assert table.rows == {} + + +@pytest.mark.asyncio +async def test_invalid_route_permission_is_rejected(): + table = _FakeTable() + with patch(_TABLE_PATH, return_value=table): + with pytest.raises(HTTPException) as exc: + await new_custom_role( + data=CustomRBACRoleCreateRequest(role_name="viewer", allowed_routes=("key/info",)), + user_api_key_dict=_admin(), + ) + assert exc.value.status_code == 400 + assert table.rows == {} + + +@pytest.mark.asyncio +async def test_duplicate_role_name_conflicts(): + table = _FakeTable(rows={"viewer": {"role_name": "viewer", "allowed_routes": [], "inherits": []}}) + with patch(_TABLE_PATH, return_value=table): + with pytest.raises(HTTPException) as exc: + await new_custom_role( + data=CustomRBACRoleCreateRequest(role_name="viewer"), + user_api_key_dict=_admin(), + ) + assert exc.value.status_code == 409 + + +@pytest.mark.asyncio +async def test_update_only_changes_provided_fields(): + table = _FakeTable( + rows={ + "viewer": { + "role_name": "viewer", + "description": "read only", + "allowed_routes": ["/key/info"], + "inherits": [], + } + } + ) + with patch(_TABLE_PATH, return_value=table): + response = await update_custom_role( + data=CustomRBACRoleUpdateRequest(role_name="viewer", allowed_routes=("/team/info",)), + user_api_key_dict=_admin(), + ) + assert response.allowed_routes == ("/team/info",) + assert table.rows["viewer"]["description"] == "read only" + + +@pytest.mark.asyncio +async def test_update_missing_role_is_404(): + with patch(_TABLE_PATH, return_value=_FakeTable()): + with pytest.raises(HTTPException) as exc: + await update_custom_role( + data=CustomRBACRoleUpdateRequest(role_name="ghost"), + user_api_key_dict=_admin(), + ) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_delete_removes_role(): + table = _FakeTable(rows={"viewer": {"role_name": "viewer", "allowed_routes": [], "inherits": []}}) + with patch(_TABLE_PATH, return_value=table): + response = await delete_custom_role( + data=CustomRBACRoleDeleteRequest(role_name="viewer"), + user_api_key_dict=_admin(), + ) + assert (response.role_name, response.status) == ("viewer", "deleted") + assert table.rows == {} + + +@pytest.mark.asyncio +async def test_inheriting_unknown_role_is_rejected(): + table = _FakeTable() + with patch(_TABLE_PATH, return_value=table): + with pytest.raises(HTTPException) as exc: + await new_custom_role( + data=CustomRBACRoleCreateRequest(role_name="viewer", inherits=("ghost",)), + user_api_key_dict=_admin(), + ) + assert exc.value.status_code == 400 + assert table.rows == {} + + +@pytest.mark.asyncio +async def test_inheriting_existing_role_is_allowed(): + table = _FakeTable(rows={"base": {"role_name": "base", "allowed_routes": ["/key/info"], "inherits": []}}) + with patch(_TABLE_PATH, return_value=table): + response = await new_custom_role( + data=CustomRBACRoleCreateRequest(role_name="viewer", inherits=("base",)), + user_api_key_dict=_admin(), + ) + assert response.inherits == ("base",)