fix(router): avoid blocking token validation across JWKS providers

router-tests/security/authentication_test.go

@@ -1487,6 +1487,257 @@ func TestHttpJwksAuthorization(t *testing.T) {
 		})
 	})
 
+	t.Run("known kid owned by a later provider does not block on earlier providers", func(t *testing.T) {
+		t.Parallel()
+
+		// Provider A is listed first but will NOT own the token's kid. Its endpoint is slow
+		// so that any cross-provider unknown-kid refresh on the request path is unmistakable
+		// in the timing. refresh_unknown_kid is enabled so a cache miss triggers a blocking
+		// refresh (this is the bug being fixed).
+		cryptoA, err := jwks.NewRSACrypto("kid-a", jwkset.AlgRS256, 2048)
+		require.NoError(t, err)
+		serverA, err := jwks.NewServerWithCrypto(t, cryptoA)
+		require.NoError(t, err)
+		t.Cleanup(serverA.Close)
+		serverA.SetRespondTime(2 * time.Second)
+
+		// Provider B is listed second and owns the token's kid. Its cache is populated by the
+		// initial synchronous refresh that NewStorageFromHTTP performs at decoder construction,
+		// so kid-b is cached before the first request — no warm-up required.
+		cryptoB, err := jwks.NewRSACrypto("kid-b", jwkset.AlgRS256, 2048)
+		require.NoError(t, err)
+		serverB, err := jwks.NewServerWithCrypto(t, cryptoB)
+		require.NoError(t, err)
+		t.Cleanup(serverB.Close)
+
+		authenticators := testutils.ConfigureAuthWithJwksConfig(t, []authentication.JWKSConfig{
+			{
+				URL:             serverA.JWKSURL(), // first
+				RefreshInterval: 10 * time.Second,
+				RefreshUnknownKID: authentication.RefreshUnknownKIDConfig{
+					Enabled:  true,
+					Interval: 1 * time.Second,
+					Burst:    1,
+					MaxWait:  5 * time.Second, // larger than the slow refresh so it proceeds into it
+				},
+			},
+			{
+				URL:             serverB.JWKSURL(), // second — owns kid-b
+				RefreshInterval: 10 * time.Second,
+			},
+		})
+
+		accessController, err := core.NewAccessController(core.AccessControllerOptions{
+			Authenticators:         authenticators,
+			AuthenticationRequired: true,
+		})
+		require.NoError(t, err)
+
+		testenv.Run(t, &testenv.Config{
+			RouterOptions: []core.Option{
+				core.WithAccessController(accessController),
+			},
+		}, func(t *testing.T, xEnv *testenv.Environment) {
+			token, err := serverB.TokenForKID("kid-b", nil, false)
+			require.NoError(t, err)
+			header := http.Header{"Authorization": []string{"Bearer " + token}}
+
+			start := time.Now()
+			res, err := xEnv.MakeRequest(http.MethodPost, "/graphql", header, strings.NewReader(employeesQuery))
+			require.NoError(t, err)
+			defer func() { _ = res.Body.Close() }()
+			elapsed := time.Since(start)
+
+			require.Equal(t, http.StatusOK, res.StatusCode)
+			require.Less(t, elapsed, 500*time.Millisecond,
+				"validation of a key owned by a later provider must not block on earlier providers")
+		})
+	})
+
+	t.Run("known kid owned by the first provider validates fast (no regression)", func(t *testing.T) {
+		t.Parallel()
+
+		// Owner is listed first.
+		cryptoA, err := jwks.NewRSACrypto("kid-a", jwkset.AlgRS256, 2048)
+		require.NoError(t, err)
+		serverA, err := jwks.NewServerWithCrypto(t, cryptoA)
+		require.NoError(t, err)
+		t.Cleanup(serverA.Close)
+
+		// A slow, penalized provider sits second — it must never be reached for a first-provider key.
+		cryptoB, err := jwks.NewRSACrypto("kid-b", jwkset.AlgRS256, 2048)
+		require.NoError(t, err)
+		serverB, err := jwks.NewServerWithCrypto(t, cryptoB)
+		require.NoError(t, err)
+		t.Cleanup(serverB.Close)
+		serverB.SetRespondTime(2 * time.Second)
+
+		authenticators := testutils.ConfigureAuthWithJwksConfig(t, []authentication.JWKSConfig{
+			{
+				URL:             serverA.JWKSURL(), // first — owns kid-a
+				RefreshInterval: 10 * time.Second,
+			},
+			{
+				URL:             serverB.JWKSURL(), // second
+				RefreshInterval: 10 * time.Second,
+				RefreshUnknownKID: authentication.RefreshUnknownKIDConfig{
+					Enabled:  true,
+					Interval: 1 * time.Second,
+					Burst:    1,
+					MaxWait:  5 * time.Second,
+				},
+			},
+		})
+
+		accessController, err := core.NewAccessController(core.AccessControllerOptions{
+			Authenticators:         authenticators,
+			AuthenticationRequired: true,
+		})
+		require.NoError(t, err)
+
+		testenv.Run(t, &testenv.Config{
+			RouterOptions: []core.Option{
+				core.WithAccessController(accessController),
+			},
+		}, func(t *testing.T, xEnv *testenv.Environment) {
+			token, err := serverA.TokenForKID("kid-a", nil, false)
+			require.NoError(t, err)
+			header := http.Header{"Authorization": []string{"Bearer " + token}}
+
+			start := time.Now()
+			res, err := xEnv.MakeRequest(http.MethodPost, "/graphql", header, strings.NewReader(employeesQuery))
+			require.NoError(t, err)
+			defer func() { _ = res.Body.Close() }()
+			elapsed := time.Since(start)
+
+			require.Equal(t, http.StatusOK, res.StatusCode)
+			require.Less(t, elapsed, 500*time.Millisecond)
+		})
+	})
+
+	t.Run("genuinely unknown kid still blocks across providers (unchanged semantics)", func(t *testing.T) {
+		t.Parallel()
+
+		// Two providers, neither owning the token's kid, both slow + refresh_unknown_kid enabled.
+		// Fix A intentionally leaves this path unchanged: a kid owned by nobody still pays the
+		// blocking refresh on each provider.
+		cryptoA, err := jwks.NewRSACrypto("kid-a", jwkset.AlgRS256, 2048)
+		require.NoError(t, err)
+		serverA, err := jwks.NewServerWithCrypto(t, cryptoA)
+		require.NoError(t, err)
+		t.Cleanup(serverA.Close)
+		serverA.SetRespondTime(1 * time.Second)
+
+		cryptoB, err := jwks.NewRSACrypto("kid-b", jwkset.AlgRS256, 2048)
+		require.NoError(t, err)
+		serverB, err := jwks.NewServerWithCrypto(t, cryptoB)
+		require.NoError(t, err)
+		t.Cleanup(serverB.Close)
+		serverB.SetRespondTime(1 * time.Second)
+
+		unknownKIDConfig := authentication.RefreshUnknownKIDConfig{
+			Enabled:  true,
+			Interval: 1 * time.Second,
+			Burst:    1,
+			MaxWait:  5 * time.Second,
+		}
+		authenticators := testutils.ConfigureAuthWithJwksConfig(t, []authentication.JWKSConfig{
+			{URL: serverA.JWKSURL(), RefreshInterval: 10 * time.Second, RefreshUnknownKID: unknownKIDConfig},
+			{URL: serverB.JWKSURL(), RefreshInterval: 10 * time.Second, RefreshUnknownKID: unknownKIDConfig},
+		})
+
+		accessController, err := core.NewAccessController(core.AccessControllerOptions{
+			Authenticators:         authenticators,
+			AuthenticationRequired: true,
+		})
+		require.NoError(t, err)
+
+		testenv.Run(t, &testenv.Config{
+			RouterOptions: []core.Option{
+				core.WithAccessController(accessController),
+			},
+		}, func(t *testing.T, xEnv *testenv.Environment) {
+			// Signed by a real key but stamped with a kid no provider publishes.
+			token, err := serverA.TokenForKID("unknown-kid", nil, true)
+			require.NoError(t, err)
+			header := http.Header{"Authorization": []string{"Bearer " + token}}
+
+			start := time.Now()
+			res, err := xEnv.MakeRequest(http.MethodPost, "/graphql", header, strings.NewReader(employeesQuery))
+			require.NoError(t, err)
+			defer func() { _ = res.Body.Close() }()
+			elapsed := time.Since(start)
+
+			require.Equal(t, http.StatusUnauthorized, res.StatusCode)
+			require.GreaterOrEqual(t, elapsed, 700*time.Millisecond,
+				"an unknown kid should still pay the blocking refresh on the providers")
+		})
+	})
+
+	t.Run("known kid owned by the third provider does not block on earlier providers", func(t *testing.T) {
+		t.Parallel()
+
+		// Providers #1 and #2 are slow + penalized and do NOT own the kid; #3 owns it. Proves
+		// the cost of validating a later-provider key does not scale with config position.
+		cryptoA, err := jwks.NewRSACrypto("kid-a", jwkset.AlgRS256, 2048)
+		require.NoError(t, err)
+		serverA, err := jwks.NewServerWithCrypto(t, cryptoA)
+		require.NoError(t, err)
+		t.Cleanup(serverA.Close)
+		serverA.SetRespondTime(2 * time.Second)
+
+		cryptoB, err := jwks.NewRSACrypto("kid-b", jwkset.AlgRS256, 2048)
+		require.NoError(t, err)
+		serverB, err := jwks.NewServerWithCrypto(t, cryptoB)
+		require.NoError(t, err)
+		t.Cleanup(serverB.Close)
+		serverB.SetRespondTime(2 * time.Second)
+
+		cryptoC, err := jwks.NewRSACrypto("kid-c", jwkset.AlgRS256, 2048)
+		require.NoError(t, err)
+		serverC, err := jwks.NewServerWithCrypto(t, cryptoC)
+		require.NoError(t, err)
+		t.Cleanup(serverC.Close)
+
+		penalized := authentication.RefreshUnknownKIDConfig{
+			Enabled:  true,
+			Interval: 1 * time.Second,
+			Burst:    1,
+			MaxWait:  5 * time.Second,
+		}
+		authenticators := testutils.ConfigureAuthWithJwksConfig(t, []authentication.JWKSConfig{
+			{URL: serverA.JWKSURL(), RefreshInterval: 10 * time.Second, RefreshUnknownKID: penalized},
+			{URL: serverB.JWKSURL(), RefreshInterval: 10 * time.Second, RefreshUnknownKID: penalized},
+			{URL: serverC.JWKSURL(), RefreshInterval: 10 * time.Second}, // owns kid-c
+		})
+
+		accessController, err := core.NewAccessController(core.AccessControllerOptions{
+			Authenticators:         authenticators,
+			AuthenticationRequired: true,
+		})
+		require.NoError(t, err)
+
+		testenv.Run(t, &testenv.Config{
+			RouterOptions: []core.Option{
+				core.WithAccessController(accessController),
+			},
+		}, func(t *testing.T, xEnv *testenv.Environment) {
+			token, err := serverC.TokenForKID("kid-c", nil, false)
+			require.NoError(t, err)
+			header := http.Header{"Authorization": []string{"Bearer " + token}}
+
+			start := time.Now()
+			res, err := xEnv.MakeRequest(http.MethodPost, "/graphql", header, strings.NewReader(employeesQuery))
+			require.NoError(t, err)
+			defer func() { _ = res.Body.Close() }()
+			elapsed := time.Since(start)
+
+			require.Equal(t, http.StatusOK, res.StatusCode)
+			require.Less(t, elapsed, 500*time.Millisecond,
+				"validation cost must not scale with the owning provider's position")
+		})
+	})
+
 }
 
 func TestNonHttpAuthorization(t *testing.T) {

router/pkg/authentication/jwks_token_decoder.go

@@ -71,7 +71,13 @@ type configKey struct {
 type audienceSet map[string]struct{}
 
 type keyFuncEntry struct {
-	jwks              keyfunc.Keyfunc
+	jwks keyfunc.Keyfunc
+	// storage is the raw, refresh-free cache handle for this provider. Reading it directly
+	// (KeyRead) is a non-blocking in-memory lookup that never triggers the rate-limited
+	// unknown-kid HTTP refresh. It is used for the cross-provider cache pre-pass in
+	// keyFuncWrapper. Do NOT use jwks.Storage() here: that returns the combined HTTP client
+	// whose KeyRead performs the blocking refresh on a miss.
+	storage           jwkset.Storage
 	aud               audienceSet
 	allowedAlgorithms []string
 	allowedUse        []string
@@ -135,6 +141,7 @@ func NewJwksTokenDecoder(ctx context.Context, logger *zap.Logger, configs []JWKS
 			}
 			entries = append(entries, keyFuncEntry{
 				jwks:              jwks,
+				storage:           store,
 				aud:               audiencesMap[key],
 				allowedAlgorithms: c.AllowedAlgorithms,
 				allowedUse:        c.AllowedUse,
@@ -195,6 +202,7 @@ func NewJwksTokenDecoder(ctx context.Context, logger *zap.Logger, configs []JWKS
 			}
 			entries = append(entries, keyFuncEntry{
 				jwks:              jwks,
+				storage:           given,
 				aud:               audiencesMap[key],
 				allowedAlgorithms: c.AllowedAlgorithms,
 				allowedUse:        c.AllowedUse,
@@ -203,43 +211,33 @@ func NewJwksTokenDecoder(ctx context.Context, logger *zap.Logger, configs []JWKS
 	}
 
 	keyFuncWrapper := jwt.Keyfunc(func(token *jwt.Token) (any, error) {
-		var errJoin error
-		for _, entry := range entries {
-			if len(entry.aud) > 0 {
-				tokenAudiences, err := token.Claims.GetAudience()
-				if err != nil {
-					errJoin = errors.Join(errJoin, fmt.Errorf("could not get audiences from token claims: %w", err))
+		// Non-blocking cache pre-pass: when the token carries a kid, ask each provider's raw
+		// storage whether it already holds that kid. This is a refresh-free in-memory lookup,
+		// so a provider that does not own the kid is skipped without ever blocking on its
+		// unknown-kid rate limiter or slow endpoint. Validating against the owning provider's
+		// cache (the steady state for known/rotated keys) therefore never pays a sibling
+		// provider's blocking refresh. Kid-less tokens fall through to the loop below, which
+		// preserves keyfunc's "try every key" behavior.
+		if kid, ok := token.Header[jwkset.HeaderKID].(string); ok && kid != "" {
+			for _, entry := range entries {
+				if entry.storage == nil {
 					continue
 				}
-				if !hasAudience(tokenAudiences, entry.aud) {
-					errJoin = errors.Join(errJoin, errUnacceptableAud)
-					continue
+				if _, err := entry.storage.KeyRead(ctx, kid); err != nil {
+					continue // not cached here — cheap, non-blocking miss
 				}
-			}
-
-			// When an algorithm is actually provided in the jwks the current keyfunc will validate the
-			// jwks algorithm with it. But when no algorithm is provided (alg: none or missing alg)
-			// the default keyfunc will not validate the algorithm as it has nothing to cross check.
-			if len(entry.allowedAlgorithms) > 0 {
-				algInter, ok := token.Header["alg"]
-				if !ok {
-					errJoin = errors.Join(errJoin, fmt.Errorf("%w: could not find alg in JWT header", keyfunc.ErrKeyfunc))
-					continue
-				}
-				alg, ok := algInter.(string)
-				if !ok {
-					errJoin = errors.Join(errJoin, fmt.Errorf(`%w: the JWT header did not contain the "alg" parameter, which is required by RFC 7515 section 4.1.1`, keyfunc.ErrKeyfunc))
-					continue
-				}
-
-				// This is a custom validation different from the original keyfunc.Keyfunc
-				if !slices.Contains(entry.allowedAlgorithms, alg) {
-					errJoin = errors.Join(errJoin, fmt.Errorf("%w: could not find alg %s in allow list", keyfunc.ErrKeyfunc, alg))
-					continue
+				// The owning provider has this kid cached, so tryEntry will not block. On
+				// success we are done; on rejection (e.g. aud/alg mismatch) we fall through
+				// to the full loop below, which aggregates errors exactly as before.
+				if pub, err := tryEntry(token, entry); err == nil {
+					return pub, nil
 				}
 			}
+		}
 
-			pub, err := entry.jwks.Keyfunc(token)
+		var errJoin error
+		for _, entry := range entries {
+			pub, err := tryEntry(token, entry)
 			if err != nil {
 				errJoin = errors.Join(errJoin, err)
 				continue
@@ -255,6 +253,45 @@ func NewJwksTokenDecoder(ctx context.Context, logger *zap.Logger, configs []JWKS
 	}, nil
 }
 
+// tryEntry runs the per-provider audience and algorithm allow-list checks and then resolves
+// the signing key via the provider's keyfunc. It contains the exact validation logic that was
+// previously inlined in the keyFuncWrapper loop, so behavior is identical whether an entry is
+// reached via the non-blocking cache pre-pass or the sequential fallback loop. When the entry's
+// cache holds the token's kid this does not block; on a cache miss for a provider with
+// refresh_unknown_kid enabled, entry.jwks.Keyfunc may block on the rate limiter as before.
+func tryEntry(token *jwt.Token, entry keyFuncEntry) (any, error) {
+	if len(entry.aud) > 0 {
+		tokenAudiences, err := token.Claims.GetAudience()
+		if err != nil {
+			return nil, fmt.Errorf("could not get audiences from token claims: %w", err)
+		}
+		if !hasAudience(tokenAudiences, entry.aud) {
+			return nil, errUnacceptableAud
+		}
+	}
+
+	// When an algorithm is actually provided in the jwks the current keyfunc will validate the
+	// jwks algorithm with it. But when no algorithm is provided (alg: none or missing alg)
+	// the default keyfunc will not validate the algorithm as it has nothing to cross check.
+	if len(entry.allowedAlgorithms) > 0 {
+		algInter, ok := token.Header["alg"]
+		if !ok {
+			return nil, fmt.Errorf("%w: could not find alg in JWT header", keyfunc.ErrKeyfunc)
+		}
+		alg, ok := algInter.(string)
+		if !ok {
+			return nil, fmt.Errorf(`%w: the JWT header did not contain the "alg" parameter, which is required by RFC 7515 section 4.1.1`, keyfunc.ErrKeyfunc)
+		}
+
+		// This is a custom validation different from the original keyfunc.Keyfunc
+		if !slices.Contains(entry.allowedAlgorithms, alg) {
+			return nil, fmt.Errorf("%w: could not find alg %s in allow list", keyfunc.ErrKeyfunc, alg)
+		}
+	}
+
+	return entry.jwks.Keyfunc(token)
+}
+
 func getAudienceSet(audiences []string) audienceSet {
 	audSet := make(audienceSet, len(audiences))
 	for _, aud := range audiences {