From 961211b21a2e002e5863343b3042cc45b530e7cb Mon Sep 17 00:00:00 2001 From: Gabriel MABILLE Date: Wed, 18 Dec 2024 14:07:19 +0100 Subject: [PATCH] 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 | 23 ++ pkg/services/authz/rbac/service.go | 131 +++++++++-- pkg/services/authz/rbac/service_test.go | 294 +++++++++++++++++++++++- pkg/services/authz/rbac/store/store.go | 18 +- 4 files changed, 435 insertions(+), 31 deletions(-) create mode 100644 pkg/services/authz/rbac/cache.go diff --git a/pkg/services/authz/rbac/cache.go b/pkg/services/authz/rbac/cache.go new file mode 100644 index 00000000000..58c46cb0a26 --- /dev/null +++ b/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) +} diff --git a/pkg/services/authz/rbac/service.go b/pkg/services/authz/rbac/service.go index 67862779f62..8ca38e2c4be 100644 --- a/pkg/services/authz/rbac/service.go +++ b/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,25 +24,42 @@ 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 { return &Service{ - store: store.NewStore(sql), - identityStore: identityStore, - actionMapper: mappers.NewK8sRbacMapper(), - logger: logger, - tracer: tracer, + store: store.NewStore(sql), + identityStore: identityStore, + 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}) - if err != nil { - return nil, fmt.Errorf("could not get basic roles: %w", err) + s.idCache.Set(uidCacheKey, userIdentifiers, 0) + s.idCache.Set(idCacheKey, userIdentifiers, 0) + + 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{ - UserID: userIdentifiers.ID, - Action: req.Action, - TeamIDs: teamIDs, - Role: basicRoles.Role, - IsServerAdmin: basicRoles.IsAdmin, - } - - return s.store.GetUserPermissions(ctx, req.Namespace, userPermQuery) + return teamIDs, nil } -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 diff --git a/pkg/services/authz/rbac/service_test.go b/pkg/services/authz/rbac/service_test.go index f7d865c5d69..043cb957fab 100644 --- a/pkg/services/authz/rbac/service_test.go +++ b/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 +} diff --git a/pkg/services/authz/rbac/store/store.go b/pkg/services/authz/rbac/store/store.go index db0a1e72b70..a81fab761bd 100644 --- a/pkg/services/authz/rbac/store/store.go +++ b/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 { - return &Store{ +func NewStore(sql legacysql.LegacyDatabaseProvider) *StoreImpl { + 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