From 3669df3f8ab0aec5f54f49dd9b0bd3ebd292f814 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=89mile=20R=C3=A9?= <emile@getprobo.com>
Date: Fri, 16 Jan 2026 16:09:04 +0400
Subject: [PATCH] Handle membership deletions and role updates in iam policies
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Émile Ré <emile@getprobo.com>
---
 .../settings/_components/MemberListItem.tsx   | 21 +++++--------
 pkg/coredata/membership.go                    | 19 +++++++++---
 pkg/iam/authorizer.go                         | 15 ++++++----
 pkg/iam/iam_actions.go                        |  3 ++
 pkg/iam/iam_policies.go                       | 30 ++++++++++++++++---
 pkg/server/api/connect/v1/v1_resolver.go      |  6 ++++
 6 files changed, 66 insertions(+), 28 deletions(-)

diff --git a/apps/console/src/pages/iam/organizations/settings/_components/MemberListItem.tsx b/apps/console/src/pages/iam/organizations/settings/_components/MemberListItem.tsx
index f6494e09d..1dc2b4568 100644
--- a/apps/console/src/pages/iam/organizations/settings/_components/MemberListItem.tsx
+++ b/apps/console/src/pages/iam/organizations/settings/_components/MemberListItem.tsx
@@ -10,7 +10,7 @@ import {
   useConfirm,
 } from "@probo/ui";
 import clsx from "clsx";
-import { use, useState } from "react";
+import { useState } from "react";
 import { useFragment } from "react-relay";
 import { graphql } from "relay-runtime";
 import type { MemberListItemFragment$key } from "/__generated__/iam/MemberListItemFragment.graphql";
@@ -18,7 +18,6 @@ import { useMutationWithToasts } from "/hooks/useMutationWithToasts";
 import { useOrganizationId } from "/hooks/useOrganizationId";
 import { sprintf } from "@probo/helpers";
 import { EditMemberDialog } from "./EditMemberDialog";
-import { CurrentUser } from "/providers/CurrentUser";
 
 const fragment = graphql`
   fragment MemberListItemFragment on Membership {
@@ -62,19 +61,15 @@ export function MemberListItem(props: {
   const [dialogOpen, setDialogOpen] = useState(false);
 
   const membership = useFragment<MemberListItemFragment$key>(fragment, fKey);
-  const { role } = use(CurrentUser);
 
   const isInactive = membership.state === "INACTIVE";
 
-  // Only OWNER can edit OWNER members
-  const canEditThisRole = membership.role === "OWNER" ? role === "OWNER" : true;
-
   const [removeMembership, isRemoving] = useMutationWithToasts(
     removeMemberMutation,
     {
       successMessage: __("Member removed successfully"),
       errorMessage: __("Failed to remove member"),
-    }
+    },
   );
 
   const handleRemove = async () => {
@@ -93,9 +88,9 @@ export function MemberListItem(props: {
       {
         message: sprintf(
           __("Are you sure you want to remove %s?"),
-          membership.profile.fullName
+          membership.profile.fullName,
         ),
-      }
+      },
     );
   };
 
@@ -104,7 +99,7 @@ export function MemberListItem(props: {
       <Tr
         className={clsx(
           isRemoving && "opacity-60 pointer-events-none",
-          isInactive && "opacity-50"
+          isInactive && "opacity-50",
         )}
       >
         <Td>
@@ -129,7 +124,7 @@ export function MemberListItem(props: {
               className="flex gap-2 justify-end"
               onClick={(e) => e.stopPropagation()}
             >
-              {membership.canUpdate && canEditThisRole && (
+              {membership.canUpdate && (
                 <Button
                   variant="secondary"
                   onClick={() => setDialogOpen(true)}
@@ -141,9 +136,7 @@ export function MemberListItem(props: {
               {isRemoving ? (
                 <Spinner size={16} />
               ) : (
-                membership.canDelete &&
-                canEditThisRole &&
-                membership.source !== "SCIM" && (
+                membership.canDelete && (
                   <Button
                     variant="danger"
                     onClick={handleRemove}
diff --git a/pkg/coredata/membership.go b/pkg/coredata/membership.go
index 27826c502..04d91e1f2 100644
--- a/pkg/coredata/membership.go
+++ b/pkg/coredata/membership.go
@@ -245,8 +245,10 @@ LEFT JOIN
 func (m *Membership) AuthorizationAttributes(ctx context.Context, conn pg.Conn) (map[string]string, error) {
 	q := `
 SELECT
-    identity_id
-    , organization_id
+    identity_id,
+    organization_id,
+    role,
+    source
 FROM
     iam_memberships
 WHERE
@@ -256,7 +258,14 @@ LIMIT 1;
 
 	var identityID gid.GID
 	var organizationID gid.GID
