From e20f8c566d8fdb251ff9539574700abce272bc32 Mon Sep 17 00:00:00 2001
From: Vardan Torosyan <vardants@gmail.com>
Date: Tue, 30 Jul 2024 09:00:47 +0200
Subject: [PATCH] RBAC sync: Fix removal of roles which need to be added
 (#91152)

* RBAC sync: Fix removal of roles which need to be added

* Optimize code

* cleanup: appease the linter

---------

Co-authored-by: Victor Cinaglia <victor@grafana.com>
---
 .../authn/authnimpl/sync/rbac_sync.go         | 29 ++++++++++---------
 .../authn/authnimpl/sync/rbac_sync_test.go    |  2 --
 2 files changed, 15 insertions(+), 16 deletions(-)

diff --git a/pkg/services/authn/authnimpl/sync/rbac_sync.go b/pkg/services/authn/authnimpl/sync/rbac_sync.go
index b1a3542af6d..95d7b0b44b5 100644
--- a/pkg/services/authn/authnimpl/sync/rbac_sync.go
+++ b/pkg/services/authn/authnimpl/sync/rbac_sync.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"errors"
 
+	"golang.org/x/exp/maps"
+
 	"github.com/grafana/grafana/pkg/apimachinery/errutil"
 	"github.com/grafana/grafana/pkg/apimachinery/identity"
 	"github.com/grafana/grafana/pkg/infra/log"
@@ -99,10 +101,6 @@ func (s *RBACSync) fetchPermissions(ctx context.Context, ident *authn.Identity)
 }
 
 func cloudRolesToAddAndRemove(ident *authn.Identity) ([]string, []string, error) {
-	const (
-		expectedRolesToAddCount = 2
-		rolesToRemoveInitialCap = 4
-	)
 	// Since Cloud Admin/Editor/Viewer roles are not yet implemented one-to-one in the Grafana, it becomes a confusing experience for users,
 	// therefore we are doing granular mapping of all available functionality in the Grafana temporary.
 	var fixedCloudRoles = map[org.RoleType][]string{
@@ -111,8 +109,8 @@ func cloudRolesToAddAndRemove(ident *authn.Identity) ([]string, []string, error)
 		org.RoleAdmin:  {accesscontrol.FixedCloudAdminRole, accesscontrol.FixedCloudSupportTicketAdmin},
 	}
 
-	rolesToAdd := make([]string, 0, expectedRolesToAddCount)
-	rolesToRemove := make([]string, 0, rolesToRemoveInitialCap)
+	rolesToAdd := make(map[string]bool)
+	rolesToRemove := make([]string, 0, 4)
 
 	currentRole := ident.GetOrgRole()
 	_, validRole := fixedCloudRoles[currentRole]
@@ -121,21 +119,24 @@ func cloudRolesToAddAndRemove(ident *authn.Identity) ([]string, []string, error)
 		return nil, nil, errInvalidCloudRole.Errorf("invalid role: %s", currentRole)
 	}
 
+	// Add roles for the current role and track them
+	for _, fixedRole := range fixedCloudRoles[currentRole] {
+		rolesToAdd[fixedRole] = true
+	}
+
+	// Add roles to remove, ensuring we don't remove any that have been added
 	for role, fixedRoles := range fixedCloudRoles {
+		if role == currentRole {
+			continue
+		}
 		for _, fixedRole := range fixedRoles {
-			if role == currentRole {
-				rolesToAdd = append(rolesToAdd, fixedRole)
-			} else {
+			if _, ok := rolesToAdd[fixedRole]; !ok {
 				rolesToRemove = append(rolesToRemove, fixedRole)
 			}
 		}
 	}
 
-	if len(rolesToAdd) != expectedRolesToAddCount {
-		return nil, nil, errInvalidCloudRole.Errorf("invalid role: %s", currentRole)
-	}
-
-	return rolesToAdd, rolesToRemove, nil
+	return maps.Keys(rolesToAdd), rolesToRemove, nil
 }
 
 func (s *RBACSync) SyncCloudRoles(ctx context.Context, ident *authn.Identity, r *authn.Request) error {
diff --git a/pkg/services/authn/authnimpl/sync/rbac_sync_test.go b/pkg/services/authn/authnimpl/sync/rbac_sync_test.go
index 4035565d8f9..1b0f54ef58e 100644
--- a/pkg/services/authn/authnimpl/sync/rbac_sync_test.go
+++ b/pkg/services/authn/authnimpl/sync/rbac_sync_test.go
@@ -189,7 +189,6 @@ func TestRBACSync_cloudRolesToAddAndRemove(t *testing.T) {
 				accesscontrol.FixedCloudViewerRole,
 				accesscontrol.FixedCloudSupportTicketReader,
 				accesscontrol.FixedCloudAdminRole,
-				accesscontrol.FixedCloudSupportTicketAdmin,
 			},
 		},
 		{
@@ -208,7 +207,6 @@ func TestRBACSync_cloudRolesToAddAndRemove(t *testing.T) {
 				accesscontrol.FixedCloudViewerRole,
 				accesscontrol.FixedCloudSupportTicketReader,
 				accesscontrol.FixedCloudEditorRole,
-				accesscontrol.FixedCloudSupportTicketAdmin,
 			},
 		},
 		{