AuthZ Service: Add caching (#98008)

* AuthZ Service: Add caching

* split in functions

* Test getUserTeams

* Add tests to getUserBasicRole

* Test getUserPermissions

* Cache user identifiers

* fix test

pkg/services/authz/rbac/cache.go

@@ -0,0 +1,23 @@
+package rbac
+
+import "fmt"
+
+func userIdentifierCacheKey(namespace, userUID string) string {
+	return fmt.Sprintf("UID_%s_%s", namespace, userUID)
+}
+
+func userIdentifierCacheKeyById(namespace, ID string) string {
+	return fmt.Sprintf("ID_%s_%s", namespace, ID)
+}
+
+func userPermCacheKey(namespace, userUID, action string) string {
+	return fmt.Sprintf("%s_%s_%s", namespace, userUID, action)
+}
+
+func userBasicRoleCacheKey(namespace, userUID string) string {
+	return fmt.Sprintf("%s_%s", namespace, userUID)
+}
+
+func userTeamCacheKey(namespace, userUID string) string {
+	return fmt.Sprintf("%s_%s", namespace, userUID)
+}

pkg/services/authz/rbac/service.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"strconv"
+	"time"
 
 	authzv1 "github.com/grafana/authlib/authz/proto/v1"
 	"github.com/grafana/authlib/claims"
@@ -11,6 +12,7 @@ import (
 	"google.golang.org/grpc/status"
 	"k8s.io/apiserver/pkg/endpoints/request"
 
+	"github.com/grafana/grafana/pkg/infra/localcache"
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/infra/tracing"
 	"github.com/grafana/grafana/pkg/registry/apis/iam/common"
@@ -22,16 +24,29 @@ import (
 	"github.com/grafana/grafana/pkg/storage/legacysql"
 )
 
+const (
+	shortCacheTTL        = 1 * time.Minute
+	shortCleanupInterval = 5 * time.Minute
+	longCacheTTL         = 5 * time.Minute
+	longCleanupInterval  = 10 * time.Minute
+)
+
 type Service struct {
 	authzv1.UnimplementedAuthzServiceServer
 	authzextv1.UnimplementedAuthzExtentionServiceServer
 
-	store         *store.Store
+	store         store.Store
 	identityStore legacy.LegacyIdentityStore
 	actionMapper  *mappers.K8sRbacMapper
 
 	logger log.Logger
 	tracer tracing.Tracer
+
+	// Cache for user permissions, user team memberships and user basic roles
+	idCache        *localcache.CacheService
+	permCache      *localcache.CacheService
+	teamCache      *localcache.CacheService
+	basicRoleCache *localcache.CacheService
 }
 
 func NewService(sql legacysql.LegacyDatabaseProvider, identityStore legacy.LegacyIdentityStore, logger log.Logger, tracer tracing.Tracer) *Service {
@@ -41,6 +56,10 @@ func NewService(sql legacysql.LegacyDatabaseProvider, identityStore legacy.Legac
 		actionMapper:   mappers.NewK8sRbacMapper(),
 		logger:         logger,
 		tracer:         tracer,
+		idCache:        localcache.New(longCacheTTL, longCleanupInterval),
+		permCache:      localcache.New(shortCacheTTL, shortCleanupInterval),
+		teamCache:      localcache.New(shortCacheTTL, shortCleanupInterval),
+		basicRoleCache: localcache.New(longCacheTTL, longCleanupInterval),
 	}
 }
 
@@ -127,7 +146,55 @@ func (s *Service) validateRequest(ctx context.Context, req *authzv1.CheckRequest
 	return checkReq, nil
 }
 
-func (s *Service) getUserPermissions(ctx context.Context, req *CheckRequest) ([]accesscontrol.Permission, error) {
+func (s *Service) getUserPermissions(ctx context.Context, req *CheckRequest) (map[string]bool, error) {
+	userIdentifiers, err := s.GetUserIdentifiers(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+
+	userPermKey := userPermCacheKey(req.Namespace.Value, userIdentifiers.UID, req.Action)
+	if cached, ok := s.permCache.Get(userPermKey); ok {
+		return cached.(map[string]bool), nil
+	}
+
+	basicRoles, err := s.getUserBasicRole(ctx, req, userIdentifiers)
+	if err != nil {
+		return nil, err
+	}
+
+	teamIDs, err := s.getUserTeams(ctx, req, userIdentifiers)
+	if err != nil {
+		return nil, err
+	}
+
+	userPermQuery := store.PermissionsQuery{
+		UserID:        userIdentifiers.ID,
+		Action:        req.Action,
+		TeamIDs:       teamIDs,
+		Role:          basicRoles.Role,
+		IsServerAdmin: basicRoles.IsAdmin,
+	}
+
+	permissions, err := s.store.GetUserPermissions(ctx, req.Namespace, userPermQuery)
+	if err != nil {
+		return nil, err
+	}
+	scopeMap := getScopeMap(permissions)
+	s.permCache.Set(userPermKey, scopeMap, 0)
+	return scopeMap, nil
+}
+
+func (s *Service) GetUserIdentifiers(ctx context.Context, req *CheckRequest) (*store.UserIdentifiers, error) {
+	uidCacheKey := userIdentifierCacheKey(req.Namespace.Value, req.UserUID)
+	if cached, ok := s.idCache.Get(uidCacheKey); ok {
+		return cached.(*store.UserIdentifiers), nil
+	}
+
+	idCacheKey := userIdentifierCacheKeyById(req.Namespace.Value, req.UserUID)
+	if cached, ok := s.idCache.Get(idCacheKey); ok {
+		return cached.(*store.UserIdentifiers), nil
+	}
+
 	var userIDQuery store.UserIdentifierQuery
 	// Assume that numeric UID is user ID
 	if userID, err := strconv.Atoi(req.UserUID); err == nil {
@@ -140,12 +207,19 @@ func (s *Service) getUserPermissions(ctx context.Context, req *CheckRequest) ([]
 		return nil, fmt.Errorf("could not get user internal id: %w", err)
 	}
 
-	basicRoles, err := s.store.GetBasicRoles(ctx, req.Namespace, store.BasicRoleQuery{UserID: userIdentifiers.ID})
+	s.idCache.Set(uidCacheKey, userIdentifiers, 0)
-	if err != nil {
+	s.idCache.Set(idCacheKey, userIdentifiers, 0)
-		return nil, fmt.Errorf("could not get basic roles: %w", err)
+
+	return userIdentifiers, nil
+}
+
+func (s *Service) getUserTeams(ctx context.Context, req *CheckRequest, userIdentifiers *store.UserIdentifiers) ([]int64, error) {
+	teamIDs := make([]int64, 0, 50)
+	teamsCacheKey := userTeamCacheKey(req.Namespace.Value, userIdentifiers.UID)
+	if cached, ok := s.teamCache.Get(teamsCacheKey); ok {
+		return cached.([]int64), nil
 	}
 
-	teamIDs := make([]int64, 0, 50)
 	teamQuery := legacy.ListUserTeamsQuery{
 		UserUID:    userIdentifiers.UID,
 		Pagination: common.Pagination{Limit: 50},
@@ -164,28 +238,37 @@ func (s *Service) getUserPermissions(ctx context.Context, req *CheckRequest) ([]
 			break
 		}
 	}
+	s.teamCache.Set(teamsCacheKey, teamIDs, 0)
 
-	userPermQuery := store.PermissionsQuery{
+	return teamIDs, nil
-		UserID:        userIdentifiers.ID,
-		Action:        req.Action,
-		TeamIDs:       teamIDs,
-		Role:          basicRoles.Role,
-		IsServerAdmin: basicRoles.IsAdmin,
-	}
-
-	return s.store.GetUserPermissions(ctx, req.Namespace, userPermQuery)
 }
 
-func (s *Service) checkPermission(ctx context.Context, permissions []accesscontrol.Permission, req *CheckRequest) (bool, error) {
+func (s *Service) getUserBasicRole(ctx context.Context, req *CheckRequest, userIdentifiers *store.UserIdentifiers) (store.BasicRole, error) {
+	basicRoleKey := userBasicRoleCacheKey(req.Namespace.Value, userIdentifiers.UID)
+	if cached, ok := s.basicRoleCache.Get(basicRoleKey); ok {
+		return cached.(store.BasicRole), nil
+	}
+
+	basicRole, err := s.store.GetBasicRoles(ctx, req.Namespace, store.BasicRoleQuery{UserID: userIdentifiers.ID})
+	if err != nil {
+		return store.BasicRole{}, fmt.Errorf("could not get basic roles: %w", err)
+	}
+	if basicRole == nil {
+		basicRole = &store.BasicRole{}
+	}
+	s.basicRoleCache.Set(basicRoleKey, *basicRole, 0)
+
+	return *basicRole, nil
+}
+
+func (s *Service) checkPermission(ctx context.Context, scopeMap map[string]bool, req *CheckRequest) (bool, error) {
 	ctxLogger := s.logger.FromContext(ctx)
 
 	// Only check action if the request doesn't specify scope
 	if req.Name == "" {
-		return len(permissions) > 0, nil
+		return len(scopeMap) > 0, nil
 	}
 
-	scopeMap := getScopeMap(permissions)
-
 	// Wildcard grant, no further checks needed
 	if scopeMap["*"] {
 		return true, nil

pkg/services/authz/rbac/service_test.go

@@ -2,14 +2,20 @@ package rbac
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
+	"github.com/grafana/authlib/claims"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/grafana/grafana/pkg/infra/localcache"
 	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/infra/tracing"
+	"github.com/grafana/grafana/pkg/registry/apis/iam/legacy"
 	"github.com/grafana/grafana/pkg/services/accesscontrol"
 	"github.com/grafana/grafana/pkg/services/authz/mappers"
+	"github.com/grafana/grafana/pkg/services/authz/rbac/store"
 )
 
 func TestService_checkPermission(t *testing.T) {
@@ -143,9 +149,295 @@ func TestService_checkPermission(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			s := &Service{logger: log.New("test"), actionMapper: mappers.NewK8sRbacMapper()}
-			got, err := s.checkPermission(context.Background(), tc.permissions, &tc.check)
+			got, err := s.checkPermission(context.Background(), getScopeMap(tc.permissions), &tc.check)
 			require.NoError(t, err)
 			assert.Equal(t, tc.expected, got)
 		})
 	}
 }
+func TestService_getUserTeams(t *testing.T) {
+	type testCase struct {
+		name          string
+		teams         []int64
+		cacheHit      bool
+		expectedTeams []int64
+		expectedError bool
+	}
+
+	testCases := []testCase{
+		{
+			name:          "should return teams from cache if available",
+			teams:         []int64{1, 2},
+			cacheHit:      true,
+			expectedTeams: []int64{1, 2},
+			expectedError: false,
+		},
+		{
+			name:          "should return teams from identity store if not in cache",
+			teams:         []int64{3, 4},
+			cacheHit:      false,
+			expectedTeams: []int64{3, 4},
+			expectedError: false,
+		},
+		{
+			name:          "should return error if identity store fails",
+			teams:         []int64{},
+			cacheHit:      false,
+			expectedTeams: nil,
+			expectedError: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx := context.Background()
+			req := &CheckRequest{Namespace: claims.NamespaceInfo{Value: "stacks-12", OrgID: 1, StackID: 12}}
+
+			userIdentifiers := &store.UserIdentifiers{UID: "test-uid"}
+			identityStore := &fakeIdentityStore{teams: tc.teams, err: tc.expectedError}
+
+			cacheService := localcache.New(shortCacheTTL, shortCleanupInterval)
+			if tc.cacheHit {
+				cacheService.Set(userTeamCacheKey(req.Namespace.Value, userIdentifiers.UID), tc.expectedTeams, 0)
+			}
+
+			s := &Service{
+				teamCache:     cacheService,
+				identityStore: identityStore,
+				logger:        log.New("test"),
+			}
+
+			teams, err := s.getUserTeams(ctx, req, userIdentifiers)
+			if tc.expectedError {
+				require.Error(t, err)
+				return
+			}
+
+			require.NoError(t, err)
+			require.ElementsMatch(t, tc.expectedTeams, teams)
+			if tc.cacheHit {
+				require.Zero(t, identityStore.calls)
+			} else {
+				require.Equal(t, 1, identityStore.calls)
+			}
+		})
+	}
+}
+
+func TestService_getUserBasicRole(t *testing.T) {
+	type testCase struct {
+		name          string
+		basicRole     store.BasicRole
+		cacheHit      bool
+		expectedRole  store.BasicRole
+		expectedError bool
+	}
+
+	testCases := []testCase{
+		{
+			name: "should return basic role from cache if available",
+			basicRole: store.BasicRole{
+				Role:    "viewer",
+				IsAdmin: false,
+			},
+			cacheHit: true,
+			expectedRole: store.BasicRole{
+				Role:    "viewer",
+				IsAdmin: false,
+			},
+			expectedError: false,
+		},
+		{
+			name: "should return basic role from store if not in cache",
+			basicRole: store.BasicRole{
+				Role:    "editor",
+				IsAdmin: false,
+			},
+			cacheHit: false,
+			expectedRole: store.BasicRole{
+				Role:    "editor",
+				IsAdmin: false,
+			},
+			expectedError: false,
+		},
+		{
+			name:          "should return error if store fails",
+			basicRole:     store.BasicRole{},
+			cacheHit:      false,
+			expectedRole:  store.BasicRole{},
+			expectedError: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx := context.Background()
+			req := &CheckRequest{Namespace: claims.NamespaceInfo{Value: "stacks-12", OrgID: 1, StackID: 12}}
+
+			userIdentifiers := &store.UserIdentifiers{UID: "test-uid", ID: 1}
+			store := &fakeStore{basicRole: &tc.basicRole, err: tc.expectedError}
+
+			cacheService := localcache.New(shortCacheTTL, shortCleanupInterval)
+			if tc.cacheHit {
+				cacheService.Set(userBasicRoleCacheKey(req.Namespace.Value, userIdentifiers.UID), tc.expectedRole, 0)
+			}
+
+			s := &Service{
+				basicRoleCache: cacheService,
+				store:          store,
+				logger:         log.New("test"),
+			}
+
+			role, err := s.getUserBasicRole(ctx, req, userIdentifiers)
+			if tc.expectedError {
+				require.Error(t, err)
+				return
+			}
+
+			require.NoError(t, err)
+			require.Equal(t, tc.expectedRole, role)
+			if tc.cacheHit {
+				require.Zero(t, store.calls)
+			} else {
+				require.Equal(t, 1, store.calls)
+			}
+		})
+	}
+}
+
+func TestService_getUserPermissions(t *testing.T) {
+	type testCase struct {
+		name          string
+		permissions   []accesscontrol.Permission
+		cacheHit      bool
+		expectedPerms map[string]bool
+	}
+
+	testCases := []testCase{
+		{
+			name: "should return permissions from cache if available",
+			permissions: []accesscontrol.Permission{
+				{Action: "dashboards:read", Scope: "dashboards:uid:some_dashboard"},
+			},
+			cacheHit:      true,
+			expectedPerms: map[string]bool{"dashboards:uid:some_dashboard": true},
+		},
+		{
+			name: "should return permissions from store if not in cache",
+			permissions: []accesscontrol.Permission{
+				{Action: "dashboards:read", Scope: "dashboards:uid:some_dashboard"},
+			},
+			cacheHit:      false,
+			expectedPerms: map[string]bool{"dashboards:uid:some_dashboard": true},
+		},
+		{
+			name:          "should return error if store fails",
+			permissions:   nil,
+			cacheHit:      false,
+			expectedPerms: map[string]bool{},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx := context.Background()
+			userID := &store.UserIdentifiers{UID: "test-uid", ID: 112}
+			req := &CheckRequest{
+				Namespace: claims.NamespaceInfo{Value: "stacks-12", OrgID: 1, StackID: 12},
+				UserUID:   userID.UID,
+				Action:    "dashboards:read",
+			}
+
+			cacheService := localcache.New(shortCacheTTL, shortCleanupInterval)
+			if tc.cacheHit {
+				cacheService.Set(userPermCacheKey(req.Namespace.Value, userID.UID, req.Action), tc.expectedPerms, 0)
+			}
+
+			store := &fakeStore{
+				userID:          userID,
+				basicRole:       &store.BasicRole{Role: "viewer", IsAdmin: false},
+				userPermissions: tc.permissions,
+			}
+
+			s := &Service{
+				store:          store,
+				identityStore:  &fakeIdentityStore{teams: []int64{1, 2}},
+				actionMapper:   mappers.NewK8sRbacMapper(),
+				logger:         log.New("test"),
+				tracer:         tracing.NewNoopTracerService(),
+				idCache:        localcache.New(longCacheTTL, longCleanupInterval),
+				permCache:      cacheService,
+				basicRoleCache: localcache.New(longCacheTTL, longCleanupInterval),
+				teamCache:      localcache.New(shortCacheTTL, shortCleanupInterval),
+			}
+
+			perms, err := s.getUserPermissions(ctx, req)
+			require.NoError(t, err)
+			require.Len(t, perms, len(tc.expectedPerms))
+			for _, perm := range tc.permissions {
+				_, ok := tc.expectedPerms[perm.Scope]
+				require.True(t, ok)
+			}
+			if tc.cacheHit {
+				require.Equal(t, 1, store.calls) // only get user id
+			} else {
+				require.Equal(t, 3, store.calls) // get user id, basic role, and permissions
+			}
+		})
+	}
+}
+
+type fakeStore struct {
+	store.Store
+	basicRole       *store.BasicRole
+	userID          *store.UserIdentifiers
+	userPermissions []accesscontrol.Permission
+	err             bool
+	calls           int
+}
+
+func (f *fakeStore) GetBasicRoles(ctx context.Context, namespace claims.NamespaceInfo, query store.BasicRoleQuery) (*store.BasicRole, error) {
+	f.calls++
+	if f.err {
+		return nil, fmt.Errorf("store error")
+	}
+	return f.basicRole, nil
+}
+
+func (f *fakeStore) GetUserIdentifiers(ctx context.Context, query store.UserIdentifierQuery) (*store.UserIdentifiers, error) {
+	f.calls++
+	if f.err {
+		return nil, fmt.Errorf("store error")
+	}
+	return f.userID, nil
+}
+
+func (f *fakeStore) GetUserPermissions(ctx context.Context, namespace claims.NamespaceInfo, query store.PermissionsQuery) ([]accesscontrol.Permission, error) {
+	f.calls++
+	if f.err {
+		return nil, fmt.Errorf("store error")
+	}
+	return f.userPermissions, nil
+}
+
+type fakeIdentityStore struct {
+	legacy.LegacyIdentityStore
+	teams []int64
+	err   bool
+	calls int
+}
+
+func (f *fakeIdentityStore) ListUserTeams(ctx context.Context, namespace claims.NamespaceInfo, query legacy.ListUserTeamsQuery) (*legacy.ListUserTeamsResult, error) {
+	f.calls++
+	if f.err {
+		return nil, fmt.Errorf("identity store error")
+	}
+	items := make([]legacy.UserTeam, 0, len(f.teams))
+	for _, teamID := range f.teams {
+		items = append(items, legacy.UserTeam{ID: teamID})
+	}
+	return &legacy.ListUserTeamsResult{
+		Items:    items,
+		Continue: 0,
+	}, nil
+}

pkg/services/authz/rbac/store/store.go

@@ -11,17 +11,23 @@ import (
 	"github.com/grafana/grafana/pkg/storage/unified/sql/sqltemplate"
 )
 
-type Store struct {
+type Store interface {
+	GetUserPermissions(ctx context.Context, ns claims.NamespaceInfo, query PermissionsQuery) ([]accesscontrol.Permission, error)
+	GetUserIdentifiers(ctx context.Context, query UserIdentifierQuery) (*UserIdentifiers, error)
+	GetBasicRoles(ctx context.Context, ns claims.NamespaceInfo, query BasicRoleQuery) (*BasicRole, error)
+}
+
+type StoreImpl struct {
 	sql legacysql.LegacyDatabaseProvider
 }
 
-func NewStore(sql legacysql.LegacyDatabaseProvider) *Store {
+func NewStore(sql legacysql.LegacyDatabaseProvider) *StoreImpl {
-	return &Store{
+	return &StoreImpl{
 		sql: sql,
 	}
 }
 
-func (s *Store) GetUserPermissions(ctx context.Context, ns claims.NamespaceInfo, query PermissionsQuery) ([]accesscontrol.Permission, error) {
+func (s *StoreImpl) GetUserPermissions(ctx context.Context, ns claims.NamespaceInfo, query PermissionsQuery) ([]accesscontrol.Permission, error) {
 	sql, err := s.sql(ctx)
 	if err != nil {
 		return nil, err
@@ -56,7 +62,7 @@ func (s *Store) GetUserPermissions(ctx context.Context, ns claims.NamespaceInfo,
 	return perms, nil
 }
 
-func (s *Store) GetUserIdentifiers(ctx context.Context, query UserIdentifierQuery) (*UserIdentifiers, error) {
+func (s *StoreImpl) GetUserIdentifiers(ctx context.Context, query UserIdentifierQuery) (*UserIdentifiers, error) {
 	sql, err := s.sql(ctx)
 	if err != nil {
 		return nil, err
@@ -90,7 +96,7 @@ func (s *Store) GetUserIdentifiers(ctx context.Context, query UserIdentifierQuer
 	return &userIDs, nil
 }
 
-func (s *Store) GetBasicRoles(ctx context.Context, ns claims.NamespaceInfo, query BasicRoleQuery) (*BasicRole, error) {
+func (s *StoreImpl) GetBasicRoles(ctx context.Context, ns claims.NamespaceInfo, query BasicRoleQuery) (*BasicRole, error) {
 	sql, err := s.sql(ctx)
 	if err != nil {
 		return nil, err