From 085ae014cd80b728862e367661d848c6167156c5 Mon Sep 17 00:00:00 2001 From: Selene Date: Fri, 29 Jul 2022 16:30:46 +0200 Subject: [PATCH] Sync: Move ActiveTokenCount to a new service (#52991) * Move ActiveTokenCount to a new service * Fixing tests * fix `RootSystem` Co-authored-by: Artur Wierzbicki --- pkg/models/user_token.go | 5 +++- pkg/server/wire.go | 3 +++ pkg/services/auth/auth_token.go | 20 ++++++++++++---- pkg/services/auth/auth_token_test.go | 34 +++++++++++++++++---------- pkg/services/quota/quotaimpl/quota.go | 8 +++---- pkg/services/store/service.go | 2 +- 6 files changed, 50 insertions(+), 22 deletions(-) diff --git a/pkg/models/user_token.go b/pkg/models/user_token.go index f1220542fef..6ce74c004f3 100644 --- a/pkg/models/user_token.go +++ b/pkg/models/user_token.go @@ -71,12 +71,15 @@ type UserTokenService interface { TryRotateToken(ctx context.Context, token *UserToken, clientIP net.IP, userAgent string) (bool, error) RevokeToken(ctx context.Context, token *UserToken, soft bool) error RevokeAllUserTokens(ctx context.Context, userId int64) error - ActiveTokenCount(ctx context.Context) (int64, error) GetUserToken(ctx context.Context, userId, userTokenId int64) (*UserToken, error) GetUserTokens(ctx context.Context, userId int64) ([]*UserToken, error) GetUserRevokedTokens(ctx context.Context, userId int64) ([]*UserToken, error) } +type ActiveTokenService interface { + ActiveTokenCount(ctx context.Context) (int64, error) +} + type UserTokenBackgroundService interface { registry.BackgroundService } diff --git a/pkg/server/wire.go b/pkg/server/wire.go index 468d9cc866c..b811979a99b 100644 --- a/pkg/server/wire.go +++ b/pkg/server/wire.go @@ -6,6 +6,7 @@ package server import ( "github.com/google/wire" sdkhttpclient "github.com/grafana/grafana-plugin-sdk-go/backend/httpclient" + "github.com/grafana/grafana/pkg/services/auth" "github.com/grafana/grafana/pkg/services/playlist/playlistimpl" "github.com/grafana/grafana/pkg/services/store/sanitizer" @@ -222,6 +223,8 @@ var wireBasicSet = wire.NewSet( influxdb.ProvideService, wire.Bind(new(social.Service), new(*social.SocialService)), oauthtoken.ProvideService, + auth.ProvideActiveAuthTokenService, + wire.Bind(new(models.ActiveTokenService), new(*auth.ActiveAuthTokenService)), wire.Bind(new(oauthtoken.OAuthTokenService), new(*oauthtoken.Service)), tempo.ProvideService, loki.ProvideService, diff --git a/pkg/services/auth/auth_token.go b/pkg/services/auth/auth_token.go index b277b30d081..e95d0c842b7 100644 --- a/pkg/services/auth/auth_token.go +++ b/pkg/services/auth/auth_token.go @@ -42,14 +42,26 @@ type UserAuthTokenService struct { log log.Logger } -func (s *UserAuthTokenService) ActiveTokenCount(ctx context.Context) (int64, error) { +type ActiveAuthTokenService struct { + cfg *setting.Cfg + sqlStore sqlstore.Store +} + +func ProvideActiveAuthTokenService(cfg *setting.Cfg, sqlStore sqlstore.Store) *ActiveAuthTokenService { + return &ActiveAuthTokenService{ + cfg: cfg, + sqlStore: sqlStore, + } +} + +func (a *ActiveAuthTokenService) ActiveTokenCount(ctx context.Context) (int64, error) { var count int64 var err error - err = s.SQLStore.WithDbSession(ctx, func(dbSession *sqlstore.DBSession) error { + err = a.sqlStore.WithDbSession(ctx, func(dbSession *sqlstore.DBSession) error { var model userAuthToken count, err = dbSession.Where(`created_at > ? AND rotated_at > ? AND revoked_at = 0`, - s.createdAfterParam(), - s.rotatedAfterParam()). + getTime().Add(-a.cfg.LoginMaxLifetime).Unix(), + getTime().Add(-a.cfg.LoginMaxInactiveLifetime).Unix()). Count(&model) return err diff --git a/pkg/services/auth/auth_token_test.go b/pkg/services/auth/auth_token_test.go index 6a062647b2d..938aca66bae 100644 --- a/pkg/services/auth/auth_token_test.go +++ b/pkg/services/auth/auth_token_test.go @@ -41,7 +41,7 @@ func TestUserAuthToken(t *testing.T) { userToken := createToken() t.Run("Can count active tokens", func(t *testing.T) { - count, err := ctx.tokenService.ActiveTokenCount(context.Background()) + count, err := ctx.activeTokenService.ActiveTokenCount(context.Background()) require.Nil(t, err) require.Equal(t, int64(1), count) }) @@ -209,7 +209,7 @@ func TestUserAuthToken(t *testing.T) { require.Nil(t, notGood) t.Run("should not find active token when expired", func(t *testing.T) { - count, err := ctx.tokenService.ActiveTokenCount(context.Background()) + count, err := ctx.activeTokenService.ActiveTokenCount(context.Background()) require.Nil(t, err) require.Equal(t, int64(0), count) }) @@ -534,25 +534,35 @@ func createTestContext(t *testing.T) *testContext { maxInactiveDurationVal, _ := time.ParseDuration("168h") maxLifetimeDurationVal, _ := time.ParseDuration("720h") sqlstore := sqlstore.InitTestDB(t) + + cfg := &setting.Cfg{ + LoginMaxInactiveLifetime: maxInactiveDurationVal, + LoginMaxLifetime: maxLifetimeDurationVal, + TokenRotationIntervalMinutes: 10, + } + tokenService := &UserAuthTokenService{ SQLStore: sqlstore, - Cfg: &setting.Cfg{ - LoginMaxInactiveLifetime: maxInactiveDurationVal, - LoginMaxLifetime: maxLifetimeDurationVal, - TokenRotationIntervalMinutes: 10, - }, - log: log.New("test-logger"), + Cfg: cfg, + log: log.New("test-logger"), + } + + activeTokenService := &ActiveAuthTokenService{ + cfg: cfg, + sqlStore: sqlstore, } return &testContext{ - sqlstore: sqlstore, - tokenService: tokenService, + sqlstore: sqlstore, + tokenService: tokenService, + activeTokenService: activeTokenService, } } type testContext struct { - sqlstore *sqlstore.SQLStore - tokenService *UserAuthTokenService + sqlstore *sqlstore.SQLStore + tokenService *UserAuthTokenService + activeTokenService *ActiveAuthTokenService } func (c *testContext) getAuthTokenByID(id int64) (*userAuthToken, error) { diff --git a/pkg/services/quota/quotaimpl/quota.go b/pkg/services/quota/quotaimpl/quota.go index fc85fc6d97f..bb2a806a480 100644 --- a/pkg/services/quota/quotaimpl/quota.go +++ b/pkg/services/quota/quotaimpl/quota.go @@ -13,17 +13,17 @@ import ( type Service struct { store store - AuthTokenService models.UserTokenService + authTokenService models.ActiveTokenService Cfg *setting.Cfg SQLStore sqlstore.Store Logger log.Logger } -func ProvideService(db db.DB, cfg *setting.Cfg, tokenService models.UserTokenService, ss *sqlstore.SQLStore) quota.Service { +func ProvideService(db db.DB, cfg *setting.Cfg, tokenService models.ActiveTokenService, ss *sqlstore.SQLStore) quota.Service { return &Service{ store: &sqlStore{db: db}, Cfg: cfg, - AuthTokenService: tokenService, + authTokenService: tokenService, SQLStore: ss, Logger: log.New("quota_service"), } @@ -71,7 +71,7 @@ func (s *Service) CheckQuotaReached(ctx context.Context, target string, scopePar return true, nil } if target == "session" { - usedSessions, err := s.AuthTokenService.ActiveTokenCount(ctx) + usedSessions, err := s.authTokenService.ActiveTokenCount(ctx) if err != nil { return false, err } diff --git a/pkg/services/store/service.go b/pkg/services/store/service.go index 1e2661a8ca2..47cf8b233c7 100644 --- a/pkg/services/store/service.go +++ b/pkg/services/store/service.go @@ -171,7 +171,7 @@ func ProvideService( storages = append(storages, newSQLStorage(RootStorageMeta{ Builtin: true, - }, RootResources, + }, RootSystem, "System", "Grafana system storage", &StorageSQLConfig{}, sql, orgId))