From b583f39557c23f7f58360be54018677ee3141694 Mon Sep 17 00:00:00 2001 From: fregataa <44239739+fregataa@users.noreply.github.com> Date: Mon, 1 Jul 2024 02:57:51 +0000 Subject: [PATCH] refactor: impl vfolder RBAC APIs (#2137) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit implement basic functions and APIs for RBAC design **Checklist:** (if applicable) - [x] Milestone metadata specifying the target backport version - [x] Update of end-to-end CLI integration tests in `ai.backend.test` - [x] Documentation - Contents in the `docs` directory - docstrings in public interfaces and type annotations ---- 📚 Documentation preview 📚: https://sorna--2137.org.readthedocs.build/en/2137/ ---- 📚 Documentation preview 📚: https://sorna-ko--2137.org.readthedocs.build/ko/2137/ --- src/ai/backend/manager/api/schema.graphql | 3 + src/ai/backend/manager/models/__init__.py | 3 + src/ai/backend/manager/models/acl.py | 3 +- src/ai/backend/manager/models/group.py | 8 +- src/ai/backend/manager/models/rbac/BUILD | 1 + .../backend/manager/models/rbac/__init__.py | 315 +++++++++++ .../backend/manager/models/rbac/exceptions.py | 6 + src/ai/backend/manager/models/user.py | 2 +- src/ai/backend/manager/models/vfolder.py | 527 +++++++++++++++++- 9 files changed, 859 insertions(+), 9 deletions(-) create mode 100644 src/ai/backend/manager/models/rbac/BUILD create mode 100644 src/ai/backend/manager/models/rbac/__init__.py create mode 100644 src/ai/backend/manager/models/rbac/exceptions.py diff --git a/src/ai/backend/manager/api/schema.graphql b/src/ai/backend/manager/api/schema.graphql index 26d42b3d5b1..c8dcc5126a4 100644 --- a/src/ai/backend/manager/api/schema.graphql +++ b/src/ai/backend/manager/api/schema.graphql @@ -527,6 +527,9 @@ type VirtualFolder implements Item { group: UUID group_name: String creator: String + + """Added in 24.09.0.""" + domain_name: String unmanaged_path: String usage_mode: String permission: String diff --git a/src/ai/backend/manager/models/__init__.py b/src/ai/backend/manager/models/__init__.py index a8bd671340a..9e5a32e4c4b 100644 --- a/src/ai/backend/manager/models/__init__.py +++ b/src/ai/backend/manager/models/__init__.py @@ -8,6 +8,7 @@ from . import image as _image from . import kernel as _kernel from . import keypair as _keypair +from . import rbac as _rbac from . import resource_policy as _rpolicy from . import resource_preset as _rpreset from . import resource_usage as _rusage @@ -33,6 +34,7 @@ *_user.__all__, *_vfolder.__all__, *_dotfile.__all__, + *_rbac.__all__, *_rusage.__all__, *_rpolicy.__all__, *_rpreset.__all__, @@ -54,6 +56,7 @@ from .image import * # noqa from .kernel import * # noqa from .keypair import * # noqa +from .rbac import * # noqa from .resource_policy import * # noqa from .resource_preset import * # noqa from .resource_usage import * # noqa diff --git a/src/ai/backend/manager/models/acl.py b/src/ai/backend/manager/models/acl.py index 4b64b14e7b9..f57f11470c8 100644 --- a/src/ai/backend/manager/models/acl.py +++ b/src/ai/backend/manager/models/acl.py @@ -1,6 +1,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, List, Mapping, Sequence +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, List, Sequence import graphene diff --git a/src/ai/backend/manager/models/group.py b/src/ai/backend/manager/models/group.py index 553007c04ca..e0b0317bbec 100644 --- a/src/ai/backend/manager/models/group.py +++ b/src/ai/backend/manager/models/group.py @@ -94,6 +94,12 @@ MAXIMUM_DOTFILE_SIZE = 64 * 1024 # 61 KiB _rx_slug = re.compile(r"^[a-zA-Z0-9]([a-zA-Z0-9._-]*[a-zA-Z0-9])?$") + +class UserRoleInProject(enum.StrEnum): + ADMIN = enum.auto() # TODO: impl project admin + USER = enum.auto() # User is associated as user + + association_groups_users = sa.Table( "association_groups_users", mapper_registry.metadata, @@ -197,7 +203,7 @@ class GroupRow(Base): users = relationship("AssocGroupUserRow", back_populates="group") resource_policy_row = relationship("ProjectResourcePolicyRow", back_populates="projects") kernels = relationship("KernelRow", back_populates="group_row") - vfolder_row = relationship("VFolderRow", back_populates="group_row") + vfolder_rows = relationship("VFolderRow", back_populates="group_row") def _build_group_query(cond: sa.sql.BinaryExpression, domain_name: str) -> sa.sql.Select: diff --git a/src/ai/backend/manager/models/rbac/BUILD b/src/ai/backend/manager/models/rbac/BUILD new file mode 100644 index 00000000000..db46e8d6c97 --- /dev/null +++ b/src/ai/backend/manager/models/rbac/BUILD @@ -0,0 +1 @@ +python_sources() diff --git a/src/ai/backend/manager/models/rbac/__init__.py b/src/ai/backend/manager/models/rbac/__init__.py new file mode 100644 index 00000000000..44561f140e0 --- /dev/null +++ b/src/ai/backend/manager/models/rbac/__init__.py @@ -0,0 +1,315 @@ +from __future__ import annotations + +import enum +import uuid +from abc import ABCMeta, abstractmethod +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Generic, Sequence, TypeVar + +import sqlalchemy as sa +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm import load_only + +from ..group import AssocGroupUserRow, GroupRow, UserRoleInProject +from ..user import UserRole + +if TYPE_CHECKING: + from ..utils import ExtendedAsyncSAEngine + + +__all__: Sequence[str] = ( + "BasePermission", + "ClientContext", + "DomainScope", + "ProjectScope", + "UserScope", + "StorageHost", + "ImageRegistry", + "ScalingGroup", + "AbstractPermissionContext", + "AbstractPermissionContextBuilder", +) + + +class BasePermission(enum.StrEnum): + pass + + +PermissionType = TypeVar("PermissionType", bound=BasePermission) + + +class Bypass(enum.Enum): + TOKEN = enum.auto() + + +bypass = Bypass.TOKEN + +ProjectContext = Mapping[uuid.UUID, UserRoleInProject] + + +@dataclass +class ClientContext: + db: ExtendedAsyncSAEngine + + domain_name: str + user_id: uuid.UUID + user_role: UserRole + + _project_ctx: ProjectContext | None = field(init=False, default=None) + _domain_project_ctx: Mapping[str, ProjectContext] | None = field(init=False, default=None) + + async def get_or_init_project_ctx_in_domain( + self, db_session: AsyncSession, domain_name: str + ) -> ProjectContext | None: + _project_ctx = await self._get_or_init_project_ctx(db_session) + if _project_ctx is bypass: + # client is superadmin or monitor + if self._domain_project_ctx is None: + self._domain_project_ctx = {} + if domain_name not in self._domain_project_ctx: + stmt = ( + sa.select(GroupRow) + .where(GroupRow.domain_name == domain_name) + .options(load_only(GroupRow.id)) + ) + self._domain_project_ctx = { + **self._domain_project_ctx, + domain_name: { + row.id: UserRoleInProject.ADMIN for row in await db_session.scalars(stmt) + }, + } + else: + # client is domain admin or user + self._domain_project_ctx = {self.domain_name: _project_ctx} + return self._domain_project_ctx.get(domain_name) + + async def get_user_role_in_project( + self, db_session: AsyncSession, project_id: uuid.UUID + ) -> UserRoleInProject | None: + _project_ctx = await self._get_or_init_project_ctx(db_session) + if _project_ctx is bypass: + return UserRoleInProject.ADMIN + else: + return _project_ctx.get(project_id) + + async def _get_or_init_project_ctx(self, db_session: AsyncSession) -> ProjectContext | Bypass: + match self.user_role: + case UserRole.SUPERADMIN | UserRole.MONITOR: + # Superadmins and monitors can access to ALL projects in the system. + # Let's not fetch all project data from DB. + return bypass + case UserRole.ADMIN: + if self._project_ctx is None: + stmt = ( + sa.select(GroupRow) + .where(GroupRow.domain_name == self.domain_name) + .options(load_only(GroupRow.id)) + ) + self._project_ctx = { + row.id: UserRoleInProject.ADMIN for row in await db_session.scalars(stmt) + } + return self._project_ctx + case UserRole.USER: + if self._project_ctx is None: + stmt = ( + sa.select(AssocGroupUserRow) + .select_from(sa.join(AssocGroupUserRow, GroupRow)) + .where( + (AssocGroupUserRow.user_id == self.user_id) + & (GroupRow.domain_name == self.domain_name) + ) + ) + self._project_ctx = { + row.group_id: UserRoleInProject.USER + for row in await db_session.scalars(stmt) + } + return self._project_ctx + + +class BaseScope(metaclass=ABCMeta): + @abstractmethod + def __str__(self) -> str: + pass + + +@dataclass(frozen=True) +class DomainScope(BaseScope): + domain_name: str + + def __str__(self) -> str: + return f"Domain(name: {self.domain_name})" + + +@dataclass(frozen=True) +class ProjectScope(BaseScope): + project_id: uuid.UUID + + def __str__(self) -> str: + return f"Project(id: {self.project_id})" + + +@dataclass(frozen=True) +class UserScope(BaseScope): + user_id: uuid.UUID + + def __str__(self) -> str: + return f"User(id: {self.user_id})" + + +# Extra scope is to address some scopes that contain specific object types +# such as registries for images, scaling groups for agents, storage hosts for vfolders etc. +class ExtraScope: + pass + + +@dataclass(frozen=True) +class StorageHost(ExtraScope): + name: str + + +@dataclass(frozen=True) +class ImageRegistry(ExtraScope): + name: str + + +@dataclass(frozen=True) +class ScalingGroup(ExtraScope): + name: str + + +ObjectType = TypeVar("ObjectType") +ObjectIDType = TypeVar("ObjectIDType") + + +@dataclass +class AbstractPermissionContext( + Generic[PermissionType, ObjectType, ObjectIDType], metaclass=ABCMeta +): + """ + Define permissions under given User, Project or Domain scopes. + Each field of this class represents a mapping of ["accessible scope id", "permissions under the scope"]. + For example, `project` field has a mapping of ["accessible project id", "permissions under the project"]. + { + "PROJECT_A_ID": {"READ", "WRITE", "DELETE"} + "PROJECT_B_ID": {"READ"} + } + + `additional` and `overriding` fields have a mapping of ["object id", "permissions applied to the object"]. + `additional` field is used to add permissions to specific objects. It can be used for admins. + `overriding` field is used to address exceptional cases such as permission overriding or cover other scopes(scaling groups or storage hosts etc). + """ + + user_id_to_permission_map: Mapping[uuid.UUID, frozenset[PermissionType]] = field( + default_factory=dict + ) + project_id_to_permission_map: Mapping[uuid.UUID, frozenset[PermissionType]] = field( + default_factory=dict + ) + domain_name_to_permission_map: Mapping[str, frozenset[PermissionType]] = field( + default_factory=dict + ) + + object_id_to_additional_permission_map: Mapping[ObjectIDType, frozenset[PermissionType]] = ( + field(default_factory=dict) + ) + object_id_to_overriding_permission_map: Mapping[ObjectIDType, frozenset[PermissionType]] = ( + field(default_factory=dict) + ) + + def filter_by_permission(self, permission_to_include: PermissionType) -> None: + self.user_id_to_permission_map = { + uid: permissions + for uid, permissions in self.user_id_to_permission_map.items() + if permission_to_include in permissions + } + self.project_id_to_permission_map = { + pid: permissions + for pid, permissions in self.project_id_to_permission_map.items() + if permission_to_include in permissions + } + self.domain_name_to_permission_map = { + dname: permissions + for dname, permissions in self.domain_name_to_permission_map.items() + if permission_to_include in permissions + } + self.object_id_to_additional_permission_map = { + obj_id: permissions + for obj_id, permissions in self.object_id_to_additional_permission_map.items() + if permission_to_include in permissions + } + self.object_id_to_overriding_permission_map = { + obj_id: permissions + for obj_id, permissions in self.object_id_to_overriding_permission_map.items() + if permission_to_include in permissions + } + + @abstractmethod + async def build_query(self) -> sa.sql.Select | None: + pass + + @abstractmethod + async def calculate_final_permission(self, acl_obj: ObjectType) -> frozenset[PermissionType]: + """ + Calculate the final permissions applied to the given object based on the fields in this class. + """ + pass + + +PermissionContextType = TypeVar("PermissionContextType", bound=AbstractPermissionContext) + + +class AbstractPermissionContextBuilder( + Generic[PermissionType, PermissionContextType], metaclass=ABCMeta +): + @classmethod + async def build( + cls, + db_session: AsyncSession, + ctx: ClientContext, + target_scope: BaseScope, + *, + permission: PermissionType | None = None, + ) -> PermissionContextType: + match target_scope: + case UserScope(user_id=user_id): + result = await cls._build_in_user_scope(db_session, ctx, user_id) + case ProjectScope(project_id=project_id): + result = await cls._build_in_project_scope(db_session, ctx, project_id) + case DomainScope(domain_name=domain_name): + result = await cls._build_in_domain_scope(db_session, ctx, domain_name) + case _: + raise RuntimeError(f"invalid scope `{target_scope}`") + if permission is not None: + result.filter_by_permission(permission) + return result + + @classmethod + @abstractmethod + async def _build_in_user_scope( + cls, + db_session: AsyncSession, + ctx: ClientContext, + user_id: uuid.UUID, + ) -> PermissionContextType: + pass + + @classmethod + @abstractmethod + async def _build_in_project_scope( + cls, + db_session: AsyncSession, + ctx: ClientContext, + project_id: uuid.UUID, + ) -> PermissionContextType: + pass + + @classmethod + @abstractmethod + async def _build_in_domain_scope( + cls, + db_session: AsyncSession, + ctx: ClientContext, + domain_name: str, + ) -> PermissionContextType: + pass diff --git a/src/ai/backend/manager/models/rbac/exceptions.py b/src/ai/backend/manager/models/rbac/exceptions.py new file mode 100644 index 00000000000..d83195ad193 --- /dev/null +++ b/src/ai/backend/manager/models/rbac/exceptions.py @@ -0,0 +1,6 @@ +class RBACException(Exception): + pass + + +class NotEnoughPermission(RBACException): + pass diff --git a/src/ai/backend/manager/models/user.py b/src/ai/backend/manager/models/user.py index 0f3783608e6..1ff6543a915 100644 --- a/src/ai/backend/manager/models/user.py +++ b/src/ai/backend/manager/models/user.py @@ -186,7 +186,7 @@ class UserRow(Base): main_keypair = relationship("KeyPairRow", foreign_keys=users.c.main_access_key) - vfolder_row = relationship("VFolderRow", back_populates="user_row") + vfolder_rows = relationship("VFolderRow", back_populates="user_row") class UserGroup(graphene.ObjectType): diff --git a/src/ai/backend/manager/models/vfolder.py b/src/ai/backend/manager/models/vfolder.py index 8327cf59376..d35d9587049 100644 --- a/src/ai/backend/manager/models/vfolder.py +++ b/src/ai/backend/manager/models/vfolder.py @@ -4,9 +4,21 @@ import logging import os.path import uuid +from collections.abc import Container, Mapping +from dataclasses import dataclass from datetime import datetime from pathlib import PurePosixPath -from typing import TYPE_CHECKING, Any, Final, List, Mapping, NamedTuple, Optional, Sequence, cast +from typing import ( + TYPE_CHECKING, + Any, + Final, + List, + NamedTuple, + Optional, + Sequence, + TypeAlias, + cast, +) import aiohttp import aiotools @@ -71,11 +83,20 @@ metadata, ) from .gql_relay import AsyncNode, Connection, ConnectionResolverResult -from .group import GroupRow, ProjectType +from .group import GroupRow, ProjectType, UserRoleInProject from .minilang.ordering import OrderSpecItem, QueryOrderParser from .minilang.queryfilter import FieldSpecItem, QueryFilterParser, enum_field_getter +from .rbac import ( + AbstractPermissionContext, + AbstractPermissionContextBuilder, + BasePermission, + BaseScope, + ClientContext, + StorageHost, +) +from .rbac.exceptions import NotEnoughPermission from .session import DEAD_SESSION_STATUSES, SessionRow -from .user import UserRole +from .user import UserRole, UserRow from .utils import ExtendedAsyncSAEngine, execute_with_retry, sql_json_merge if TYPE_CHECKING: @@ -88,6 +109,7 @@ "vfolder_invitations", "vfolder_permissions", "VirtualFolder", + "VFolderRBACPermission", "VFolderOwnershipType", "VFolderInvitationState", "VFolderPermission", @@ -116,6 +138,11 @@ "SOFT_DELETED_VFOLDER_STATUSES", "HARD_DELETED_VFOLDER_STATUSES", "VFolderPermissionSetAlias", + "get_vfolders", + "VFolderWithPermissionSet", + "OWNER_PERMISSIONS", + "PermissionContext", + "PermissionContextBuilder", ) @@ -308,7 +335,9 @@ class VFolderCloneInfo(NamedTuple): nullable=False, index=True, ), - sa.Column("permission", EnumValueType(VFolderPermission), default=VFolderPermission.READ_WRITE), + sa.Column( + "permission", EnumValueType(VFolderPermission), default=VFolderPermission.READ_WRITE + ), # legacy sa.Column("max_files", sa.Integer(), default=1000), sa.Column("max_size", sa.Integer(), default=None), # in MBytes sa.Column("num_files", sa.Integer(), default=0), @@ -426,8 +455,8 @@ class VFolderRow(Base): __table__ = vfolders endpoints = relationship("EndpointRow", back_populates="model_row") - user_row = relationship("UserRow", back_populates="vfolder_row") - group_row = relationship("GroupRow", back_populates="vfolder_row") + user_row = relationship("UserRow", back_populates="vfolder_rows") + group_row = relationship("GroupRow", back_populates="vfolder_rows") @classmethod async def get( @@ -653,6 +682,488 @@ async def _append_entries(_query, _is_owner=True): return entries +class VFolderRBACPermission(BasePermission): + # Only owners can do + CLONE = enum.auto() + ASSIGN_PERMISSION_TO_OTHERS = enum.auto() # Invite, share + + # `create_vfolder` action should be in {Domain, Project, or User} permissions, not here + READ_ATTRIBUTE = enum.auto() + UPDATE_ATTRIBUTE = enum.auto() + DELETE_VFOLDER = enum.auto() + + READ_CONTENT = enum.auto() + WRITE_CONTENT = enum.auto() + DELETE_CONTENT = enum.auto() + + MOUNT_RO = enum.auto() + MOUNT_RW = enum.auto() + MOUNT_WD = enum.auto() + + +WhereClauseType: TypeAlias = ( + sa.sql.expression.BinaryExpression | sa.sql.expression.BooleanClauseList +) +# TypeAlias is deprecated since 3.12 + +OWNER_PERMISSIONS: frozenset[VFolderRBACPermission] = frozenset([ + perm for perm in VFolderRBACPermission +]) +ADMIN_PERMISSIONS: frozenset[VFolderRBACPermission] = frozenset([ + VFolderRBACPermission.READ_ATTRIBUTE, + VFolderRBACPermission.UPDATE_ATTRIBUTE, + VFolderRBACPermission.DELETE_VFOLDER, +]) +ADMIN_PERMISSIONS_ON_OTHER_USER_INVITED_FOLDERS: frozenset[VFolderRBACPermission] = frozenset([ + VFolderRBACPermission.READ_ATTRIBUTE, +]) # Admins are allowed to READ folders that other users are invited to. +USER_PERMISSIONS_ON_PROJECT_FOLDERS: frozenset[VFolderRBACPermission] = frozenset([ + VFolderRBACPermission.READ_ATTRIBUTE, + VFolderRBACPermission.READ_CONTENT, + VFolderRBACPermission.WRITE_CONTENT, + VFolderRBACPermission.DELETE_CONTENT, + VFolderRBACPermission.MOUNT_RO, + VFolderRBACPermission.MOUNT_RW, + VFolderRBACPermission.MOUNT_WD, +]) +# `ADMIN_PERMISSIONS_ON_PROJECT_FOLDERS == OWNER_PERMISSIONS` is true +# but it doesn't mean that admins are the owner of the project folders. +ADMIN_PERMISSIONS_ON_PROJECT_FOLDERS: frozenset[VFolderRBACPermission] = ( + ADMIN_PERMISSIONS + | USER_PERMISSIONS_ON_PROJECT_FOLDERS + | {VFolderRBACPermission.CLONE, VFolderRBACPermission.ASSIGN_PERMISSION_TO_OTHERS} +) + +# TODO: Change type of `vfolder_permissions.permission` to VFolderRBACPermission +PERMISSION_TO_RBAC_PERMISSION_MAP: Mapping[VFolderPermission, frozenset[VFolderRBACPermission]] = { + VFolderPermission.READ_ONLY: frozenset([ + VFolderRBACPermission.READ_ATTRIBUTE, + VFolderRBACPermission.READ_CONTENT, + ]), + VFolderPermission.READ_WRITE: frozenset([ + VFolderRBACPermission.READ_ATTRIBUTE, + VFolderRBACPermission.UPDATE_ATTRIBUTE, + VFolderRBACPermission.DELETE_VFOLDER, + VFolderRBACPermission.READ_CONTENT, + VFolderRBACPermission.WRITE_CONTENT, + VFolderRBACPermission.DELETE_CONTENT, + VFolderRBACPermission.MOUNT_RO, + VFolderRBACPermission.MOUNT_RW, + ]), + VFolderPermission.RW_DELETE: frozenset([ + VFolderRBACPermission.READ_ATTRIBUTE, + VFolderRBACPermission.UPDATE_ATTRIBUTE, + VFolderRBACPermission.DELETE_VFOLDER, + VFolderRBACPermission.READ_CONTENT, + VFolderRBACPermission.WRITE_CONTENT, + VFolderRBACPermission.DELETE_CONTENT, + VFolderRBACPermission.MOUNT_RO, + VFolderRBACPermission.MOUNT_RW, + VFolderRBACPermission.MOUNT_WD, + ]), + VFolderPermission.OWNER_PERM: OWNER_PERMISSIONS, +} + + +@dataclass +class PermissionContext(AbstractPermissionContext[VFolderRBACPermission, VFolderRow, uuid.UUID]): + @property + def query_condition(self) -> WhereClauseType | None: + cond: WhereClauseType | None = None + + def _OR_coalesce( + base_cond: WhereClauseType | None, + _cond: sa.sql.expression.BinaryExpression, + ) -> WhereClauseType: + return base_cond | _cond if base_cond is not None else _cond + + def _AND_coalesce( + base_cond: WhereClauseType | None, + _cond: sa.sql.expression.BinaryExpression, + ) -> WhereClauseType: + return base_cond & _cond if base_cond is not None else _cond + + if self.user_id_to_permission_map: + cond = _OR_coalesce(cond, VFolderRow.user.in_(self.user_id_to_permission_map.keys())) + if self.project_id_to_permission_map: + cond = _OR_coalesce( + cond, VFolderRow.group.in_(self.project_id_to_permission_map.keys()) + ) + if self.domain_name_to_permission_map: + cond = _OR_coalesce( + cond, VFolderRow.domain_name.in_(self.domain_name_to_permission_map.keys()) + ) + if self.object_id_to_additional_permission_map: + cond = _OR_coalesce( + cond, VFolderRow.id.in_(self.object_id_to_additional_permission_map.keys()) + ) + if self.object_id_to_overriding_permission_map: + cond = _OR_coalesce( + cond, VFolderRow.id.in_(self.object_id_to_overriding_permission_map.keys()) + ) + + return cond + + async def build_query(self) -> sa.sql.Select | None: + cond = self.query_condition + if cond is None: + return None + return sa.select(VFolderRow).where(cond) + + async def calculate_final_permission( + self, acl_obj: VFolderRow + ) -> frozenset[VFolderRBACPermission]: + vfolder_row = acl_obj + vfolder_id = cast(uuid.UUID, acl_obj.id) + if ( + overriding_perm := self.object_id_to_overriding_permission_map.get(vfolder_id) + ) is not None: + return overriding_perm + permissions: set[VFolderRBACPermission] = set() + permissions |= self.object_id_to_additional_permission_map.get(vfolder_id, set()) + permissions |= self.user_id_to_permission_map.get(vfolder_row.user, set()) + permissions |= self.project_id_to_permission_map.get(vfolder_row.group, set()) + permissions |= self.domain_name_to_permission_map.get(vfolder_row.domain_name, set()) + return frozenset(permissions) + + +class PermissionContextBuilder( + AbstractPermissionContextBuilder[VFolderRBACPermission, PermissionContext] +): + @classmethod + async def _build_in_user_scope( + cls, + db_session: SASession, + ctx: ClientContext, + user_id: uuid.UUID, + ) -> PermissionContext: + user_id_to_permission_map: Mapping[uuid.UUID, frozenset[VFolderRBACPermission]] = {} + object_id_to_additional_permission_map: Mapping[ + uuid.UUID, frozenset[VFolderRBACPermission] + ] = {} + object_id_to_overriding_permission_map: Mapping[ + uuid.UUID, frozenset[VFolderRBACPermission] + ] = {} + match ctx.user_role: + case UserRole.SUPERADMIN | UserRole.MONITOR: + if ctx.user_id == user_id: + additional_stmt = ( + sa.select(VFolderPermissionRow) + .select_from(sa.join(VFolderPermissionRow, VFolderRow)) + .where( + (VFolderPermissionRow.user == ctx.user_id) + & ( + VFolderRow.ownership_type == VFolderOwnershipType.USER + ) # filter out project vfolders + ) + ) + object_id_to_additional_permission_map = { + row.vfolder: PERMISSION_TO_RBAC_PERMISSION_MAP[row.permission] + for row in await db_session.scalars(additional_stmt) + } + user_id_to_permission_map = {user_id: OWNER_PERMISSIONS} + else: + additional_stmt = ( + sa.select(VFolderPermissionRow) + .select_from(sa.join(VFolderPermissionRow, VFolderRow)) + .where( + (VFolderPermissionRow.user == ctx.user_id) + & ( + VFolderRow.ownership_type == VFolderOwnershipType.USER + ) # filter out project vfolders + ) + ) + object_id_to_additional_permission_map = { + row.vfolder: ADMIN_PERMISSIONS_ON_OTHER_USER_INVITED_FOLDERS + for row in await db_session.scalars(additional_stmt) + } + user_id_to_permission_map = {user_id: ADMIN_PERMISSIONS} + case UserRole.ADMIN: + if ctx.user_id == user_id: + additional_stmt = ( + sa.select(VFolderPermissionRow) + .select_from(sa.join(VFolderPermissionRow, VFolderRow)) + .where( + (VFolderPermissionRow.user == ctx.user_id) + & ( + VFolderRow.ownership_type == VFolderOwnershipType.USER + ) # filter out project vfolders + ) + ) + object_id_to_additional_permission_map = { + row.vfolder: PERMISSION_TO_RBAC_PERMISSION_MAP[row.permission] + for row in await db_session.scalars(additional_stmt) + } + + user_id_to_permission_map = {user_id: OWNER_PERMISSIONS} + else: + # domain admins cannot access to users in another domain + user_domain_stmt = ( + sa.select(UserRow) + .where(UserRow.uuid == user_id) + .options(load_only(UserRow.domain_name)) + ) + user_row = cast(UserRow | None, await db_session.scalar(user_domain_stmt)) + if user_row is not None: + if user_row.domain_name == ctx.domain_name: + additional_stmt = ( + sa.select(VFolderPermissionRow) + .select_from(sa.join(VFolderPermissionRow, VFolderRow)) + .where( + (VFolderPermissionRow.user == ctx.user_id) + & ( + VFolderRow.ownership_type == VFolderOwnershipType.USER + ) # filter out project vfolders + ) + ) + object_id_to_additional_permission_map = { + row.vfolder: ADMIN_PERMISSIONS_ON_OTHER_USER_INVITED_FOLDERS + for row in await db_session.scalars(additional_stmt) + } + user_id_to_permission_map = {user_id: ADMIN_PERMISSIONS} + case UserRole.USER: + if ctx.user_id == user_id: + overriding_stmt = ( + sa.select(VFolderPermissionRow) + .select_from(sa.join(VFolderPermissionRow, VFolderRow)) + .where( + (VFolderPermissionRow.user == ctx.user_id) + & ( + VFolderRow.ownership_type == VFolderOwnershipType.USER + ) # filter out project vfolders + ) + ) + object_id_to_overriding_permission_map = { + row.vfolder: PERMISSION_TO_RBAC_PERMISSION_MAP[row.permission] + for row in await db_session.scalars(overriding_stmt) + } + + user_id_to_permission_map = {ctx.user_id: OWNER_PERMISSIONS} + case _: + raise RuntimeError("should not reach here") + return PermissionContext( + user_id_to_permission_map, + object_id_to_additional_permission_map=object_id_to_additional_permission_map, + object_id_to_overriding_permission_map=object_id_to_overriding_permission_map, + ) + + @classmethod + async def _build_in_project_scope( + cls, + db_session: SASession, + ctx: ClientContext, + project_id: uuid.UUID, + ) -> PermissionContext: + project_id_to_permission_map: Mapping[uuid.UUID, frozenset[VFolderRBACPermission]] = {} + object_id_to_overriding_permission_map: Mapping[ + uuid.UUID, frozenset[VFolderRBACPermission] + ] = {} + + role_in_project = await ctx.get_user_role_in_project(db_session, project_id) + match role_in_project: + case UserRoleInProject.ADMIN: + project_id_to_permission_map = {project_id: ADMIN_PERMISSIONS_ON_PROJECT_FOLDERS} + case UserRoleInProject.USER: + project_id_to_permission_map = {project_id: USER_PERMISSIONS_ON_PROJECT_FOLDERS} + overriding_stmt = ( + sa.select(VFolderPermissionRow) + .select_from(sa.join(VFolderPermissionRow, VFolderRow)) + .where( + (VFolderPermissionRow.user == ctx.user_id) + & (VFolderRow.group == project_id) + ) + ) + object_id_to_overriding_permission_map = { + row.vfolder: PERMISSION_TO_RBAC_PERMISSION_MAP[row.permission] + for row in await db_session.scalars(overriding_stmt) + } + case _: + pass + return PermissionContext( + project_id_to_permission_map=project_id_to_permission_map, + object_id_to_overriding_permission_map=object_id_to_overriding_permission_map, + ) + + @classmethod + async def _build_in_domain_scope( + cls, + db_session: SASession, + ctx: ClientContext, + domain_name: str, + ) -> PermissionContext: + user_id_to_permission_map: Mapping[uuid.UUID, frozenset[VFolderRBACPermission]] = {} + project_id_to_permission_map: Mapping[uuid.UUID, frozenset[VFolderRBACPermission]] = {} + domain_name_to_permission_map: Mapping[str, frozenset[VFolderRBACPermission]] = {} + object_id_to_additional_permission_map: Mapping[ + uuid.UUID, frozenset[VFolderRBACPermission] + ] = {} + object_id_to_overriding_permission_map: Mapping[ + uuid.UUID, frozenset[VFolderRBACPermission] + ] = {} + match ctx.user_role: + case UserRole.SUPERADMIN | UserRole.MONITOR: + domain_name_to_permission_map = {domain_name: ADMIN_PERMISSIONS} + + project_ctx = await ctx.get_or_init_project_ctx_in_domain(db_session, domain_name) + if project_ctx is not None: + project_id_to_permission_map = { + project_id: ADMIN_PERMISSIONS_ON_PROJECT_FOLDERS + for project_id, _ in project_ctx.items() + } + additional_stmt = ( + sa.select(VFolderPermissionRow) + .select_from(sa.join(VFolderPermissionRow, VFolderRow)) + .where( + (VFolderPermissionRow.user == ctx.user_id) + & (VFolderRow.domain_name == domain_name) + ) + ) + object_id_to_additional_permission_map = { + row.vfolder: PERMISSION_TO_RBAC_PERMISSION_MAP[row.permission] + for row in await db_session.scalars(additional_stmt) + } + + user_id_to_permission_map = {ctx.user_id: OWNER_PERMISSIONS} + case UserRole.ADMIN: + if ctx.domain_name == domain_name: + domain_name_to_permission_map = {domain_name: ADMIN_PERMISSIONS} + project_ctx = await ctx.get_or_init_project_ctx_in_domain( + db_session, domain_name + ) + if project_ctx is not None: + project_id_to_permission_map = { + project_id: ADMIN_PERMISSIONS_ON_PROJECT_FOLDERS + for project_id, _ in project_ctx.items() + } + additional_stmt = ( + sa.select(VFolderPermissionRow) + .select_from(sa.join(VFolderPermissionRow, VFolderRow)) + .where( + (VFolderPermissionRow.user == ctx.user_id) + & (VFolderRow.domain_name == domain_name) + ) + ) + object_id_to_additional_permission_map = { + row.vfolder: PERMISSION_TO_RBAC_PERMISSION_MAP[row.permission] + for row in await db_session.scalars(additional_stmt) + } + + user_id_to_permission_map = {ctx.user_id: OWNER_PERMISSIONS} + else: + # Only superadmin can access to another domains + pass + case UserRole.USER: + if ctx.domain_name == domain_name: + project_ctx = await ctx.get_or_init_project_ctx_in_domain( + db_session, domain_name + ) + if project_ctx is not None: + project_id_to_permission_map = { + project_id: ADMIN_PERMISSIONS_ON_PROJECT_FOLDERS + if role == UserRoleInProject.ADMIN + else USER_PERMISSIONS_ON_PROJECT_FOLDERS + for project_id, role in project_ctx.items() + } + overriding_stmt = ( + sa.select(VFolderPermissionRow) + .select_from(sa.join(VFolderPermissionRow, VFolderRow)) + .where( + (VFolderPermissionRow.user == ctx.user_id) + & (VFolderRow.domain_name == domain_name) + ) + ) + object_id_to_overriding_permission_map = { + row.vfolder: PERMISSION_TO_RBAC_PERMISSION_MAP[row.permission] + for row in await db_session.scalars(overriding_stmt) + } + + user_id_to_permission_map = {ctx.user_id: OWNER_PERMISSIONS} + else: + # Only superadmin can access to another domains + pass + + case _: + raise RuntimeError("should not reach here") + return PermissionContext( + user_id_to_permission_map, + project_id_to_permission_map, + domain_name_to_permission_map, + object_id_to_additional_permission_map, + object_id_to_overriding_permission_map, + ) + + +class VFolderWithPermissionSet(NamedTuple): + vfolder_row: VFolderRow + permissions: frozenset[VFolderRBACPermission] + + +async def get_vfolders( + db_conn: SAConnection, + ctx: ClientContext, + target_scope: BaseScope, + extra_scope: StorageHost | None = None, + requested_permission: VFolderRBACPermission | None = None, + *, + vfolder_id: uuid.UUID | None = None, + vfolder_name: str | None = None, + usage_mode: VFolderUsageMode | None = None, + allowed_status: Container[VFolderOperationStatus] | None = None, + blocked_status: Container[VFolderOperationStatus] | None = None, +) -> list[VFolderWithPermissionSet]: + async with ctx.db.begin_readonly_session(db_conn) as db_session: + permission_ctx = await PermissionContextBuilder.build( + db_session, ctx, target_scope, permission=requested_permission + ) + query_stmt = await permission_ctx.build_query() + if query_stmt is None: + return [] + if vfolder_id is not None: + query_stmt = query_stmt.where(VFolderRow.id == vfolder_id) + if vfolder_name is not None: + query_stmt = query_stmt.where(VFolderRow.name == vfolder_name) + if usage_mode is not None: + query_stmt = query_stmt.where(VFolderRow.usage_mode == usage_mode) + if allowed_status is not None: + query_stmt = query_stmt.where(VFolderRow.status.in_(allowed_status)) + if blocked_status is not None: + query_stmt = query_stmt.where(VFolderRow.status.not_in(blocked_status)) + + result: list[VFolderWithPermissionSet] = [] + for row in await db_session.scalars(query_stmt): + row = cast(VFolderRow, row) + permissions = await permission_ctx.calculate_final_permission(row) + result.append(VFolderWithPermissionSet(row, permissions)) + return result + + +async def validate_permission( + db_conn: SAConnection, + ctx: ClientContext, + target_scope: BaseScope, + extra_scope: StorageHost | None = None, + *, + permission: VFolderRBACPermission, + vfolder_id: uuid.UUID, +) -> None: + async with ctx.db.begin_readonly_session(db_conn) as db_session: + permission_ctx = await PermissionContextBuilder.build( + db_session, ctx, target_scope, permission=permission + ) + query_stmt = await permission_ctx.build_query() + if query_stmt is None: + raise NotEnoughPermission(f"'{permission.name}' not allowed in {str(target_scope)}") + query_stmt = query_stmt.where(VFolderRow.id == vfolder_id) + vfolder_row = cast(VFolderRow | None, await db_session.scalar(query_stmt)) + if vfolder_row is None: + raise VFolderNotFound( + f"VFolder not found (id:{vfolder_id}, permission:{permission.name})" + ) + final_perms = await permission_ctx.calculate_final_permission(vfolder_row) + if permission not in final_perms: + raise NotEnoughPermission(f"'{permission.name}' not allowed in {str(target_scope)}") + + async def get_allowed_vfolder_hosts_by_group( conn: SAConnection, resource_policy, @@ -1297,6 +1808,7 @@ class Meta: group = graphene.UUID() # Group.id (current owner, null in user vfolders) group_name = graphene.String() # Group.name (current owenr, null in user vfolders) creator = graphene.String() # User.email (always set) + domain_name = graphene.String(description="Added in 24.09.0.") unmanaged_path = graphene.String() usage_mode = graphene.String() permission = graphene.String() @@ -1333,6 +1845,7 @@ def _get_field(name: str) -> Any: group=row["group"], group_name=_get_field("groups_name"), creator=row["creator"], + domain_name=row["domain_name"], unmanaged_path=row["unmanaged_path"], usage_mode=row["usage_mode"], permission=row["permission"], @@ -1361,6 +1874,7 @@ async def resolve_num_files(self, info: graphene.ResolveInfo) -> int: "user": ("vfolders_user", uuid.UUID), "user_email": ("users_email", None), "creator": ("vfolders_creator", None), + "domain_name": ("vfolders_domain_name", None), "unmanaged_path": ("vfolders_unmanaged_path", None), "usage_mode": ( "vfolders_usage_mode", @@ -1392,6 +1906,7 @@ async def resolve_num_files(self, info: graphene.ResolveInfo) -> int: "name": ("vfolders_name", None), "group": ("vfolders_group", None), "group_name": ("groups_name", None), + "domain_name": ("domain_name", None), "user": ("vfolders_user", None), "user_email": ("users_email", None), "creator": ("vfolders_creator", None),