From 86f2bf294044129e802ddeab6f4ac8e9f79dee9e Mon Sep 17 00:00:00 2001 From: xavi <114113189+volcanonoodle@users.noreply.github.com> Date: Tue, 3 Jun 2025 08:59:40 +0200 Subject: [PATCH] IAM: Skip token rotation if it has been rotated recently (#106075) --- .../src/types/featureToggles.gen.ts | 5 ++ pkg/services/auth/authimpl/auth_token.go | 48 +++++++++++++++++-- pkg/services/auth/authimpl/auth_token_test.go | 38 ++++++++++++++- pkg/services/featuremgmt/registry.go | 9 ++++ pkg/services/featuremgmt/toggles_gen.csv | 1 + pkg/services/featuremgmt/toggles_gen.go | 4 ++ pkg/services/featuremgmt/toggles_gen.json | 18 +++++++ pkg/services/quota/quotaimpl/quota_test.go | 2 +- 8 files changed, 119 insertions(+), 6 deletions(-) diff --git a/packages/grafana-data/src/types/featureToggles.gen.ts b/packages/grafana-data/src/types/featureToggles.gen.ts index 90fb2564f8c..7feebb276e6 100644 --- a/packages/grafana-data/src/types/featureToggles.gen.ts +++ b/packages/grafana-data/src/types/featureToggles.gen.ts @@ -1017,4 +1017,9 @@ export interface FeatureToggles { * @default false */ restoreDashboards?: boolean; + /** + * Skip token rotation if it was already rotated less than 5 seconds ago + * @default false + */ + skipTokenRotationIfRecent?: boolean; } diff --git a/pkg/services/auth/authimpl/auth_token.go b/pkg/services/auth/authimpl/auth_token.go index 241c919033c..bf95502b11e 100644 --- a/pkg/services/auth/authimpl/auth_token.go +++ b/pkg/services/auth/authimpl/auth_token.go @@ -9,6 +9,8 @@ import ( "strings" "time" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/codes" "golang.org/x/sync/singleflight" "github.com/grafana/grafana/pkg/infra/db" @@ -17,6 +19,7 @@ import ( "github.com/grafana/grafana/pkg/infra/tracing" "github.com/grafana/grafana/pkg/models/usertoken" "github.com/grafana/grafana/pkg/services/auth" + "github.com/grafana/grafana/pkg/services/featuremgmt" "github.com/grafana/grafana/pkg/services/quota" "github.com/grafana/grafana/pkg/services/secrets" "github.com/grafana/grafana/pkg/setting" @@ -29,12 +32,14 @@ var ( errUserIDInvalid = errors.New("invalid user ID") ) +const SkipRotationTime = 5 * time.Second + var _ auth.UserTokenService = (*UserAuthTokenService)(nil) func ProvideUserAuthTokenService(sqlStore db.DB, serverLockService *serverlock.ServerLockService, quotaService quota.Service, secretService secrets.Service, - cfg *setting.Cfg, tracer tracing.Tracer, + cfg *setting.Cfg, tracer tracing.Tracer, features featuremgmt.FeatureToggles, ) (*UserAuthTokenService, error) { s := &UserAuthTokenService{ sqlStore: sqlStore, @@ -42,6 +47,8 @@ func ProvideUserAuthTokenService(sqlStore db.DB, cfg: cfg, log: log.New("auth"), singleflight: new(singleflight.Group), + features: features, + tracer: tracer, } s.externalSessionStore = provideExternalSessionStore(sqlStore, secretService, tracer) @@ -68,6 +75,8 @@ type UserAuthTokenService struct { log log.Logger externalSessionStore auth.ExternalSessionStore singleflight *singleflight.Group + features featuremgmt.FeatureToggles + tracer tracing.Tracer } func (s *UserAuthTokenService) CreateToken(ctx context.Context, cmd *auth.CreateTokenCommand) (*auth.UserToken, error) { @@ -261,28 +270,55 @@ func (s *UserAuthTokenService) UpdateExternalSession(ctx context.Context, extern } func (s *UserAuthTokenService) RotateToken(ctx context.Context, cmd auth.RotateCommand) (*auth.UserToken, error) { + ctx, span := s.tracer.Start(ctx, "authtoken.RotateToken") + defer span.End() + if cmd.UnHashedToken == "" { return nil, auth.ErrInvalidSessionToken } - res, err, _ := s.singleflight.Do(cmd.UnHashedToken, func() (any, error) { + rotate := func(ctx context.Context) (*auth.UserToken, error) { token, err := s.LookupToken(ctx, cmd.UnHashedToken) if err != nil { return nil, err } - s.log.FromContext(ctx).Debug("Rotating token", "tokenID", token.Id, "userID", token.UserId, "createdAt", token.CreatedAt, "rotatedAt", token.RotatedAt) + log := s.log.FromContext(ctx).New("tokenID", token.Id, "userID", token.UserId, "createdAt", token.CreatedAt, "rotatedAt", token.RotatedAt) + + // Avoid multiple instances in HA mode rotating at the same time. + if s.features.IsEnabled(ctx, featuremgmt.FlagSkipTokenRotationIfRecent) && time.Unix(token.RotatedAt, 0).Add(SkipRotationTime).After(getTime()) { + log.Debug("Token was last rotated very recently, skipping rotation") + span.SetAttributes(attribute.Bool("skipped", true)) + return token, nil + } + log.Debug("Rotating token") newToken, err := s.rotateToken(ctx, token, cmd.IP, cmd.UserAgent) if errors.Is(err, errTokenNotRotated) { + span.SetAttributes(attribute.Bool("rotated", false)) return token, nil } if err != nil { + span.SetStatus(codes.Error, "token rotation failed") + span.RecordError(err) return nil, err } return newToken, nil + } + + res, err, _ := s.singleflight.Do(cmd.UnHashedToken, func() (any, error) { + if s.features.IsEnabled(ctx, featuremgmt.FlagSkipTokenRotationIfRecent) { + var token *auth.UserToken + err := s.sqlStore.InTransaction(ctx, func(ctx context.Context) error { + var err error + token, err = rotate(ctx) + return err + }) + return token, err + } + return rotate(ctx) }) if err != nil { @@ -318,7 +354,11 @@ func (s *UserAuthTokenService) rotateToken(ctx context.Context, token *auth.User now := getTime() var affected int64 - err = s.sqlStore.WithTransactionalDbSession(ctx, func(dbSession *db.Session) error { + withDbSession := s.sqlStore.WithDbSession + if !s.features.IsEnabled(ctx, featuremgmt.FlagSkipTokenRotationIfRecent) { + withDbSession = s.sqlStore.WithTransactionalDbSession + } + err = withDbSession(ctx, func(dbSession *db.Session) error { res, err := dbSession.Exec(sql, userAgent, clientIPStr, hashedToken, s.sqlStore.GetDialect().BooleanValue(false), now.Unix(), token.Id) if err != nil { return err diff --git a/pkg/services/auth/authimpl/auth_token_test.go b/pkg/services/auth/authimpl/auth_token_test.go index 86d579c49c6..df585feeeb5 100644 --- a/pkg/services/auth/authimpl/auth_token_test.go +++ b/pkg/services/auth/authimpl/auth_token_test.go @@ -20,6 +20,7 @@ import ( "github.com/grafana/grafana/pkg/infra/tracing" "github.com/grafana/grafana/pkg/services/auth" "github.com/grafana/grafana/pkg/services/auth/authtest" + "github.com/grafana/grafana/pkg/services/featuremgmt" "github.com/grafana/grafana/pkg/services/quota" "github.com/grafana/grafana/pkg/services/secrets/fakes" "github.com/grafana/grafana/pkg/services/user" @@ -491,6 +492,13 @@ func TestIntegrationUserAuthToken(t *testing.T) { }) t.Run("RotateToken", func(t *testing.T) { + advanceTime := func(d time.Duration) { + currentTime := getTime() + getTime = func() time.Time { + return currentTime.Add(d) + } + } + var prev string token, err := ctx.tokenService.CreateToken(context.Background(), &auth.CreateTokenCommand{ User: usr, @@ -499,6 +507,7 @@ func TestIntegrationUserAuthToken(t *testing.T) { }) require.NoError(t, err) t.Run("should rotate token when called with current auth token", func(t *testing.T) { + advanceTime(SkipRotationTime + 1*time.Second) prev = token.UnhashedToken token, err = ctx.tokenService.RotateToken(context.Background(), auth.RotateCommand{UnHashedToken: token.UnhashedToken}) require.NoError(t, err) @@ -507,6 +516,7 @@ func TestIntegrationUserAuthToken(t *testing.T) { }) t.Run("should rotate token when called with previous", func(t *testing.T) { + advanceTime(SkipRotationTime + 1*time.Second) newPrev := token.UnhashedToken token, err = ctx.tokenService.RotateToken(context.Background(), auth.RotateCommand{UnHashedToken: prev}) require.NoError(t, err) @@ -514,10 +524,33 @@ func TestIntegrationUserAuthToken(t *testing.T) { }) t.Run("should not rotate token when called with old previous", func(t *testing.T) { + advanceTime(SkipRotationTime + 1*time.Second) _, err = ctx.tokenService.RotateToken(context.Background(), auth.RotateCommand{UnHashedToken: prev}) require.ErrorIs(t, err, auth.ErrUserTokenNotFound) }) + t.Run("should not rotate token when last rotation happened recently", func(t *testing.T) { + advanceTime(SkipRotationTime + 1*time.Second) + prevToken, err := ctx.tokenService.CreateToken(context.Background(), &auth.CreateTokenCommand{ + User: usr, + ClientIP: nil, + UserAgent: "", + }) + require.NoError(t, err) + + advanceTime(SkipRotationTime + 1*time.Second) + rotatedToken, err := ctx.tokenService.RotateToken(context.Background(), auth.RotateCommand{UnHashedToken: prevToken.UnhashedToken}) + require.NoError(t, err) + assert.True(t, rotatedToken.UnhashedToken != prevToken.UnhashedToken) + assert.True(t, rotatedToken.PrevAuthToken == hashToken("", prevToken.UnhashedToken)) + + // Should not rotate because it already rotated less than 5s ago + skippedToken, err := ctx.tokenService.RotateToken(context.Background(), auth.RotateCommand{UnHashedToken: rotatedToken.UnhashedToken}) + require.NoError(t, err) + assert.True(t, skippedToken.UnhashedToken == rotatedToken.UnhashedToken) + assert.True(t, skippedToken.PrevAuthToken == hashToken("", prevToken.UnhashedToken)) + }) + t.Run("should return error when token is revoked", func(t *testing.T) { revokedToken, err := ctx.tokenService.CreateToken(context.Background(), &auth.CreateTokenCommand{ User: usr, @@ -669,6 +702,7 @@ func createTestContext(t *testing.T) *testContext { maxInactiveDurationVal, _ := time.ParseDuration("168h") maxLifetimeDurationVal, _ := time.ParseDuration("720h") sqlstore := db.InitTestDB(t) + tracer := tracing.InitializeTracerForTest() cfg := &setting.Cfg{ LoginMaxInactiveLifetime: maxInactiveDurationVal, @@ -676,7 +710,7 @@ func createTestContext(t *testing.T) *testContext { TokenRotationIntervalMinutes: 10, } - extSessionStore := provideExternalSessionStore(sqlstore, &fakes.FakeSecretsService{}, tracing.InitializeTracerForTest()) + extSessionStore := provideExternalSessionStore(sqlstore, &fakes.FakeSecretsService{}, tracer) tokenService := &UserAuthTokenService{ sqlStore: sqlstore, @@ -684,6 +718,8 @@ func createTestContext(t *testing.T) *testContext { log: log.New("test-logger"), singleflight: new(singleflight.Group), externalSessionStore: extSessionStore, + features: featuremgmt.WithFeatures(featuremgmt.FlagSkipTokenRotationIfRecent), + tracer: tracer, } return &testContext{ diff --git a/pkg/services/featuremgmt/registry.go b/pkg/services/featuremgmt/registry.go index 8df14ece300..16702138f45 100644 --- a/pkg/services/featuremgmt/registry.go +++ b/pkg/services/featuremgmt/registry.go @@ -1750,6 +1750,15 @@ var ( HideFromAdminPage: true, Expression: "false", }, + { + Name: "skipTokenRotationIfRecent", + Description: "Skip token rotation if it was already rotated less than 5 seconds ago", + Stage: FeatureStagePrivatePreview, + Owner: identityAccessTeam, + HideFromAdminPage: true, + HideFromDocs: true, + Expression: "false", + }, } ) diff --git a/pkg/services/featuremgmt/toggles_gen.csv b/pkg/services/featuremgmt/toggles_gen.csv index b17fcd520a7..55743a7a27c 100644 --- a/pkg/services/featuremgmt/toggles_gen.csv +++ b/pkg/services/featuremgmt/toggles_gen.csv @@ -229,3 +229,4 @@ alertRuleUseFiredAtForStartsAt,experimental,@grafana/alerting-squad,false,false, alertingBulkActionsInUI,GA,@grafana/alerting-squad,false,false,true extensionsReadOnlyProxy,experimental,@grafana/plugins-platform-backend,false,false,true restoreDashboards,experimental,@grafana/grafana-frontend-platform,false,false,false +skipTokenRotationIfRecent,privatePreview,@grafana/identity-access-team,false,false,false diff --git a/pkg/services/featuremgmt/toggles_gen.go b/pkg/services/featuremgmt/toggles_gen.go index 461311710ba..53a6c1f5d63 100644 --- a/pkg/services/featuremgmt/toggles_gen.go +++ b/pkg/services/featuremgmt/toggles_gen.go @@ -926,4 +926,8 @@ const ( // FlagRestoreDashboards // Enables restore deleted dashboards feature FlagRestoreDashboards = "restoreDashboards" + + // FlagSkipTokenRotationIfRecent + // Skip token rotation if it was already rotated less than 5 seconds ago + FlagSkipTokenRotationIfRecent = "skipTokenRotationIfRecent" ) diff --git a/pkg/services/featuremgmt/toggles_gen.json b/pkg/services/featuremgmt/toggles_gen.json index 7f96297126d..3be2bce0f85 100644 --- a/pkg/services/featuremgmt/toggles_gen.json +++ b/pkg/services/featuremgmt/toggles_gen.json @@ -3020,6 +3020,24 @@ "codeowner": "@grafana/dashboards-squad" } }, + { + "metadata": { + "name": "skipTokenRotationIfRecent", + "resourceVersion": "1748363818965", + "creationTimestamp": "2025-05-27T16:30:53Z", + "annotations": { + "grafana.app/updatedTimestamp": "2025-05-27 16:36:58.965629 +0000 UTC" + } + }, + "spec": { + "description": "Skip token rotation if it was already rotated less than 5 seconds ago", + "stage": "privatePreview", + "codeowner": "@grafana/identity-access-team", + "hideFromAdminPage": true, + "hideFromDocs": true, + "expression": "false" + } + }, { "metadata": { "name": "sqlDatasourceDatabaseSelection", diff --git a/pkg/services/quota/quotaimpl/quota_test.go b/pkg/services/quota/quotaimpl/quota_test.go index 6fca4b44c0a..6cfab09e36d 100644 --- a/pkg/services/quota/quotaimpl/quota_test.go +++ b/pkg/services/quota/quotaimpl/quota_test.go @@ -494,7 +494,7 @@ func setupEnv(t *testing.T, sqlStore db.DB, cfg *setting.Cfg, b bus.Bus, quotaSe tracer := tracing.InitializeTracerForTest() _, err := apikeyimpl.ProvideService(sqlStore, cfg, quotaService) require.NoError(t, err) - _, err = authimpl.ProvideUserAuthTokenService(sqlStore, nil, quotaService, fakes.NewFakeSecretsService(), cfg, tracing.InitializeTracerForTest()) + _, err = authimpl.ProvideUserAuthTokenService(sqlStore, nil, quotaService, fakes.NewFakeSecretsService(), cfg, tracing.InitializeTracerForTest(), featuremgmt.WithFeatures()) require.NoError(t, err) folderStore := folderimpl.ProvideDashboardFolderStore(sqlStore) fStore := folderimpl.ProvideStore(sqlStore)