AuthN: Add in-memory cache for oauth token refresh hook (#63569)

* OAuthSyncHook: Add in-memory cache so we don't have to perform the check on every request

Co-authored-by: Jo <joao.guerreiro@grafana.com>

pkg/services/authn/authnimpl/service.go

@@ -148,7 +148,7 @@ func ProvideService(
 	s.RegisterPostAuthHook(userSyncService.SyncLastSeenHook, 40)
 
 	if features.IsEnabled(featuremgmt.FlagAccessTokenExpirationCheck) {
-		s.RegisterPostAuthHook(sync.ProvideOauthTokenSync(oauthTokenService, sessionService).SyncOauthTokenHook, 60)
+		s.RegisterPostAuthHook(sync.ProvideOAuthTokenSync(oauthTokenService, sessionService).SyncOauthTokenHook, 60)
 	}
 
 	s.RegisterPostAuthHook(userSyncService.FetchSyncedUserHook, 100)

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

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"time"
 
+	"github.com/grafana/grafana/pkg/infra/localcache"
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/services/auth"
 	"github.com/grafana/grafana/pkg/services/authn"
@@ -17,21 +18,23 @@ var (
 	errExpiredAccessToken = errutil.NewBase(errutil.StatusUnauthorized, "oauth.expired-token")
 )
 
-func ProvideOauthTokenSync(service oauthtoken.OAuthTokenService, sessionService auth.UserTokenService) *OauthTokenSync {
+func ProvideOAuthTokenSync(service oauthtoken.OAuthTokenService, sessionService auth.UserTokenService) *OAuthTokenSync {
-	return &OauthTokenSync{
+	return &OAuthTokenSync{
 		log.New("oauth_token.sync"),
+		localcache.New(maxOAuthTokenCacheTTL, 15*time.Minute),
 		service,
 		sessionService,
 	}
 }
 
-type OauthTokenSync struct {
+type OAuthTokenSync struct {
 	log            log.Logger
+	cache          *localcache.CacheService
 	service        oauthtoken.OAuthTokenService
 	sessionService auth.UserTokenService
 }
 
-func (s *OauthTokenSync) SyncOauthTokenHook(ctx context.Context, identity *authn.Identity, _ *authn.Request) error {
+func (s *OAuthTokenSync) SyncOauthTokenHook(ctx context.Context, identity *authn.Identity, _ *authn.Request) error {
 	namespace, id := identity.NamespacedID()
 	// only perform oauth token check if identity is a user
 	if namespace != authn.NamespaceUser {
@@ -43,6 +46,11 @@ func (s *OauthTokenSync) SyncOauthTokenHook(ctx context.Context, identity *authn
 		return nil
 	}
 
+	// if we recently have performed this it would be cached, so we can skip the hook
+	if _, ok := s.cache.Get(identity.ID); ok {
+		return nil
+	}
+
 	token, exists, _ := s.service.HasOAuthEntry(ctx, &user.SignedInUser{UserID: id})
 	// user is not authenticated through oauth so skip further checks
 	if !exists {
@@ -51,11 +59,16 @@ func (s *OauthTokenSync) SyncOauthTokenHook(ctx context.Context, identity *authn
 
 	// token has no expire time configured, so we don't have to refresh it
 	if token.OAuthExpiry.IsZero() {
+		// cache the token check, so we don't perform it on every request
+		s.cache.Set(identity.ID, struct{}{}, getOAuthTokenCacheTTL(token.OAuthExpiry))
 		return nil
 	}
 
+	expires := token.OAuthExpiry.Round(0).Add(-oauthtoken.ExpiryDelta)
 	// token has not expired, so we don't have to refresh it
-	if !token.OAuthExpiry.Round(0).Add(-oauthtoken.ExpiryDelta).Before(time.Now()) {
+	if !expires.Before(time.Now()) {
+		// cache the token check, so we don't perform it on every request
+		s.cache.Set(identity.ID, struct{}{}, getOAuthTokenCacheTTL(expires))
 		return nil
 	}
 
@@ -77,3 +90,18 @@ func (s *OauthTokenSync) SyncOauthTokenHook(ctx context.Context, identity *authn
 
 	return nil
 }
+
+const maxOAuthTokenCacheTTL = 10 * time.Minute
+
+func getOAuthTokenCacheTTL(t time.Time) time.Duration {
+	if t.IsZero() {
+		return maxOAuthTokenCacheTTL
+	}
+
+	ttl := time.Until(t)
+	if ttl > maxOAuthTokenCacheTTL {
+		return maxOAuthTokenCacheTTL
+	}
+
+	return ttl
+}

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

@@ -8,6 +8,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 
+	"github.com/grafana/grafana/pkg/infra/localcache"
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/services/auth"
 	"github.com/grafana/grafana/pkg/services/auth/authtest"
@@ -17,7 +18,7 @@ import (
 	"github.com/grafana/grafana/pkg/services/user"
 )
 
-func TestOauthTokenSync_SyncOauthTokenHook(t *testing.T) {
+func TestOauthTokenSync_SyncOAuthTokenHook(t *testing.T) {
 	type testCase struct {
 		desc     string
 		identity *authn.Identity
@@ -117,8 +118,9 @@ func TestOauthTokenSync_SyncOauthTokenHook(t *testing.T) {
 				},
 			}
 
-			sync := &OauthTokenSync{
+			sync := &OAuthTokenSync{
 				log:            log.NewNopLogger(),
+				cache:          localcache.New(0, 0),
 				service:        service,
 				sessionService: sessionService,
 			}