From 045733aed68ac45d629c1ffd30553b8dac63f24d Mon Sep 17 00:00:00 2001
From: xavi <114113189+volcanonoodle@users.noreply.github.com>
Date: Wed, 19 Mar 2025 10:06:58 +0100
Subject: [PATCH] [IAM] Clear user's permission cache after login (#102311)

---
 pkg/services/authn/authnimpl/registration.go  |  1 +
 .../authn/authnimpl/sync/rbac_sync.go         | 19 +++++++
 .../authn/authnimpl/sync/rbac_sync_test.go    | 56 +++++++++++++++++++
 3 files changed, 76 insertions(+)

diff --git a/pkg/services/authn/authnimpl/registration.go b/pkg/services/authn/authnimpl/registration.go
index c5af37a3098..bad45e88ef5 100644
--- a/pkg/services/authn/authnimpl/registration.go
+++ b/pkg/services/authn/authnimpl/registration.go
@@ -142,6 +142,7 @@ func ProvideRegistration(
 
 	authnSvc.RegisterPostAuthHook(rbacSync.SyncPermissionsHook, 120)
 	authnSvc.RegisterPostLoginHook(orgSync.SetDefaultOrgHook, 140)
+	authnSvc.RegisterPostLoginHook(rbacSync.ClearUserPermissionCacheHook, 170)
 
 	nsSync := sync.ProvideNamespaceSync(cfg)
 	authnSvc.RegisterPostAuthHook(nsSync.SyncNamespace, 150)
diff --git a/pkg/services/authn/authnimpl/sync/rbac_sync.go b/pkg/services/authn/authnimpl/sync/rbac_sync.go
index 8af08a38f51..75f6324382f 100644
--- a/pkg/services/authn/authnimpl/sync/rbac_sync.go
+++ b/pkg/services/authn/authnimpl/sync/rbac_sync.go
@@ -186,3 +186,22 @@ func (s *RBACSync) SyncCloudRoles(ctx context.Context, ident *authn.Identity, r
 		RolesToRemove: rolesToRemove,
 	})
 }
+
+// ClearUserPermissionCacheHook clears a user's permission cache if user Login succeeded. Necessary so that if a user logs in
+// through different SSO providers with different roles assigned in each, they do not get the wrong permissions.
+func (s *RBACSync) ClearUserPermissionCacheHook(ctx context.Context, ident *authn.Identity, r *authn.Request, err error) {
+	ctx, span := s.tracer.Start(ctx, "rbac.sync.ClearUserPermissionCacheHook")
+	defer span.End()
+
+	if err != nil {
+		return
+	}
+
+	ctxLogger := s.log.FromContext(ctx)
+	if !ident.IsIdentityType(claims.TypeUser) {
+		ctxLogger.Debug("Skipping user permission cache clear, not a user", "type", ident.GetIdentityType())
+		return
+	}
+
+	s.ac.ClearUserPermissionCache(ident)
+}
diff --git a/pkg/services/authn/authnimpl/sync/rbac_sync_test.go b/pkg/services/authn/authnimpl/sync/rbac_sync_test.go
index 594e9d80a1c..97c27942f22 100644
--- a/pkg/services/authn/authnimpl/sync/rbac_sync_test.go
+++ b/pkg/services/authn/authnimpl/sync/rbac_sync_test.go
@@ -2,6 +2,7 @@ package sync
 
 import (
 	"context"
+	"errors"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -365,6 +366,61 @@ func TestRBACSync_cloudRolesToAddAndRemove(t *testing.T) {
 	}
 }
 
+func TestRBACSync_ClearUserPermissionCacheHook(t *testing.T) {
+	type testCase struct {
+		desc           string
+		identityType   claims.IdentityType
+		loginErr       error
+		expectedCalled bool
+	}
+
+	tests := []testCase{
+		{
+			desc:           "should clear the permission cache when the user logged in successfully",
+			identityType:   claims.TypeUser,
+			loginErr:       nil,
+			expectedCalled: true,
+		},
+		{
+			desc:           "should skip clearing the permission cache when the user failed to log in",
+			identityType:   claims.TypeUser,
+			loginErr:       errors.New("failed to log in"),
+			expectedCalled: false,
+		},
+		{
+			desc:           "should skip clearing the permission cache when the identity is not a user",
+			identityType:   claims.TypeServiceAccount,
+			loginErr:       nil,
+			expectedCalled: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.desc, func(t *testing.T) {
+			var called bool
+			s := &RBACSync{
+				ac: &acmock.Mock{
+					ClearUserPermissionCacheFunc: func(_ identity.Requester) {
+						called = true
+					},
+				},
+				log:    log.NewNopLogger(),
+				tracer: tracing.InitializeTracerForTest(),
+			}
+			identity := &authn.Identity{
+				ID:       "1",
+				Type:     tt.identityType,
+				OrgID:    1,
+				OrgRoles: map[int64]org.RoleType{1: org.RoleViewer},
+			}
+			req := &authn.Request{}
+
+			s.ClearUserPermissionCacheHook(context.Background(), identity, req, tt.loginErr)
+			assert.Equal(t, tt.expectedCalled, called)
+		})
+	}
+}
+
 func setupTestEnv(t *testing.T) *RBACSync {
 	acMock := &acmock.Mock{
 		GetUserPermissionsFunc: func(ctx context.Context, siu identity.Requester, o accesscontrol.Options) ([]accesscontrol.Permission, error) {