OAuth: Use the attached external session data in OAuthToken and OAuthTokenSync (#96655)

* wip

* wip + tests

* wip

* wip opt2

* Use authn.Identity struct's SessionToken

* Merge fixes

* Handle disabling the feature flag correctly

* Fix test

* Cleanup

* Remove HasOAuthEntry from the OAuthTokenService interface

* Remove unused function
This commit is contained in:
Misi
2024-11-27 11:06:39 +01:00
committed by GitHub
parent 6e2d3cae5e
commit 84b8296ffb
27 changed files with 1032 additions and 310 deletions
@@ -17,12 +17,15 @@ import (
"github.com/grafana/grafana/pkg/login/social"
"github.com/grafana/grafana/pkg/services/auth"
"github.com/grafana/grafana/pkg/services/authn"
"github.com/grafana/grafana/pkg/services/featuremgmt"
"github.com/grafana/grafana/pkg/services/oauthtoken"
)
const maxOAuthTokenCacheTTL = 5 * time.Minute
func ProvideOAuthTokenSync(service oauthtoken.OAuthTokenService, sessionService auth.UserTokenService, socialService social.Service, tracer tracing.Tracer) *OAuthTokenSync {
func ProvideOAuthTokenSync(service oauthtoken.OAuthTokenService, sessionService auth.UserTokenService, socialService social.Service, tracer tracing.Tracer,
features featuremgmt.FeatureToggles,
) *OAuthTokenSync {
return &OAuthTokenSync{
log.New("oauth_token.sync"),
service,
@@ -31,6 +34,7 @@ func ProvideOAuthTokenSync(service oauthtoken.OAuthTokenService, sessionService
new(singleflight.Group),
tracer,
localcache.New(maxOAuthTokenCacheTTL, 15*time.Minute),
features,
}
}
@@ -42,6 +46,7 @@ type OAuthTokenSync struct {
singleflightGroup *singleflight.Group
tracer tracing.Tracer
cache *localcache.CacheService
features featuremgmt.FeatureToggles
}
func (s *OAuthTokenSync) SyncOauthTokenHook(ctx context.Context, id *authn.Identity, _ *authn.Request) error {
@@ -72,6 +77,10 @@ func (s *OAuthTokenSync) SyncOauthTokenHook(ctx context.Context, id *authn.Ident
ctxLogger := s.log.FromContext(ctx).New("userID", userID)
cacheKey := fmt.Sprintf("token-check-%s", id.GetID())
if s.features.IsEnabledGlobally(featuremgmt.FlagImprovedExternalSessionHandling) {
cacheKey = fmt.Sprintf("token-check-%s-%d", id.GetID(), id.SessionToken.Id)
}
if _, ok := s.cache.Get(cacheKey); ok {
ctxLogger.Debug("Expiration check has been cached, no need to refresh")
return nil
@@ -83,7 +92,7 @@ func (s *OAuthTokenSync) SyncOauthTokenHook(ctx context.Context, id *authn.Ident
updateCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 15*time.Second)
defer cancel()
token, refreshErr := s.service.TryTokenRefresh(updateCtx, id)
token, refreshErr := s.service.TryTokenRefresh(updateCtx, id, id.SessionToken)
if refreshErr != nil {
if errors.Is(refreshErr, context.Canceled) {
return nil, nil
@@ -20,6 +20,9 @@ import (
"github.com/grafana/grafana/pkg/services/auth"
"github.com/grafana/grafana/pkg/services/auth/authtest"
"github.com/grafana/grafana/pkg/services/authn"
"github.com/grafana/grafana/pkg/services/contexthandler/ctxkey"
contextmodel "github.com/grafana/grafana/pkg/services/contexthandler/model"
"github.com/grafana/grafana/pkg/services/featuremgmt"
"github.com/grafana/grafana/pkg/services/login"
"github.com/grafana/grafana/pkg/services/oauthtoken/oauthtokentest"
)
@@ -85,7 +88,7 @@ func TestOAuthTokenSync_SyncOAuthTokenHook(t *testing.T) {
)
service := &oauthtokentest.MockOauthTokenService{
TryTokenRefreshFunc: func(ctx context.Context, usr identity.Requester) (*oauth2.Token, error) {
TryTokenRefreshFunc: func(ctx context.Context, usr identity.Requester, _ *auth.UserToken) (*oauth2.Token, error) {
tryRefreshCalled = true
return nil, tt.expectedTryRefreshErr
},
@@ -116,9 +119,13 @@ func TestOAuthTokenSync_SyncOAuthTokenHook(t *testing.T) {
singleflightGroup: new(singleflight.Group),
tracer: tracing.InitializeTracerForTest(),
cache: localcache.New(maxOAuthTokenCacheTTL, 15*time.Minute),
features: featuremgmt.WithFeatures(),
}
err := sync.SyncOauthTokenHook(context.Background(), tt.identity, nil)
ctx := context.Background()
reqCtx := context.WithValue(ctx, ctxkey.Key{}, &contextmodel.ReqContext{UserToken: nil})
err := sync.SyncOauthTokenHook(reqCtx, tt.identity, nil)
assert.ErrorIs(t, err, tt.expectedErr)
assert.Equal(t, tt.expectTryRefreshTokenCalled, tryRefreshCalled)
assert.Equal(t, tt.expectRevokeTokenCalled, revokeTokenCalled)
+22 -8
View File
@@ -12,6 +12,7 @@ import (
"github.com/grafana/grafana/pkg/infra/log"
"github.com/grafana/grafana/pkg/infra/tracing"
"github.com/grafana/grafana/pkg/services/authn"
"github.com/grafana/grafana/pkg/services/featuremgmt"
"github.com/grafana/grafana/pkg/services/login"
"github.com/grafana/grafana/pkg/services/org"
"github.com/grafana/grafana/pkg/services/quota"
@@ -51,7 +52,9 @@ var (
errSignupNotAllowed = errors.New("system administrator has disabled signup")
)
func ProvideUserSync(userService user.Service, userProtectionService login.UserProtectionService, authInfoService login.AuthInfoService, quotaService quota.Service, tracer tracing.Tracer) *UserSync {
func ProvideUserSync(userService user.Service, userProtectionService login.UserProtectionService, authInfoService login.AuthInfoService,
quotaService quota.Service, tracer tracing.Tracer, features featuremgmt.FeatureToggles,
) *UserSync {
return &UserSync{
userService: userService,
authInfoService: authInfoService,
@@ -59,6 +62,7 @@ func ProvideUserSync(userService user.Service, userProtectionService login.UserP
quotaService: quotaService,
log: log.New("user.sync"),
tracer: tracer,
features: features,
}
}
@@ -69,6 +73,7 @@ type UserSync struct {
quotaService quota.Service
log log.Logger
tracer tracing.Tracer
features featuremgmt.FeatureToggles
}
// SyncUserHook syncs a user with the database
@@ -223,21 +228,30 @@ func (s *UserSync) upsertAuthConnection(ctx context.Context, userID int64, ident
// This can happen when: using multiple auth client where the same user exists in several or
// changing to new auth client
if createConnection {
return s.authInfoService.SetAuthInfo(ctx, &login.SetAuthInfoCommand{
setAuthInfoCmd := &login.SetAuthInfoCommand{
UserId: userID,
AuthModule: identity.AuthenticatedBy,
AuthId: identity.AuthID,
OAuthToken: identity.OAuthToken,
})
}
if !s.features.IsEnabledGlobally(featuremgmt.FlagImprovedExternalSessionHandling) {
setAuthInfoCmd.OAuthToken = identity.OAuthToken
}
return s.authInfoService.SetAuthInfo(ctx, setAuthInfoCmd)
}
s.log.FromContext(ctx).Debug("Updating auth connection for user", "id", identity.ID)
return s.authInfoService.UpdateAuthInfo(ctx, &login.UpdateAuthInfoCommand{
updateAuthInfoCmd := &login.UpdateAuthInfoCommand{
UserId: userID,
AuthId: identity.AuthID,
AuthModule: identity.AuthenticatedBy,
OAuthToken: identity.OAuthToken,
})
}
if !s.features.IsEnabledGlobally(featuremgmt.FlagImprovedExternalSessionHandling) {
updateAuthInfoCmd.OAuthToken = identity.OAuthToken
}
s.log.FromContext(ctx).Debug("Updating auth connection for user", "id", identity.ID)
return s.authInfoService.UpdateAuthInfo(ctx, updateAuthInfoCmd)
}
func (s *UserSync) updateUserAttributes(ctx context.Context, usr *user.User, id *authn.Identity, userAuth *login.UserAuth) error {
@@ -10,6 +10,7 @@ import (
"github.com/grafana/grafana/pkg/infra/tracing"
"github.com/grafana/grafana/pkg/services/authn"
"github.com/grafana/grafana/pkg/services/featuremgmt"
"github.com/grafana/grafana/pkg/services/login"
"github.com/grafana/grafana/pkg/services/login/authinfoimpl"
"github.com/grafana/grafana/pkg/services/login/authinfotest"
@@ -45,7 +46,9 @@ func TestUserSync_SyncUserHook(t *testing.T) {
AuthModule: "oauth",
AuthId: "2032",
UserId: 1,
Id: 1}}
Id: 1,
},
}
userService := &usertest.FakeUserService{ExpectedUser: &user.User{
ID: 1,
@@ -434,7 +437,7 @@ func TestUserSync_SyncUserHook(t *testing.T) {
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
s := ProvideUserSync(tt.fields.userService, userProtection, tt.fields.authInfoService, tt.fields.quotaService, tracing.InitializeTracerForTest())
s := ProvideUserSync(tt.fields.userService, userProtection, tt.fields.authInfoService, tt.fields.quotaService, tracing.InitializeTracerForTest(), featuremgmt.WithFeatures())
err := s.SyncUserHook(tt.args.ctx, tt.args.id, nil)
if tt.wantErr {
require.Error(t, err)