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

pkg/services/authn/authnimpl/sync/oauth_token_sync.go

@@ -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