IAM: Skip token rotation if it has been rotated recently (#106075)

This commit is contained in:
xavi
2025-06-03 08:59:40 +02:00
committed by GitHub
parent 439b8c01b3
commit 86f2bf2940
8 changed files with 119 additions and 6 deletions
+44 -4
View File
@@ -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
+37 -1
View File
@@ -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{