diff --git a/changes/11891.feature.md b/changes/11891.feature.md new file mode 100644 index 00000000000..7ffe88cf0ae --- /dev/null +++ b/changes/11891.feature.md @@ -0,0 +1 @@ +Add a superadmin operation that ensures a user has an RBAC system role, creating it (with permissions and the user-role mapping) when missing. diff --git a/src/ai/backend/manager/data/permission/user_role.py b/src/ai/backend/manager/data/permission/user_role.py index d2b0d624e59..4af1c2c5548 100644 --- a/src/ai/backend/manager/data/permission/user_role.py +++ b/src/ai/backend/manager/data/permission/user_role.py @@ -39,3 +39,9 @@ class UserRoleDataWithRole: deleted_at: datetime | None mapped_role_data: RoleData + + +@dataclass(frozen=True) +class RoleMappingData: + user_id: uuid.UUID + role_data: RoleData diff --git a/src/ai/backend/manager/repositories/permission_controller/db_source/db_source.py b/src/ai/backend/manager/repositories/permission_controller/db_source/db_source.py index 5332036eb1c..308579a98b4 100644 --- a/src/ai/backend/manager/repositories/permission_controller/db_source/db_source.py +++ b/src/ai/backend/manager/repositories/permission_controller/db_source/db_source.py @@ -13,6 +13,7 @@ from ai.backend.common.data.permission.types import ( RBACElementType, RelationType, + RoleSource, ) from ai.backend.logging.utils import BraceStyleAdapter from ai.backend.manager.actions.action.rbac_role_invitation import ( @@ -39,6 +40,7 @@ BulkUserRoleRevocationInput, PermissionResolutionKey, ProjectRoleCount, + RoleData, RoleListResult, RolePermissionsUpdateInput, RoleRevocationResult, @@ -58,6 +60,7 @@ ScopeListResult, ScopeType, ) +from ai.backend.manager.data.permission.user_role import RoleMappingData from ai.backend.manager.data.role_invitation.types import ( RoleInvitationData, RoleInvitationState, @@ -106,6 +109,10 @@ PermissionCreatorSpec, UserRoleCreatorSpec, ) +from ai.backend.manager.repositories.permission_controller.role_manager import ( + RoleManager, + UserSystemRoleSpec, +) from ai.backend.manager.repositories.permission_controller.types import ( PermissionSearchScope, ScopedRoleSearchScope, @@ -148,9 +155,56 @@ class _PermissionGroupKey: class PermissionDBSource: _db: ExtendedAsyncSAEngine + _role_manager: RoleManager def __init__(self, db: ExtendedAsyncSAEngine) -> None: self._db = db + self._role_manager = RoleManager() + + async def ensure_user_system_roles( + self, specs: Collection[UserSystemRoleSpec] + ) -> list[RoleMappingData]: + """Ensure the SYSTEM role(s) for the given scope exist. + + Idempotent: roles already present in the scope (matched by name) are + reused; missing ones are created together with their permissions. For + the USER scope, the user-to-role mapping is ensured as well. + """ + if not specs: + return [] + async with self._db.begin_session_read_committed() as db_session: + existing = await self._existing_user_system_roles(db_session, specs) + results: list[RoleMappingData] = [] + for spec in specs: + role = existing.get(spec.user_id) + if role is None: + # The user has no SYSTEM role yet: create it and map the + # user to it. Existing mappings are left untouched so the + # operation stays idempotent on repeated calls. + role = await self._role_manager.create_system_role(db_session, spec) + await execute_creator( + db_session, + Creator(spec=UserRoleCreatorSpec(user_id=spec.user_id, role_id=role.id)), + ) + results.append(RoleMappingData(user_id=spec.user_id, role_data=role)) + return results + + async def _existing_user_system_roles( + self, db_session: SASession, specs: Iterable[UserSystemRoleSpec] + ) -> dict[uuid.UUID, RoleData]: + stmt = ( + sa.select(UserRoleRow.user_id, RoleRow) + .select_from(UserRoleRow) + .join(RoleRow, UserRoleRow.role_id == RoleRow.id) + .where( + sa.and_( + UserRoleRow.user_id.in_([spec.user_id for spec in specs]), + RoleRow.source == RoleSource.SYSTEM, + ) + ) + ) + result = await db_session.execute(stmt) + return {row.user_id: row.RoleRow.to_data() for row in result} @staticmethod async def _sync_user_scopes_on_assign( diff --git a/src/ai/backend/manager/repositories/permission_controller/repository.py b/src/ai/backend/manager/repositories/permission_controller/repository.py index f7fd7986b1d..a2d029c86bc 100644 --- a/src/ai/backend/manager/repositories/permission_controller/repository.py +++ b/src/ai/backend/manager/repositories/permission_controller/repository.py @@ -50,6 +50,7 @@ ScopeListResult, ScopeType, ) +from ai.backend.manager.data.permission.user_role import RoleMappingData from ai.backend.manager.data.role_invitation.types import RoleInvitationData from ai.backend.manager.models.rbac_models.permission.permission import PermissionRow from ai.backend.manager.models.rbac_models.role import RoleRow @@ -66,6 +67,7 @@ PermissionCreatorSpec, UserRoleCreatorSpec, ) +from ai.backend.manager.repositories.permission_controller.role_manager import UserSystemRoleSpec from ai.backend.manager.repositories.permission_controller.types import ( PermissionSearchScope, ScopedRoleSearchScope, @@ -114,6 +116,17 @@ async def create_role(self, input_data: CreateRoleInput) -> RoleData: role_row = await self._db_source.create_role(input_data) return role_row.to_data() + @permission_controller_repository_resilience.apply() + async def ensure_user_system_roles( + self, specs: Collection[UserSystemRoleSpec] + ) -> list[RoleMappingData]: + """ + Ensure the SYSTEM role(s) for the given scope exist (idempotent). + + Returns the ensured roles. + """ + return await self._db_source.ensure_user_system_roles(specs) + @permission_controller_repository_resilience.apply() async def create_permission( self, diff --git a/src/ai/backend/manager/services/permission_contoller/actions/__init__.py b/src/ai/backend/manager/services/permission_contoller/actions/__init__.py index f1ebfa2a923..6af67b50e5a 100644 --- a/src/ai/backend/manager/services/permission_contoller/actions/__init__.py +++ b/src/ai/backend/manager/services/permission_contoller/actions/__init__.py @@ -12,6 +12,7 @@ from .check_permission import CheckPermissionAction, CheckPermissionActionResult from .create_role import CreateRoleAction, CreateRoleActionResult from .delete_role import DeleteRoleAction, DeleteRoleActionResult +from .ensure_system_role import EnsureSystemRoleAction, EnsureSystemRoleActionResult from .get_permission_matrix import GetPermissionMatrixAction, GetPermissionMatrixActionResult from .get_role_detail import GetRoleDetailAction, GetRoleDetailActionResult from .purge_role import PurgeRoleAction, PurgeRoleActionResult @@ -61,6 +62,8 @@ "CreateRoleActionResult", "DeleteRoleAction", "DeleteRoleActionResult", + "EnsureSystemRoleAction", + "EnsureSystemRoleActionResult", "GetPermissionMatrixAction", "GetPermissionMatrixActionResult", "GetRoleDetailAction", diff --git a/src/ai/backend/manager/services/permission_contoller/actions/ensure_system_role.py b/src/ai/backend/manager/services/permission_contoller/actions/ensure_system_role.py new file mode 100644 index 00000000000..318d240af36 --- /dev/null +++ b/src/ai/backend/manager/services/permission_contoller/actions/ensure_system_role.py @@ -0,0 +1,34 @@ +from collections.abc import Collection +from dataclasses import dataclass +from typing import override + +from ai.backend.manager.actions.action import BaseActionResult +from ai.backend.manager.actions.types import ActionOperationType +from ai.backend.manager.data.permission.user_role import RoleMappingData +from ai.backend.manager.repositories.permission_controller.role_manager import UserSystemRoleSpec +from ai.backend.manager.services.permission_contoller.actions.base import RoleAction + + +@dataclass +class EnsureSystemRoleAction(RoleAction): + """Ensure the SYSTEM role(s) for the given users exist (superadmin only).""" + + specs: Collection[UserSystemRoleSpec] + + @override + def entity_id(self) -> str | None: + return None + + @override + @classmethod + def operation_type(cls) -> ActionOperationType: + return ActionOperationType.CREATE + + +@dataclass +class EnsureSystemRoleActionResult(BaseActionResult): + data: list[RoleMappingData] + + @override + def entity_id(self) -> str | None: + return None diff --git a/src/ai/backend/manager/services/permission_contoller/processors.py b/src/ai/backend/manager/services/permission_contoller/processors.py index d458bce1e83..1e774437b09 100644 --- a/src/ai/backend/manager/services/permission_contoller/processors.py +++ b/src/ai/backend/manager/services/permission_contoller/processors.py @@ -22,6 +22,8 @@ CreateRoleActionResult, DeleteRoleAction, DeleteRoleActionResult, + EnsureSystemRoleAction, + EnsureSystemRoleActionResult, GetRoleDetailAction, GetRoleDetailActionResult, PurgeRoleAction, @@ -107,6 +109,7 @@ class PermissionControllerProcessors(AbstractProcessorPackage): """Processor package for RBAC permission controller operations.""" create_role: ActionProcessor[CreateRoleAction, CreateRoleActionResult] + ensure_system_role: ActionProcessor[EnsureSystemRoleAction, EnsureSystemRoleActionResult] update_role: ActionProcessor[UpdateRoleAction, UpdateRoleActionResult] delete_role: ActionProcessor[DeleteRoleAction, DeleteRoleActionResult] assign_role: ActionProcessor[AssignRoleAction, AssignRoleActionResult] @@ -179,6 +182,7 @@ def __init__( validators: ActionValidators, ) -> None: self.create_role = ActionProcessor(service.create_role, action_monitors) + self.ensure_system_role = ActionProcessor(service.ensure_system_role, action_monitors) self.update_role = ActionProcessor(service.update_role, action_monitors) self.delete_role = ActionProcessor(service.delete_role, action_monitors) self.purge_role = ActionProcessor(service.purge_role, action_monitors) @@ -256,6 +260,7 @@ def __init__( def supported_actions(self) -> list[ActionSpec]: return [ CreateRoleAction.spec(), + EnsureSystemRoleAction.spec(), UpdateRoleAction.spec(), DeleteRoleAction.spec(), PurgeRoleAction.spec(), diff --git a/src/ai/backend/manager/services/permission_contoller/service.py b/src/ai/backend/manager/services/permission_contoller/service.py index d5271f7d78c..a0f7e8d2f84 100644 --- a/src/ai/backend/manager/services/permission_contoller/service.py +++ b/src/ai/backend/manager/services/permission_contoller/service.py @@ -49,6 +49,10 @@ DeleteRoleAction, DeleteRoleActionResult, ) +from ai.backend.manager.services.permission_contoller.actions.ensure_system_role import ( + EnsureSystemRoleAction, + EnsureSystemRoleActionResult, +) from ai.backend.manager.services.permission_contoller.actions.get_entity_types import ( GetEntityTypesAction, GetEntityTypesActionResult, @@ -196,6 +200,15 @@ async def create_role(self, action: CreateRoleAction) -> CreateRoleActionResult: data=result, ) + async def ensure_system_role( + self, action: EnsureSystemRoleAction + ) -> EnsureSystemRoleActionResult: + """ + Ensures the SYSTEM role(s) for the given users exist (idempotent). + """ + roles = await self._repository.ensure_user_system_roles(action.specs) + return EnsureSystemRoleActionResult(data=roles) + async def create_permission( self, action: CreatePermissionAction ) -> CreatePermissionActionResult: diff --git a/tests/unit/manager/repositories/permission_controller/test_ensure_user_system_roles.py b/tests/unit/manager/repositories/permission_controller/test_ensure_user_system_roles.py new file mode 100644 index 00000000000..d36e7a1fff9 --- /dev/null +++ b/tests/unit/manager/repositories/permission_controller/test_ensure_user_system_roles.py @@ -0,0 +1,258 @@ +""" +Tests for PermissionControllerRepository.ensure_user_system_roles(). + +Verifies that the operation creates a SYSTEM role and user-role mapping for a +user that has none, and that it is idempotent (existing roles are reused, no +duplicate role or mapping is created on repeated calls). +""" + +from __future__ import annotations + +import uuid +from collections.abc import AsyncGenerator, Sequence + +import pytest +import sqlalchemy as sa + +from ai.backend.common.data.permission.types import RoleSource +from ai.backend.common.types import BinarySize, ResourceSlot, VFolderHostPermissionMap +from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm +from ai.backend.manager.models.domain import DomainRow +from ai.backend.manager.models.hasher.types import PasswordInfo +from ai.backend.manager.models.keypair import KeyPairRow +from ai.backend.manager.models.rbac_models import RoleRow, UserRoleRow +from ai.backend.manager.models.rbac_models.association_scopes_entities import ( + AssociationScopesEntitiesRow, +) +from ai.backend.manager.models.rbac_models.permission.permission import PermissionRow +from ai.backend.manager.models.resource_policy import ( + KeyPairResourcePolicyRow, + UserResourcePolicyRow, +) +from ai.backend.manager.models.user import UserRole, UserRow, UserStatus +from ai.backend.manager.models.utils import ExtendedAsyncSAEngine +from ai.backend.manager.repositories.permission_controller.repository import ( + PermissionControllerRepository, +) +from ai.backend.manager.repositories.permission_controller.role_manager import UserSystemRoleSpec +from ai.backend.testutils.db import TableOrORM, with_tables + +ENSURE_TABLES: Sequence[TableOrORM] = [ + DomainRow, + UserResourcePolicyRow, + KeyPairResourcePolicyRow, + RoleRow, + UserRoleRow, + UserRow, + KeyPairRow, + AssociationScopesEntitiesRow, + PermissionRow, +] + + +def _password() -> PasswordInfo: + return PasswordInfo( + password="dummy", + algorithm=PasswordHashAlgorithm.PBKDF2_SHA256, + rounds=600_000, + salt_size=32, + ) + + +class TestEnsureUserSystemRoles: + """Behavior of ensure_user_system_roles (create-if-missing, idempotent).""" + + @pytest.fixture + async def db_with_tables( + self, + database_connection: ExtendedAsyncSAEngine, + ) -> AsyncGenerator[ExtendedAsyncSAEngine, None]: + async with with_tables(database_connection, ENSURE_TABLES): + yield database_connection + + @pytest.fixture + def repository( + self, + db_with_tables: ExtendedAsyncSAEngine, + ) -> PermissionControllerRepository: + return PermissionControllerRepository(db_with_tables) + + @pytest.fixture + async def domain_name(self, db_with_tables: ExtendedAsyncSAEngine) -> str: + name = f"test-domain-{uuid.uuid4().hex[:8]}" + async with db_with_tables.begin_session() as session: + session.add( + DomainRow( + name=name, + description="", + is_active=True, + total_resource_slots=ResourceSlot(), + allowed_vfolder_hosts=VFolderHostPermissionMap(), + allowed_docker_registries=[], + ) + ) + await session.flush() + return name + + @pytest.fixture + async def user_resource_policy(self, db_with_tables: ExtendedAsyncSAEngine) -> str: + policy_name = f"user-policy-{uuid.uuid4().hex[:8]}" + async with db_with_tables.begin_session() as session: + session.add( + UserResourcePolicyRow( + name=policy_name, + max_vfolder_count=10, + max_quota_scope_size=BinarySize.finite_from_str("10GiB"), + max_session_count_per_model_session=5, + max_customized_image_count=3, + ) + ) + await session.flush() + return policy_name + + async def _create_user( + self, + db: ExtendedAsyncSAEngine, + *, + domain_name: str, + resource_policy: str, + with_system_role: bool = False, + ) -> uuid.UUID: + user_uuid = uuid.uuid4() + async with db.begin_session() as session: + session.add( + UserRow( + uuid=user_uuid, + username=f"user-{user_uuid.hex[:8]}", + email=f"user-{user_uuid.hex[:8]}@example.com", + password=_password(), + need_password_change=False, + status=UserStatus.ACTIVE, + status_info="active", + domain_name=domain_name, + role=UserRole.USER, + resource_policy=resource_policy, + ) + ) + await session.flush() + if with_system_role: + role_id = uuid.uuid4() + session.add( + RoleRow( + id=role_id, + name=f"user-{user_uuid.hex[:8]}", + source=RoleSource.SYSTEM, + ) + ) + await session.flush() + session.add(UserRoleRow(user_id=user_uuid, role_id=role_id)) + await session.flush() + return user_uuid + + async def _count_system_roles(self, db: ExtendedAsyncSAEngine, user_id: uuid.UUID) -> int: + async with db.begin_readonly_session() as session: + count = await session.scalar( + sa.select(sa.func.count()) + .select_from(UserRoleRow) + .join(RoleRow, UserRoleRow.role_id == RoleRow.id) + .where( + sa.and_( + UserRoleRow.user_id == user_id, + RoleRow.source == RoleSource.SYSTEM, + ) + ) + ) + return count or 0 + + async def test_creates_role_and_mapping_when_missing( + self, + repository: PermissionControllerRepository, + db_with_tables: ExtendedAsyncSAEngine, + domain_name: str, + user_resource_policy: str, + ) -> None: + """A user without a SYSTEM role gets one created together with its mapping.""" + user_id = await self._create_user( + db_with_tables, + domain_name=domain_name, + resource_policy=user_resource_policy, + ) + assert await self._count_system_roles(db_with_tables, user_id) == 0 + + results = await repository.ensure_user_system_roles([UserSystemRoleSpec(user_id=user_id)]) + + assert len(results) == 1 + assert results[0].user_id == user_id + assert results[0].role_data.source == RoleSource.SYSTEM + assert await self._count_system_roles(db_with_tables, user_id) == 1 + async with db_with_tables.begin_readonly_session() as session: + perm_count = await session.scalar( + sa.select(sa.func.count()) + .select_from(PermissionRow) + .where(PermissionRow.role_id == results[0].role_data.id) + ) + assert (perm_count or 0) >= 1 + + async def test_idempotent_reuses_existing_role( + self, + repository: PermissionControllerRepository, + db_with_tables: ExtendedAsyncSAEngine, + domain_name: str, + user_resource_policy: str, + ) -> None: + """When a SYSTEM role already exists, it is reused; repeated calls do not duplicate.""" + user_id = await self._create_user( + db_with_tables, + domain_name=domain_name, + resource_policy=user_resource_policy, + with_system_role=True, + ) + async with db_with_tables.begin_readonly_session() as session: + existing_role_id = await session.scalar( + sa.select(UserRoleRow.role_id).where(UserRoleRow.user_id == user_id) + ) + + first = await repository.ensure_user_system_roles([UserSystemRoleSpec(user_id=user_id)]) + assert first[0].role_data.id == existing_role_id + assert await self._count_system_roles(db_with_tables, user_id) == 1 + + # A second call must not raise (no duplicate mapping) and stays stable. + second = await repository.ensure_user_system_roles([UserSystemRoleSpec(user_id=user_id)]) + assert second[0].role_data.id == existing_role_id + assert await self._count_system_roles(db_with_tables, user_id) == 1 + + async def test_batch_mixed_users( + self, + repository: PermissionControllerRepository, + db_with_tables: ExtendedAsyncSAEngine, + domain_name: str, + user_resource_policy: str, + ) -> None: + """A batch with one existing and one missing role is fully ensured.""" + existing_user = await self._create_user( + db_with_tables, + domain_name=domain_name, + resource_policy=user_resource_policy, + with_system_role=True, + ) + missing_user = await self._create_user( + db_with_tables, + domain_name=domain_name, + resource_policy=user_resource_policy, + ) + + results = await repository.ensure_user_system_roles([ + UserSystemRoleSpec(user_id=existing_user), + UserSystemRoleSpec(user_id=missing_user), + ]) + + assert {r.user_id for r in results} == {existing_user, missing_user} + assert await self._count_system_roles(db_with_tables, existing_user) == 1 + assert await self._count_system_roles(db_with_tables, missing_user) == 1 + + async def test_empty_specs_returns_empty( + self, + repository: PermissionControllerRepository, + ) -> None: + """An empty spec collection is a no-op.""" + assert await repository.ensure_user_system_roles([]) == []