-	if err := conn.QueryRow(ctx, q, m.ID).Scan(&identityID, &organizationID); err != nil {
+	var role MembershipRole
+	var source MembershipSource
+	if err := conn.QueryRow(ctx, q, m.ID).Scan(
+		&identityID,
+		&organizationID,
+		&role,
+		&source,
+	); err != nil {
 		if errors.Is(err, pgx.ErrNoRows) {
 			return nil, ErrResourceNotFound
 		}
@@ -266,6 +275,8 @@ LIMIT 1;
 	return map[string]string{
 		"identity_id":     identityID.String(),
 		"organization_id": organizationID.String(),
+		"role":            role.String(),
+		"source":          source.String(),
 	}, nil
 }
 
@@ -449,7 +460,7 @@ JOIN
 LEFT JOIN
     iam_membership_profiles mp ON mp.membership_id = mbr.id
 WHERE
-	%s
+    %s
 `
 
 	query = fmt.Sprintf(query, scope.SQLFragment(), cursor.SQLFragment())
diff --git a/pkg/iam/authorizer.go b/pkg/iam/authorizer.go
index 3eb7b3890..625e9bafb 100644
--- a/pkg/iam/authorizer.go
+++ b/pkg/iam/authorizer.go
@@ -112,12 +112,15 @@ func (a *Authorizer) authorize(ctx context.Context, conn pg.Conn, params Authori
 	}
 
 	// Only set principal.organization_id if they have a role in this org
-	var principalOrgID string
+	var scopedPrincipalAttrs map[string]string
 	if membership != nil && role != "" {
-		principalOrgID = membership.OrganizationID.String()
+		scopedPrincipalAttrs = map[string]string{
+			"organization_id": membership.OrganizationID.String(),
+			"role":            membership.Role.String(),
+		}
 	}
 
-	principalAttrs, err := a.buildPrincipalAttributes(ctx, conn, params.Principal, principalOrgID)
+	principalAttrs, err := a.buildPrincipalAttributes(ctx, conn, params.Principal, scopedPrincipalAttrs)
 	if err != nil {
 		return fmt.Errorf("cannot build principal attributes: %w", err)
 	}
@@ -176,12 +179,12 @@ func (a *Authorizer) buildPrincipalAttributes(
 	ctx context.Context,
 	conn pg.Conn,
 	principalID gid.GID,
-	organizationID string,
+	defaultAttrs map[string]string,
 ) (map[string]string, error) {
 	attrs := map[string]string{
-		"id":              principalID.String(),
-		"organization_id": organizationID,
+		"id": principalID.String(),
 	}
+	maps.Copy(attrs, defaultAttrs)
 
 	if entity, ok := coredata.NewEntityFromID(principalID); ok {
 		if attributer, ok := entity.(AuthorizationAttributer); ok {
diff --git a/pkg/iam/iam_actions.go b/pkg/iam/iam_actions.go
index 8d48b62f2..3bd292567 100644
--- a/pkg/iam/iam_actions.go
+++ b/pkg/iam/iam_actions.go
@@ -48,6 +48,9 @@ const (
 	ActionMembershipUpdate = "iam:membership:update"
 	ActionMembershipDelete = "iam:membership:delete"
 
+	// Membership role actions
+	ActionMembershipRoleSetOwner = "iam:membership-role:set-owner"
+
 	// Membership Profile actions
 	ActionMembershipProfileGet = "iam:membership-profile:get"
 
diff --git a/pkg/iam/iam_policies.go b/pkg/iam/iam_policies.go
index b211bb265..6fb3b080a 100644
--- a/pkg/iam/iam_policies.go
+++ b/pkg/iam/iam_policies.go
@@ -123,9 +123,23 @@ var IAMOwnerPolicy = policy.NewPolicy(
 		WithSID("full-org-access").
 		When(policy.Equals("principal.organization_id", "resource.organization_id")),
 
-	// Full access to member management (scoped to own organization)
-	policy.Allow("iam:membership:*").
-		WithSID("full-membership-access").
+	// Full access to member management (scoped to own organization), except deletion on SCIM sourced memberships
+	policy.Allow(
+		ActionMembershipGet,
+		ActionMembershipList,
+		ActionMembershipUpdate,
+	).
+		WithSID("membership-owner-access").
+		When(policy.Equals("principal.organization_id", "resource.organization_id")),
+	policy.Allow(
+		ActionMembershipDelete,
+	).
+		WithSID("membership-deletion-owner-access").
+		When(policy.NotEquals("resource.source", "SCIM")),
+
+	// Can set other members OWNER
+	policy.Allow(ActionMembershipRoleSetOwner).
+		WithSID("membership-role-owner-access").
 		When(policy.Equals("principal.organization_id", "resource.organization_id")),
 
 	// Full access to membership profiles (scoped to own organization)
@@ -185,11 +199,19 @@ var IAMAdminPolicy = policy.NewPolicy(
 	// Can manage memberships (scoped to own organization)
 	policy.Allow(
 		ActionMembershipGet,
-		ActionMembershipUpdate,
 	).
 		WithSID("membership-admin-access").
 		When(policy.Equals("principal.organization_id", "resource.organization_id")),
 
+	policy.Allow(
+		ActionMembershipUpdate,
+	).
+		WithSID("membership-role-admin-access").
+		When(
+			policy.Equals("principal.organization_id", "resource.organization_id"),
+			policy.NotEquals("resource.role", "OWNER"),
+		),
+
 	// Can view membership profiles (scoped to own organization)
 	policy.Allow(ActionMembershipProfileGet).
 		WithSID("membership-profile-admin-access").
diff --git a/pkg/server/api/connect/v1/v1_resolver.go b/pkg/server/api/connect/v1/v1_resolver.go
index 9d02bac74..2248310d7 100644
--- a/pkg/server/api/connect/v1/v1_resolver.go
+++ b/pkg/server/api/connect/v1/v1_resolver.go
@@ -950,6 +950,12 @@ func (r *mutationResolver) UpdateMembership(ctx context.Context, input types.Upd
 		return nil, err
 	}
 
+	if input.Role == coredata.MembershipRoleOwner {
+		if err := r.authorize(ctx, input.MembershipID, iam.ActionMembershipRoleSetOwner); err != nil {
+			return nil, err
+		}
+	}
+
 	membership, err := r.iam.OrganizationService.UpdateMempership(ctx, input.OrganizationID, input.MembershipID, input.Role)
 	if err != nil {
 		r.logger.ErrorCtx(ctx, "cannot update membership", log.Error(err))