Repository: auth0/go-jwt-middleware

Issue:
Race condition in JWK Caching Provider implementation?
### Checklist

- [X] I have looked into the [README](https://github.com/auth0/go-jwt-middleware#readme) and have not found a suitable solution or answer.
- [X] I have looked into the [documentation](https://pkg.go.dev/github.com/auth0/go-jwt-middleware/v2) and have not found a suitable solution or answer.
- [X] I have searched the [issues](https://github.com/auth0/go-jwt-middleware/issues) and have not found a suitable solution or answer.
- [X] I have upgraded to the [latest version](https://github.com/auth0/go-jwt-middleware/releases/latest) of this SDK and the issue still persists.
- [X] I have searched the [Auth0 Community](https://community.auth0.com/c/sdks/5) forums and have not found a suitable solution or answer.
- [X] I agree to the terms within the [Auth0 Code of Conduct](https://github.com/auth0/open-source-template/blob/master/CODE-OF-CONDUCT.md).

### Description

It seems like this library suffers a small race condition, when an IDP uses rolling keys. Under circumstances the validation of an JWT/JWK fails the first time.

### Reproduction

1. Use a JWKS Caching Provider with an OIDC conformant IDP
2. Validate a JWT with this library and with a valid JWK -> works fine
3. Wait until the IDP invalidates the current JWK 
4. Validate another JWT with the new JWK --> May fail with the error that the JWK type is not supported

### Go JWT Middleware version

2.2.1

### Go version

1.23

Discussion:
@jofleck: The commit history of the caching function says that refreshing in the "background" is done intentionally rather than blocking until the keys are refreshed. But this can lead to the misbehavior I mentioned before :) 

https://github.com/auth0/go-jwt-middleware/blob/f5f0a00e029f875cc2c23584955877f4cae13c57/jwks/provider.go#L136

Maybe we can let the developers decide if the key refreshing should be done blocking or non-blocking? 
@developerkunal: Hi @jofleck,

I hope you're doing great!

Sorry for the delayed response.

I’ve been working on this solution and have opened a PR for it. Would it be possible for you to test it out and check if there are any issues after the new changes? I’d really appreciate your feedback!

Thanks!@schlapzz: I have the same issue. Upgraded to v2.3.0 didn't fix it for me.
Our IDP is rotating the keys every hour. 
For example, if I use the caching provider with a cache TTL of 1 minute, as soon the keys rotate, the validation fails with the exception `go-jose/go-jose: unsupported key type/format\` until the cache expires.@developerkunal: Hi @schlapzz ,
have you used the `WithSynchronousRefresh` option after upgrade? and its still throwing error?

Solution patch:
diff --git a/error_handler.go b/error_handler.go
index 44d6c1d..816387f 100644
--- a/error_handler.go
+++ b/error_handler.go
@@ -28,7 +28,7 @@ type ErrorHandler func(w http.ResponseWriter, r *http.Request, err error)
 // DefaultErrorHandler is the default error handler implementation for the
 // JWTMiddleware. If an error handler is not provided via the WithErrorHandler
 // option this will be used.
-func DefaultErrorHandler(w http.ResponseWriter, r *http.Request, err error) {
+func DefaultErrorHandler(w http.ResponseWriter, _ *http.Request, err error) {
 	w.Header().Set("Content-Type", "application/json")
 
 	switch {
diff --git a/jwks/provider.go b/jwks/provider.go
index 808cae7..40d4f78 100644
--- a/jwks/provider.go
+++ b/jwks/provider.go
@@ -102,12 +102,16 @@ func (p *Provider) KeyFunc(ctx context.Context) (interface{}, error) {
 // in the background and the existing cached JWKS will be returned until the
 // JWKS cache is updated, or if the request errors then it will be evicted from
 // the cache.
+// The cache is keyed by the issuer's hostname. The synchronousRefresh
+// field determines whether the refresh is done synchronously or asynchronously.
+// This can be set using the WithSynchronousRefresh option.
 type CachingProvider struct {
 	*Provider
-	CacheTTL time.Duration
-	mu       sync.RWMutex
-	cache    map[string]cachedJWKS
-	sem      semaphore.Weighted
+	CacheTTL           time.Duration
+	mu                 sync.RWMutex
+	cache              map[string]cachedJWKS
+	sem                *semaphore.Weighted
+	synchronousRefresh bool
 }
 
 type cachedJWKS struct {
@@ -115,19 +119,41 @@ type cachedJWKS struct {
 	expiresAt time.Time
 }
 
+type CachingProviderOption func(*CachingProvider)
+
 // NewCachingProvider builds and returns a new CachingProvider.
 // If cacheTTL is zero then a default value of 1 minute will be used.
-func NewCachingProvider(issuerURL *url.URL, cacheTTL time.Duration, opts ...ProviderOption) *CachingProvider {
+func NewCachingProvider(issuerURL *url.URL, cacheTTL time.Duration, opts ...interface{}) *CachingProvider {
 	if cacheTTL == 0 {
 		cacheTTL = 1 * time.Minute
 	}
 
-	return &CachingProvider{
-		Provider: NewProvider(issuerURL, opts...),
-		CacheTTL: cacheTTL,
-		cache:    map[string]cachedJWKS{},
-		sem:      *semaphore.NewWeighted(1),
+	var providerOpts []ProviderOption
+	var cachingOpts []CachingProviderOption
+
+	for _, opt := range opts {
+		switch o := opt.(type) {
+		case ProviderOption:
+			providerOpts = append(providerOpts, o)
+		case CachingProviderOption:
+			cachingOpts = append(cachingOpts, o)
+		default:
+			panic(fmt.Sprintf("invalid option type: %T", o))
+		}
+	}
+	cp := &CachingProvider{
+		Provider:           NewProvider(issuerURL, providerOpts...),
+		CacheTTL:           cacheTTL,
+		cache:              map[string]cachedJWKS{},
+		sem:                semaphore.NewWeighted(1),
+		synchronousRefresh: false,
 	}
+
+	for _, opt := range cachingOpts {
+		opt(cp)
+	}
+
+	return cp
 }
 
 // KeyFunc adheres to the keyFunc signature that the Validator requires.
@@ -140,18 +166,26 @@ func (c *CachingProvider) KeyFunc(ctx context.Context) (interface{}, error) {
 
 	if cached, ok := c.cache[issuer]; ok {
 		if time.Now().After(cached.expiresAt) && c.sem.TryAcquire(1) {
-			go func() {
+			if !c.synchronousRefresh {
+				go func() {
+					defer c.sem.Release(1)
+					refreshCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+					defer cancel()
+					_, err := c.refreshKey(refreshCtx, issuer)
+
+					if err != nil {
+						c.mu.Lock()
+						delete(c.cache, issuer)
+						c.mu.Unlock()
+					}
+				}()
+				c.mu.RUnlock()
+				return cached.jwks, nil
+			} else {
+				c.mu.RUnlock()
 				defer c.sem.Release(1)
-				refreshCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
-				defer cancel()
-				_, err := c.refreshKey(refreshCtx, issuer)
-
-				if err != nil {
-					c.mu.Lock()
-					delete(c.cache, issuer)
-					c.mu.Unlock()
-				}
-			}()
+				return c.refreshKey(ctx, issuer)
+			}
 		}
 		c.mu.RUnlock()
 		return cached.jwks, nil
@@ -161,6 +195,15 @@ func (c *CachingProvider) KeyFunc(ctx context.Context) (interface{}, error) {
 	return c.refreshKey(ctx, issuer)
 }
 
+// WithSynchronousRefresh sets whether the CachingProvider blocks on refresh.
+// If set to true, it will block and wait for the refresh to complete.
+// If set to false (default), it will return the cached JWKS and trigger a background refresh.
+func WithSynchronousRefresh(blocking bool) CachingProviderOption {
+	return func(cp *CachingProvider) {
+		cp.synchronousRefresh = blocking
+	}
+}
+
 func (c *CachingProvider) refreshKey(ctx context.Context, issuer string) (interface{}, error) {
 	c.mu.Lock()
 	defer c.mu.Unlock()



Test patch:
diff --git a/jwks/provider_test.go b/jwks/provider_test.go
index 7682010..05484e7 100644
--- a/jwks/provider_test.go
+++ b/jwks/provider_test.go
@@ -240,6 +240,80 @@ func Test_JWKSProvider(t *testing.T) {
 			assert.Nil(t, cachedJWKS)
 		}, 1*time.Second, 250*time.Millisecond, "JWKS did not get uncached")
 	})
+	t.Run("It only calls the API once when multiple requests come in when using the CachingProvider with expired cache (WithSynchronousRefresh)", func(t *testing.T) {
+		initialJWKS, err := generateJWKS()
+		require.NoError(t, err)
+		atomic.StoreInt32(&requestCount, 0)
+
+		provider := NewCachingProvider(testServerURL, 5*time.Minute, WithSynchronousRefresh(true))
+		provider.cache[testServerURL.Hostname()] = cachedJWKS{
+			jwks:      initialJWKS,
+			expiresAt: time.Now(),
+		}
+
+		var wg sync.WaitGroup
+		for i := 0; i < 50; i++ {
+			wg.Add(1)
+			go func() {
+				_, _ = provider.KeyFunc(context.Background())
+				wg.Done()
+			}()
+		}
+		wg.Wait()
+		time.Sleep(2 * time.Second)
+		// No need for Eventually since we're not blocking on refresh.
+		returnedJWKS, err := provider.KeyFunc(context.Background())
+		require.NoError(t, err)
+		assert.True(t, cmp.Equal(expectedJWKS, returnedJWKS))
+
+		// Non-blocking behavior may allow extra API calls before the cache updates.
+		assert.Equal(t, int32(2), atomic.LoadInt32(&requestCount), "only wanted 2 requests (well known and jwks), but we got %d requests", atomic.LoadInt32(&requestCount))
+	})
+
+	t.Run("It only calls the API once when multiple requests come in when using the CachingProvider with no cache (WithSynchronousRefresh)", func(t *testing.T) {
+		provider := NewCachingProvider(testServerURL, 5*time.Minute, WithSynchronousRefresh(true))
+		atomic.StoreInt32(&requestCount, 0)
+
+		var wg sync.WaitGroup
+		for i := 0; i < 50; i++ {
+			wg.Add(1)
+			go func() {
+				_, _ = provider.KeyFunc(context.Background())
+				wg.Done()
+			}()
+		}
+		wg.Wait()
+
+		assert.Equal(t, int32(2), atomic.LoadInt32(&requestCount), "only wanted 2 requests (well known and jwks), but we got %d requests")
+	})
+	t.Run("It correctly applies both ProviderOptions and CachingProviderOptions when using the CachingProvider without breaking", func(t *testing.T) {
+		issuerURL, _ := url.Parse("https://example.com")
+		jwksURL, _ := url.Parse("https://example.com/jwks")
+		customClient := &http.Client{Timeout: 10 * time.Second}
+
+		provider := NewCachingProvider(
+			issuerURL,
+			30*time.Second,
+			WithCustomJWKSURI(jwksURL),
+			WithCustomClient(customClient),
+			WithSynchronousRefresh(true),
+		)
+
+		assert.Equal(t, jwksURL, provider.CustomJWKSURI, "CustomJWKSURI should be set correctly")
+		assert.Equal(t, customClient, provider.Client, "Custom HTTP client should be set correctly")
+		assert.True(t, provider.synchronousRefresh, "Synchronous refresh should be enabled")
+	})
+	t.Run("It panics when an invalid option type is provided when using the CachingProvider", func(t *testing.T) {
+		issuerURL, _ := url.Parse("https://example.com")
+
+		assert.Panics(t, func() {
+			NewCachingProvider(
+				issuerURL,
+				30*time.Second,
+				"invalid_option",
+			)
+		}, "Expected panic when passing an invalid option type")
+	})
 }
 
 func generateJWKS() (*jose.JSONWebKeySet, error) {