From dc501a8fe6f3f43477f92cac9cc6338269426201 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 6 Nov 2024 06:39:56 +0000 Subject: [PATCH 1/9] feat: Implement image RBAC project scope --- src/ai/backend/manager/models/image.py | 35 +++++++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/src/ai/backend/manager/models/image.py b/src/ai/backend/manager/models/image.py index 57fe7e63ee..344d36b67a 100644 --- a/src/ai/backend/manager/models/image.py +++ b/src/ai/backend/manager/models/image.py @@ -798,13 +798,46 @@ async def _in_domain_scope( object_id_to_additional_permission_map=image_id_permission_map ) + async def _in_project_scope( + self, + ctx: ClientContext, + scope: ProjectScope, + ) -> ImagePermissionContext: + from .domain import DomainRow + from .group import GroupRow + + permissions = await self.calculate_permission(ctx, scope) + image_id_permission_map: dict[UUID, frozenset[ImagePermission]] = {} + + _group_query_stmt = ( + sa.select(GroupRow) + .where(GroupRow.id == scope.project_id) + .options(joinedload(GroupRow.domain)) + ) + group_row = cast(Optional[GroupRow], await self.db_session.scalar(_group_query_stmt)) + if group_row is None: + raise InvalidScope(f"Project not found (n:{scope.project_id})") + + domain_row = sa.select(DomainRow).where(DomainRow.id == group_row.domain.name) + + allowed_registries: set[str] = set(domain_row.allowed_docker_registries) + _img_query_stmt = sa.select(ImageRow).options(load_only(ImageRow.id, ImageRow.registry)) + for row in await self.db_session.scalars(_img_query_stmt): + _row = cast(ImageRow, row) + if _row.registry in allowed_registries: + image_id_permission_map[_row.id] = permissions + + return ImagePermissionContext( + object_id_to_additional_permission_map=image_id_permission_map + ) + @override async def build_ctx_in_project_scope( self, ctx: ClientContext, scope: ProjectScope, ) -> ImagePermissionContext: - return ImagePermissionContext() + return await self._in_project_scope(ctx, scope) @override async def build_ctx_in_user_scope( From 4ce77942f8ba0cd5cdfe0cf3d805d7390f298774 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Wed, 6 Nov 2024 07:59:12 +0000 Subject: [PATCH 2/9] fix: Wrong impl of `_in_project_scope` --- src/ai/backend/manager/models/image.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/ai/backend/manager/models/image.py b/src/ai/backend/manager/models/image.py index 344d36b67a..821b88f7b3 100644 --- a/src/ai/backend/manager/models/image.py +++ b/src/ai/backend/manager/models/image.py @@ -818,7 +818,10 @@ async def _in_project_scope( if group_row is None: raise InvalidScope(f"Project not found (n:{scope.project_id})") - domain_row = sa.select(DomainRow).where(DomainRow.id == group_row.domain.name) + _domain_query_stmt = sa.select(DomainRow).where(DomainRow.name == group_row.domain.name) + domain_row = cast(Optional[DomainRow], await self.db_session.scalar(_domain_query_stmt)) + if domain_row is None: + raise InvalidScope(f"Domain not found (n:{scope.domain_name})") allowed_registries: set[str] = set(domain_row.allowed_docker_registries) _img_query_stmt = sa.select(ImageRow).options(load_only(ImageRow.id, ImageRow.registry)) From 4056a61b8ff05939431047fb26d6b46ed8592d65 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 11 Nov 2024 02:47:05 +0000 Subject: [PATCH 3/9] feat: Implement `build_ctx_in_project_scope` using `AssociationContainerRegistriesGroups` --- src/ai/backend/manager/models/image.py | 46 +++++++++++++++++--------- 1 file changed, 30 insertions(+), 16 deletions(-) diff --git a/src/ai/backend/manager/models/image.py b/src/ai/backend/manager/models/image.py index 821b88f7b3..c049c2ab5f 100644 --- a/src/ai/backend/manager/models/image.py +++ b/src/ai/backend/manager/models/image.py @@ -37,6 +37,9 @@ ) from ai.backend.common.utils import join_non_empty from ai.backend.logging import BraceStyleAdapter +from ai.backend.manager.models.association_container_registries_groups import ( + AssociationContainerRegistriesGroupsRow, +) from ai.backend.manager.models.container_registry import ContainerRegistryRow from ..api.exceptions import ImageNotFound @@ -803,32 +806,43 @@ async def _in_project_scope( ctx: ClientContext, scope: ProjectScope, ) -> ImagePermissionContext: - from .domain import DomainRow from .group import GroupRow permissions = await self.calculate_permission(ctx, scope) image_id_permission_map: dict[UUID, frozenset[ImagePermission]] = {} - _group_query_stmt = ( - sa.select(GroupRow) - .where(GroupRow.id == scope.project_id) - .options(joinedload(GroupRow.domain)) - ) - group_row = cast(Optional[GroupRow], await self.db_session.scalar(_group_query_stmt)) + group_query_stmt = sa.select(GroupRow).where(GroupRow.id == scope.project_id) + group_row = cast(Optional[GroupRow], await self.db_session.scalar(group_query_stmt)) if group_row is None: raise InvalidScope(f"Project not found (n:{scope.project_id})") - _domain_query_stmt = sa.select(DomainRow).where(DomainRow.name == group_row.domain.name) - domain_row = cast(Optional[DomainRow], await self.db_session.scalar(_domain_query_stmt)) - if domain_row is None: - raise InvalidScope(f"Domain not found (n:{scope.domain_name})") + image_select_stmt = ( + sa.select(ImageRow) + .select_from( + sa.join( + ImageRow, ContainerRegistryRow, ImageRow.registry_id == ContainerRegistryRow.id + ) + ) + .where( + sa.or_( + ContainerRegistryRow.is_global, + sa.and_( + not ContainerRegistryRow.is_global, + sa.exists().where( + (AssociationContainerRegistriesGroupsRow.group_id == scope.project_id) + & ( + AssociationContainerRegistriesGroupsRow.container_registry_id + == ImageRow.registry_id + ) + ), + ), + ) + ) + ) - allowed_registries: set[str] = set(domain_row.allowed_docker_registries) - _img_query_stmt = sa.select(ImageRow).options(load_only(ImageRow.id, ImageRow.registry)) - for row in await self.db_session.scalars(_img_query_stmt): + for row in await self.db_session.scalars(image_select_stmt): _row = cast(ImageRow, row) - if _row.registry in allowed_registries: - image_id_permission_map[_row.id] = permissions + image_id_permission_map[_row.id] = permissions return ImagePermissionContext( object_id_to_additional_permission_map=image_id_permission_map From 0bd95ce7f8e03c232cefc6a4c2458ecb83fdae4a Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 11 Nov 2024 02:57:29 +0000 Subject: [PATCH 4/9] fix: Wrong check of boolean column --- src/ai/backend/manager/models/image.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/ai/backend/manager/models/image.py b/src/ai/backend/manager/models/image.py index c049c2ab5f..083364a3d5 100644 --- a/src/ai/backend/manager/models/image.py +++ b/src/ai/backend/manager/models/image.py @@ -24,6 +24,7 @@ import trafaret as t from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession from sqlalchemy.orm import joinedload, load_only, relationship, selectinload +from sqlalchemy.sql.expression import false, true from ai.backend.common.docker import ImageRef from ai.backend.common.exception import UnknownImageReference @@ -825,9 +826,9 @@ async def _in_project_scope( ) .where( sa.or_( - ContainerRegistryRow.is_global, + ContainerRegistryRow.is_global == true(), sa.and_( - not ContainerRegistryRow.is_global, + ContainerRegistryRow.is_global == false(), sa.exists().where( (AssociationContainerRegistriesGroupsRow.group_id == scope.project_id) & ( From 736c812e29db885666822f5bb41143cf50e856bd Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 11 Nov 2024 03:05:37 +0000 Subject: [PATCH 5/9] fix: Rename column `container_registry_id` -> `registry_id` --- src/ai/backend/manager/models/image.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ai/backend/manager/models/image.py b/src/ai/backend/manager/models/image.py index 083364a3d5..69ef778c0d 100644 --- a/src/ai/backend/manager/models/image.py +++ b/src/ai/backend/manager/models/image.py @@ -832,7 +832,7 @@ async def _in_project_scope( sa.exists().where( (AssociationContainerRegistriesGroupsRow.group_id == scope.project_id) & ( - AssociationContainerRegistriesGroupsRow.container_registry_id + AssociationContainerRegistriesGroupsRow.registry_id == ImageRow.registry_id ) ), From 5617f70d5c59a8f0ef0bbfdd7ff61fa37c95c1a9 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Mon, 11 Nov 2024 08:42:00 +0000 Subject: [PATCH 6/9] fix: Error message --- src/ai/backend/manager/models/image.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ai/backend/manager/models/image.py b/src/ai/backend/manager/models/image.py index 69ef778c0d..45c9191fd6 100644 --- a/src/ai/backend/manager/models/image.py +++ b/src/ai/backend/manager/models/image.py @@ -815,7 +815,7 @@ async def _in_project_scope( group_query_stmt = sa.select(GroupRow).where(GroupRow.id == scope.project_id) group_row = cast(Optional[GroupRow], await self.db_session.scalar(group_query_stmt)) if group_row is None: - raise InvalidScope(f"Project not found (n:{scope.project_id})") + raise InvalidScope(f"Project not found (project_id:{scope.project_id})") image_select_stmt = ( sa.select(ImageRow) From d46a0026a093cd0fe376800cb3f8a1f2ee750798 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Tue, 12 Nov 2024 05:22:21 +0000 Subject: [PATCH 7/9] chore: Add news fragment --- changes/3035.feature.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/3035.feature.md diff --git a/changes/3035.feature.md b/changes/3035.feature.md new file mode 100644 index 0000000000..e915e35e81 --- /dev/null +++ b/changes/3035.feature.md @@ -0,0 +1 @@ +Add project scope implementation to Image RBAC. From de062c5dd894263aa3b17adda7bc1d25c631b2e7 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 06:05:30 +0000 Subject: [PATCH 8/9] chore: Fix import statements --- src/ai/backend/manager/models/image.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/ai/backend/manager/models/image.py b/src/ai/backend/manager/models/image.py index 45c9191fd6..994ed8dc85 100644 --- a/src/ai/backend/manager/models/image.py +++ b/src/ai/backend/manager/models/image.py @@ -38,13 +38,13 @@ ) from ai.backend.common.utils import join_non_empty from ai.backend.logging import BraceStyleAdapter -from ai.backend.manager.models.association_container_registries_groups import ( - AssociationContainerRegistriesGroupsRow, -) -from ai.backend.manager.models.container_registry import ContainerRegistryRow from ..api.exceptions import ImageNotFound from ..container_registry import get_container_registry_cls +from ..models.association_container_registries_groups import ( + AssociationContainerRegistriesGroupsRow, +) +from ..models.container_registry import ContainerRegistryRow from .base import ( GUID, Base, From b15ab98d1a161816708c2fbd4b96e13830e59ba0 Mon Sep 17 00:00:00 2001 From: Gyubong Lee Date: Thu, 14 Nov 2024 06:37:27 +0000 Subject: [PATCH 9/9] fix: Improve `_in_project_scope` using relationship --- src/ai/backend/manager/models/image.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/src/ai/backend/manager/models/image.py b/src/ai/backend/manager/models/image.py index 994ed8dc85..19dcb147ac 100644 --- a/src/ai/backend/manager/models/image.py +++ b/src/ai/backend/manager/models/image.py @@ -41,9 +41,6 @@ from ..api.exceptions import ImageNotFound from ..container_registry import get_container_registry_cls -from ..models.association_container_registries_groups import ( - AssociationContainerRegistriesGroupsRow, -) from ..models.container_registry import ContainerRegistryRow from .base import ( GUID, @@ -829,12 +826,8 @@ async def _in_project_scope( ContainerRegistryRow.is_global == true(), sa.and_( ContainerRegistryRow.is_global == false(), - sa.exists().where( - (AssociationContainerRegistriesGroupsRow.group_id == scope.project_id) - & ( - AssociationContainerRegistriesGroupsRow.registry_id - == ImageRow.registry_id - ) + ContainerRegistryRow.association_container_registries_groups_rows.any( + group_id=scope.project_id ), ), )