Lazily register JWKS endpoint with a 5 second timeout instead of on s… (#1472)

Author: jhrozek

pkg/auth/token.go

@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/golang-jwt/jwt/v5"
@@ -57,7 +58,10 @@ type TokenValidator struct {
 	introspectURL string       // Optional introspection endpoint
 	client        *http.Client // HTTP client for making requests
 
-	// No need for additional caching as jwk.Cache handles it
+	// Lazy JWKS registration
+	jwksRegistered      bool
+	jwksRegistrationMu  sync.Mutex
+	jwksRegistrationErr error
 }
 
 // TokenValidatorConfig contains configuration for the token validator.
@@ -176,7 +180,6 @@ func NewTokenValidator(ctx context.Context, config TokenValidatorConfig) (*Token
 			return nil, fmt.Errorf("%w: %v", ErrFailedToDiscoverOIDC, err)
 		}
 		jwksURL = doc.JWKSURI
-
 	}
 
 	// Ensure we have a JWKS URL either provided or discovered
@@ -203,11 +206,7 @@ func NewTokenValidator(ctx context.Context, config TokenValidatorConfig) (*Token
 		return nil, fmt.Errorf("failed to create JWKS cache: %w", err)
 	}
 
-	// Register the JWKS URL with the cache
-	err = cache.Register(ctx, jwksURL)
-	if err != nil {
-		return nil, fmt.Errorf("failed to register JWKS URL: %w", err)
-	}
+	// Skip synchronous JWKS registration - will be done lazily on first use
 
 	return &TokenValidator{
 		issuer:        config.Issuer,
@@ -221,8 +220,40 @@ func NewTokenValidator(ctx context.Context, config TokenValidatorConfig) (*Token
 	}, nil
 }
 
+// ensureJWKSRegistered ensures that the JWKS URL is registered with the cache.
+// This is called lazily on first use to avoid blocking startup.
+func (v *TokenValidator) ensureJWKSRegistered(ctx context.Context) error {
+	v.jwksRegistrationMu.Lock()
+	defer v.jwksRegistrationMu.Unlock()
+
+	// Check if already registered or failed
+	if v.jwksRegistered {
+		return v.jwksRegistrationErr
+	}
+
+	// Create context with 5-second timeout for JWKS registration
+	registrationCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
+
+	// Attempt registration
+	err := v.jwksClient.Register(registrationCtx, v.jwksURL)
+	if err != nil {
+		v.jwksRegistrationErr = fmt.Errorf("failed to register JWKS URL: %w", err)
+	} else {
+		v.jwksRegistrationErr = nil
+	}
+
+	v.jwksRegistered = true
+	return v.jwksRegistrationErr
+}
+
 // getKeyFromJWKS gets the key from the JWKS.
 func (v *TokenValidator) getKeyFromJWKS(ctx context.Context, token *jwt.Token) (interface{}, error) {
+	// Ensure JWKS is registered before attempting to use it
+	if err := v.ensureJWKSRegistered(ctx); err != nil {
+		return nil, fmt.Errorf("JWKS registration failed: %w", err)
+	}
+
 	// Validate the signing method
 	if _, ok := token.Method.(*jwt.SigningMethodRSA); !ok {
 		return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])

pkg/auth/token_test.go

@@ -75,6 +75,12 @@ func TestTokenValidator(t *testing.T) {
 		t.Fatalf("Failed to create token validator: %v", err)
 	}
 
+	// Ensure JWKS is registered before lookup
+	err = validator.ensureJWKSRegistered(ctx)
+	if err != nil {
+		t.Fatalf("Failed to register JWKS: %v", err)
+	}
+
 	// Force a refresh of the JWKS cache
 	_, err = validator.jwksClient.Lookup(ctx, jwksServer.URL)
 	if err != nil {
@@ -217,6 +223,12 @@ func TestTokenValidatorMiddleware(t *testing.T) {
 		t.Fatalf("Failed to create token validator: %v", err)
 	}
 
+	// Ensure JWKS is registered before lookup
+	err = validator.ensureJWKSRegistered(ctx)
+	if err != nil {
+		t.Fatalf("Failed to register JWKS: %v", err)
+	}
+
 	// Force a refresh of the JWKS cache
 	_, err = validator.jwksClient.Lookup(ctx, jwksServer.URL)
 	if err != nil {
@@ -607,6 +619,12 @@ func TestNewTokenValidatorWithOIDCDiscovery(t *testing.T) {
 			t.Fatalf("Failed to sign token: %v", err)
 		}
 
+		// Ensure JWKS is registered before lookup
+		err = validator.ensureJWKSRegistered(ctx)
+		if err != nil {
+			t.Fatalf("Failed to register JWKS: %v", err)
+		}
+
 		// Force a refresh of the JWKS cache
 		_, err = validator.jwksClient.Lookup(ctx, validator.jwksURL)
 		if err != nil {