[GH-288] Add Support for WithSynchronousRefresh Option in CachingProvider for Blocking/Non-Blocking Key Refresh (#314)

Author: developerkunal

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 {

jwks/provider.go

@@ -102,32 +102,58 @@ 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 {
 	jwks      *jose.JSONWebKeySet
 	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()

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) {