Change for PR

Author: Thibaut Cholley

go.mod

@@ -1,4 +1,4 @@
-module github.com/Cafeine42/caddy-jwt
+module github.com/ggicci/caddy-jwt
 
 go 1.20
 

jwt.go

@@ -30,6 +30,23 @@ func init() {
 type User = caddyauth.User
 type Token = jwt.Token
 
+// jwkCacheEntry stores the JWK cache information for a specific URL
+type jwkCacheEntry struct {
+	URL       string
+	Cache     *jwk.Cache
+	CachedSet jwk.Set
+}
+
+// refresh refreshes the JWK cache for this entry
+func (entry *jwkCacheEntry) refresh(ctx context.Context, logger *zap.Logger) error {
+	_, err := entry.Cache.Refresh(ctx, entry.URL)
+	if err != nil {
+		logger.Warn("failed to refresh JWK cache", zap.Error(err), zap.String("url", entry.URL))
+		return err
+	}
+	return nil
+}
+
 // JWTAuth facilitates JWT (JSON Web Token) authentication.
 type JWTAuth struct {
 	// SignKey is the key used by the signing algorithm to verify the signature.
@@ -151,10 +168,7 @@ type JWTAuth struct {
 	parsedSignKey interface{} // can be []byte, *rsa.PublicKey, *ecdsa.PublicKey, etc.
 
 	// JWK cache by resolved URL to support placeholders in JWKURL
-	jwkCaches map[string]*struct {
-		cache     *jwk.Cache
-		cachedSet jwk.Set
-	}
+	jwkCaches map[string]*jwkCacheEntry
 }
 
 // CaddyModule implements caddy.Module interface.
@@ -182,19 +196,13 @@ func (ja *JWTAuth) usingJWK() bool {
 }
 
 func (ja *JWTAuth) setupJWKLoader() {
-	// Initialiser le cache pour toutes les URL possibles
-	ja.jwkCaches = make(map[string]*struct {
-		cache     *jwk.Cache
-		cachedSet jwk.Set
-	})
-	ja.logger.Info("JWK cache initialized for placeholder URL", zap.String("placeholder_url", ja.JWKURL))
+	// Initialize cache for all possible URLs
+	ja.jwkCaches = make(map[string]*jwkCacheEntry)
+	ja.logger.Info("JWK cache initialized for JWK URL", zap.String("jwk_url", ja.JWKURL))
 }
 
 // getOrCreateJWKCache retrieves or creates a cache for a specific JWK URL
-func (ja *JWTAuth) getOrCreateJWKCache(resolvedURL string) (*struct {
-	cache     *jwk.Cache
-	cachedSet jwk.Set
-}, error) {
+func (ja *JWTAuth) getOrCreateJWKCache(resolvedURL string) (*jwkCacheEntry, error) {
 	// If the URL is empty, return an error
 	if resolvedURL == "" {
 		return nil, fmt.Errorf("resolved JWK URL is empty")
@@ -212,44 +220,27 @@ func (ja *JWTAuth) getOrCreateJWKCache(resolvedURL string) (*struct {
 		return nil, fmt.Errorf("failed to register JWK URL: %w", err)
 	}
 
+	// Create cache entry before attempting refresh
+	entry := &jwkCacheEntry{
+		URL:       resolvedURL,
+		Cache:     cache,
+		CachedSet: jwk.NewCachedSet(cache, resolvedURL),
+	}
+
 	// Try to refresh the cache immediately
-	_, err = cache.Refresh(context.Background(), resolvedURL)
+	err = entry.refresh(context.Background(), ja.logger)
 	if err != nil {
 		ja.logger.Warn("failed to refresh JWK cache during initialization", zap.Error(err), zap.String("url", resolvedURL))
 		// Continue even in case of error, the important thing is that the URL is registered
 	}
 
-	cachedSet := jwk.NewCachedSet(cache, resolvedURL)
-	entry := &struct {
-		cache     *jwk.Cache
-		cachedSet jwk.Set
-	}{
-		cache:     cache,
-		cachedSet: cachedSet,
-	}
-
 	// Register the entry in the cache
 	ja.jwkCaches[resolvedURL] = entry
-	ja.logger.Info("new JWK cache created", zap.String("url", resolvedURL), zap.Int("loaded_keys", cachedSet.Len()))
+	ja.logger.Info("new JWK cache created", zap.String("url", resolvedURL), zap.Int("loaded_keys", entry.CachedSet.Len()))
 
 	return entry, nil
 }
 
-// refreshJWKCache refreshes the JWK cache for a specific URL. It validates the JWKs from the given URL.
-func (ja *JWTAuth) refreshJWKCache(resolvedURL string) error {
-	entry, err := ja.getOrCreateJWKCache(resolvedURL)
-	if err != nil {
-		return err
-	}
-
-	_, err = entry.cache.Refresh(context.Background(), resolvedURL)
-	if err != nil {
-		ja.logger.Warn("failed to refresh JWK cache", zap.Error(err), zap.String("url", resolvedURL))
-		return err
-	}
-	return nil
-}
-
 // Validate implements caddy.Validator interface.
 func (ja *JWTAuth) Validate() error {
 	if !ja.SkipVerification {
@@ -300,15 +291,18 @@ func (ja *JWTAuth) validateSignatureKeys() error {
 	return nil
 }
 
+// resolveJWKURL prend une requête HTTP et résout l'URL JWK avec des espaces réservés
+func (ja *JWTAuth) resolveJWKURL(request *http.Request) string {
+	replacer := request.Context().Value(caddy.ReplacerCtxKey).(*caddy.Replacer)
+	return replacer.ReplaceAll(ja.JWKURL, "")
+}
+
 func (ja *JWTAuth) keyProvider(request *http.Request) jws.KeyProviderFunc {
-	return func(context context.Context, sink jws.KeySink, sig *jws.Signature, _ *jws.Message) error {
+	return func(curContext context.Context, sink jws.KeySink, sig *jws.Signature, _ *jws.Message) error {
 		if ja.usingJWK() {
-			// Resolve JWKURL with placeholders
-			replacer := request.Context().Value(caddy.ReplacerCtxKey).(*caddy.Replacer)
-			resolvedURL := replacer.ReplaceAll(ja.JWKURL, "")
+			resolvedURL := ja.resolveJWKURL(request)
 
-			ja.logger.Info("JWK unresolved", zap.String("placeholder_url", ja.JWKURL))
-			ja.logger.Info("JWK resolved", zap.String("placeholder_url", resolvedURL))
+			ja.logger.Info("JWK URL", zap.String("unresolved", ja.JWKURL), zap.String("resolved", resolvedURL))
 
 			// Get or create the cache for this URL
 			cacheEntry, err := ja.getOrCreateJWKCache(resolvedURL)
@@ -318,10 +312,10 @@ func (ja *JWTAuth) keyProvider(request *http.Request) jws.KeyProviderFunc {
 
 			// Use the key set associated with this URL
 			kid := sig.ProtectedHeaders().KeyID()
-			key, found := cacheEntry.cachedSet.LookupKeyID(kid)
+			key, found := cacheEntry.CachedSet.LookupKeyID(kid)
 			if !found {
 				// Trigger an asynchronous refresh if the key is not found
-				go ja.refreshJWKCache(resolvedURL)
+				go cacheEntry.refresh(context.Background(), ja.logger)
 
 				if kid == "" {
 					return fmt.Errorf("missing kid in JWT header")

jwt_test.go

@@ -806,7 +806,7 @@ func TestJWK(t *testing.T) {
 	cachedEntry, exists := ja.jwkCaches[TestJWKURL]
 	assert.True(t, exists, "Le cache devrait exister pour l'URL de test")
 	assert.NotNil(t, cachedEntry)
-	assert.Equal(t, 1, cachedEntry.cachedSet.Len())
+	assert.Equal(t, 1, cachedEntry.CachedSet.Len())
 }
 
 func TestJWKSet(t *testing.T) {
@@ -834,7 +834,7 @@ func TestJWKSet(t *testing.T) {
 	cachedEntry, exists := ja.jwkCaches[TestJWKSetURL]
 	assert.True(t, exists, "Le cache devrait exister pour l'URL du set JWK")
 	assert.NotNil(t, cachedEntry)
-	assert.Equal(t, 2, cachedEntry.cachedSet.Len())
+	assert.Equal(t, 2, cachedEntry.CachedSet.Len())
 }
 
 func TestJWKSet_KeyNotFound(t *testing.T) {
@@ -859,7 +859,7 @@ func TestJWKSet_KeyNotFound(t *testing.T) {
 	cachedEntry, exists := ja.jwkCaches[TestJWKSetURLInapplicable]
 	assert.True(t, exists, "Le cache devrait exister pour l'URL inapplicable")
 	assert.NotNil(t, cachedEntry)
-	assert.Equal(t, 2, cachedEntry.cachedSet.Len())
+	assert.Equal(t, 2, cachedEntry.CachedSet.Len())
 
 	// Vérifier que l'authentification a échoué car la clé n'est pas trouvée
 	assert.Error(t, err)