From 84024355e278ae8128eafba5ef3557cda958499a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tim=20M=C3=B6hlmann?= <tim+github@zitadel.com>
Date: Thu, 25 Jan 2024 19:07:03 +0200
Subject: [PATCH] feat(op): split the access and ID token hint verifiers

In zitadel we require different behaviors wrt public key expiry between access tokens and ID token hints.
This change splits the two verifiers in the OP.
The default is still based on Storage and passed to both verifier fields.
---
 pkg/op/op.go | 55 +++++++++++++++++++++++++++++-----------------------
 1 file changed, 31 insertions(+), 24 deletions(-)

diff --git a/pkg/op/op.go b/pkg/op/op.go
index fdc073c..5c88f77 100644
--- a/pkg/op/op.go
+++ b/pkg/op/op.go
@@ -257,13 +257,16 @@ func NewForwardedOpenIDProvider(path string, config *Config, storage Storage, op
 // op.AuthCallbackURL(provider) which is probably /callback. On the redirect back
 // to the AuthCallbackURL, the request id should be passed as the "id" parameter.
 func NewProvider(config *Config, storage Storage, issuer func(insecure bool) (IssuerFromRequest, error), opOpts ...Option) (_ *Provider, err error) {
+	keySet := &openIDKeySet{storage}
 	o := &Provider{
-		config:    config,
-		storage:   storage,
-		endpoints: DefaultEndpoints,
-		timer:     make(<-chan time.Time),
-		corsOpts:  &defaultCORSOptions,
-		logger:    slog.Default(),
+		config:            config,
+		storage:           storage,
+		accessTokenKeySet: keySet,
+		idTokenHinKeySet:  keySet,
+		endpoints:         DefaultEndpoints,
+		timer:             make(<-chan time.Time),
+		corsOpts:          &defaultCORSOptions,
+		logger:            slog.Default(),
 	}
 
 	for _, optFunc := range opOpts {
@@ -276,19 +279,11 @@ func NewProvider(config *Config, storage Storage, issuer func(insecure bool) (Is
 	if err != nil {
 		return nil, err
 	}
-
 	o.Handler = CreateRouter(o, o.interceptors...)
-
 	o.decoder = schema.NewDecoder()
 	o.decoder.IgnoreUnknownKeys(true)
-
 	o.encoder = oidc.NewEncoder()
-
 	o.crypto = NewAESCrypto(config.CryptoKey)
-
-	// Avoid potential race conditions by calling these early
-	_ = o.openIDKeySet() // sets keySet
-
 	return o, nil
 }
 
@@ -299,7 +294,8 @@ type Provider struct {
 	insecure                bool
 	endpoints               *Endpoints
 	storage                 Storage
-	keySet                  *openIDKeySet
+	accessTokenKeySet       oidc.KeySet
+	idTokenHinKeySet        oidc.KeySet
 	crypto                  Crypto
 	decoder                 *schema.Decoder
 	encoder                 *schema.Encoder
@@ -435,7 +431,7 @@ func (o *Provider) Encoder() httphelper.Encoder {
 }
 
 func (o *Provider) IDTokenHintVerifier(ctx context.Context) *IDTokenHintVerifier {
-	return NewIDTokenHintVerifier(IssuerFromContext(ctx), o.openIDKeySet(), o.idTokenHintVerifierOpts...)
+	return NewIDTokenHintVerifier(IssuerFromContext(ctx), o.idTokenHinKeySet, o.idTokenHintVerifierOpts...)
 }
 
 func (o *Provider) JWTProfileVerifier(ctx context.Context) *JWTProfileVerifier {
@@ -443,14 +439,7 @@ func (o *Provider) JWTProfileVerifier(ctx context.Context) *JWTProfileVerifier {
 }
 
 func (o *Provider) AccessTokenVerifier(ctx context.Context) *AccessTokenVerifier {
-	return NewAccessTokenVerifier(IssuerFromContext(ctx), o.openIDKeySet(), o.accessTokenVerifierOpts...)
-}
-
-func (o *Provider) openIDKeySet() oidc.KeySet {
-	if o.keySet == nil {
-		o.keySet = &openIDKeySet{o.Storage()}
-	}
-	return o.keySet
+	return NewAccessTokenVerifier(IssuerFromContext(ctx), o.accessTokenKeySet, o.accessTokenVerifierOpts...)
 }
 
 func (o *Provider) Crypto() Crypto {
@@ -617,6 +606,15 @@ func WithHttpInterceptors(interceptors ...HttpInterceptor) Option {
 	}
 }
 
+// WithAccessTokenKeySet allows passing a KeySet with public keys for Access Token verification.
+// The default KeySet uses the [Storage] interface
+func WithAccessTokenKeySet(keySet oidc.KeySet) Option {
+	return func(o *Provider) error {
+		o.accessTokenKeySet = keySet
+		return nil
+	}
+}
+
 func WithAccessTokenVerifierOpts(opts ...AccessTokenVerifierOpt) Option {
 	return func(o *Provider) error {
 		o.accessTokenVerifierOpts = opts
@@ -624,6 +622,15 @@ func WithAccessTokenVerifierOpts(opts ...AccessTokenVerifierOpt) Option {
 	}
 }
 
+// WithIDTokenHintKeySet allows passing a KeySet with public keys for ID Token Hint verification.
+// The default KeySet uses the [Storage] interface.
+func WithIDTokenHintKeySet(keySet oidc.KeySet) Option {
+	return func(o *Provider) error {
+		o.idTokenHinKeySet = keySet
+		return nil
+	}
+}
+
 func WithIDTokenHintVerifierOpts(opts ...IDTokenHintVerifierOpt) Option {
 	return func(o *Provider) error {
 		o.idTokenHintVerifierOpts = opts