R54: refresh JWKS on unknown-kid to pick up key rotation

2026-10-12 04:04:26 +00:00 · 2026-08-13 15:30:58 +01:00
parent 08f474b551
commit 4c5991f8e6
3 changed files with 195 additions and 3 deletions
@@ -107,7 +107,7 @@ func (cm *CacheManager) GetSharedMetadataCache() *MetadataCache {
 func (cm *CacheManager) GetSharedJWKCache() JWKCacheInterface {
 	cm.mu.RLock()
 	defer cm.mu.RUnlock()
-	return &JWKCache{cache: cm.manager.GetJWKCache()}
+	return &JWKCache{cache: cm.manager.GetJWKCache(), lastForceRefresh: make(map[string]time.Time)}
 }

 // GetSharedIntrospectionCache returns the shared token introspection cache
@@ -65,6 +65,14 @@ type JWKSet struct {
 type JWKCache struct {
 	cache           *UniversalCache
 	inflightFetches sync.Map // map[jwksURL string]*jwksFetch
+
+	// forceMu guards lastForceRefresh: per-URL timestamps used to bound the
+	// live refresh triggered when a kid is missing from the cached derived
+	// keyset (see forceJWKSRefresh). Without the bound, an attacker who
+	// presents many tokens with unknown kid values could amplify N requests
+	// into N upstream JWKS fetches.
+	forceMu          sync.Mutex
+	lastForceRefresh map[string]time.Time
 }

 // jwksFetch represents an in-flight JWKS fetch. Done is closed when the fetch
@@ -87,7 +95,8 @@ type JWKCacheInterface interface {
 func NewJWKCache() *JWKCache {
 	manager := GetUniversalCacheManager(nil)
 	return &JWKCache{
-		cache: manager.GetJWKCache(),
+		cache:            manager.GetJWKCache(),
+		lastForceRefresh: make(map[string]time.Time),
 	}
 }

@@ -154,6 +163,63 @@ func (c *JWKCache) GetJWKS(ctx context.Context, jwksURL string, httpClient *http
 	return jwks, nil
 }

+// jwksForceRefreshCooldown bounds how often GetPublicKey forces a live JWKS
+// fetch when the requested kid is missing from the cached derived keyset.
+// Rotation is picked up on the first new-kid request; the cooldown
+// prevents an attacker presenting many bogus kids from amplifying a flood of
+// requests into a flood of upstream JWKS fetches.
+const jwksForceRefreshCooldown = 30 * time.Second
+
+// forceJWKSRefresh performs a live singleflighted JWKS fetch, bypassing the
+// fast-path cache, so an IdP key rotation is picked up promptly. It is
+// bounded by jwksForceRefreshCooldown per URL: within the window it falls
+// back to the (freshly refetched) cached set rather than fetching again.
+// On success the result is written to the raw cache so subsequent fast-path
+// GetJWKS calls and other replicas-benefits see the rotated keys.
+func (c *JWKCache) forceJWKSRefresh(ctx context.Context, jwksURL string, httpClient *http.Client) (*JWKSet, error) {
+	c.forceMu.Lock()
+	if c.lastForceRefresh == nil {
+		c.lastForceRefresh = make(map[string]time.Time)
+	}
+	last := c.lastForceRefresh[jwksURL]
+	now := time.Now()
+	if now.Sub(last) < jwksForceRefreshCooldown {
+		c.forceMu.Unlock()
+		return c.GetJWKS(ctx, jwksURL, httpClient)
+	}
+	c.lastForceRefresh[jwksURL] = now
+	c.forceMu.Unlock()
+
+	// Singleflight a live fetch (no fast-path cache read).
+	candidate := &jwksFetch{done: make(chan struct{})}
+	if existing, loaded := c.inflightFetches.LoadOrStore(jwksURL, candidate); loaded {
+		f, _ := existing.(*jwksFetch)
+		select {
+		case <-f.done:
+			return f.jwks, f.err
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		}
+	}
+	defer func() {
+		c.inflightFetches.Delete(jwksURL)
+		close(candidate.done)
+	}()
+
+	jwks, err := fetchJWKS(ctx, jwksURL, httpClient)
+	if err != nil {
+		candidate.err = err
+		return nil, err
+	}
+	if len(jwks.Keys) == 0 {
+		candidate.err = fmt.Errorf("JWKS response contains no keys")
+		return nil, candidate.err
+	}
+	_ = c.cache.SetLocal(jwksURL, jwks, 1*time.Hour) // Safe to ignore: cache failures are non-critical
+	candidate.jwks = jwks
+	return jwks, nil
+}
+
 // GetPublicKey returns the parsed public key for a given kid, fetching and
 // caching the JWKS plus its derived parsedJWKS on miss. The parsed entry is
 // stored alongside the raw JWKSet under a sibling cache key with the same
@@ -166,6 +232,12 @@ func (c *JWKCache) GetJWKS(ctx context.Context, jwksURL string, httpClient *http
 // try to fit that into float64 and fail with UnmarshalTypeError. Under yaegi
 // the unexported parsedJWKS.keys field is exposed via an X-prefixed name on
 // Marshal, leaking the modulus into the cached payload (issue #134).
+//
+// If the requested kid is not in the cached derived keyset, GetJWKS would
+// still hand back the stale (within-TTL) cached JWKS, so a post-rotation
+// token with a NEW kid would fail "no matching public key" for up to the
+// 1-hour TTL. Instead, on a kid miss we force a bounded live refresh so
+// rotation is picked up on the first new-kid request.
 func (c *JWKCache) GetPublicKey(ctx context.Context, jwksURL, kid string, httpClient *http.Client) (crypto.PublicKey, error) {
 	parsedKey := jwksURL + parsedKeysSuffix
 	if v, found := c.cache.GetLocal(parsedKey); found {
@@ -176,7 +248,7 @@ func (c *JWKCache) GetPublicKey(ctx context.Context, jwksURL, kid string, httpCl
 		}
 	}

-	jwks, err := c.GetJWKS(ctx, jwksURL, httpClient)
+	jwks, err := c.forceJWKSRefresh(ctx, jwksURL, httpClient)
 	if err != nil {
 		return nil, err
 	}
@@ -0,0 +1,120 @@
+package traefikoidc
+
+import (
+	"context"
+	"crypto/rand"
+	"crypto/rsa"
+	"encoding/base64"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+// TestJWKSRotationPickedUpOnKidMiss reproduces the key-rotation availability
+// bug: after the IdP rotates signing keys and issues a NEW kid, GetPublicKey
+// must pick the new key up promptly (bounded live refresh) instead of
+// returning "no matching public key found for kid" for up to the 1h cache
+// TTL because the stale cached keyset is reused.
+func TestJWKSRotationPickedUpOnKidMiss(t *testing.T) {
+	require := require.New(t)
+
+	// Two RSA key pairs: keyA is the "pre-rotation" key, keyB the rotated one.
+	keyA, err := rsa.GenerateKey(rand.Reader, 2048)
+	require.NoError(err)
+	keyB, err := rsa.GenerateKey(rand.Reader, 2048)
+	require.NoError(err)
+
+	// The server serves kidA until the atomic flag is set, then kidB.
+	var rotated int32
+	var mu sync.Mutex
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		mu.Lock()
+		defer mu.Unlock()
+		if atomic.LoadInt32(&rotated) == 1 {
+			writeJWKS(t, w, "kidB", &keyB.PublicKey)
+			return
+		}
+		writeJWKS(t, w, "kidA", &keyA.PublicKey)
+	}))
+	defer srv.Close()
+
+	cache := NewJWKCache()
+
+	// Seed the raw cache with the pre-rotation keyset (as happens on a cold
+	// start), WITHOUT taking the kid-miss force path.
+	_, err = cache.GetJWKS(context.Background(), srv.URL, http.DefaultClient)
+	require.NoError(err)
+
+	// Pre-rotation token for kidA resolves fine.
+	_, err = cache.GetPublicKey(context.Background(), srv.URL, "kidA", http.DefaultClient)
+	require.NoError(err, "pre-rotation kid should resolve")
+
+	// Simulate a long-lived process: the key rotation happens well after the
+	// process start (and thus well after any earlier live fetch), so the
+	// force-refresh cooldown has long expired.
+	cache.forceMu.Lock()
+	cache.lastForceRefresh[srv.URL] = time.Time{}
+	cache.forceMu.Unlock()
+
+	// Provider rotates keys; now issues a token with the NEW kid.
+	atomic.StoreInt32(&rotated, 1)
+
+	// The new kid must be resolved now (bounded live refresh on kid-miss),
+	// not after the 1-hour cache TTL.
+	pub, err := cache.GetPublicKey(context.Background(), srv.URL, "kidB", http.DefaultClient)
+	require.NoError(err, "post-rotation new kid must be picked up promptly")
+	require.NotNil(pub)
+	require.Equal(&keyB.PublicKey, pub, "resolved key must be the rotated key")
+}
+
+// TestJWKSRotationBoundRefetch verifies the cooldown bounds repeated live
+// fetches when an attacker presents bogus unknown kids, while a real rotation
+// still succeeds on the first new-kid request.
+func TestJWKSRotationBoundRefetch(t *testing.T) {
+	require := require.New(t)
+
+	keyA, err := rsa.GenerateKey(rand.Reader, 2048)
+	require.NoError(err)
+
+	var fetches int32
+	var mu sync.Mutex
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		mu.Lock()
+		defer mu.Unlock()
+		atomic.AddInt32(&fetches, 1)
+		writeJWKS(t, w, "kidA", &keyA.PublicKey)
+	}))
+	defer srv.Close()
+
+	cache := NewJWKCache()
+
+	// First kid-miss forces exactly one live fetch and returns "no key".
+	_, err = cache.GetPublicKey(context.Background(), srv.URL, "bogus1", http.DefaultClient)
+	require.Error(err)
+
+	// Within the cooldown window, further bogus kids must NOT each trigger a
+	// live fetch (amplification bound) — they reuse the freshly cached set.
+	_, err = cache.GetPublicKey(context.Background(), srv.URL, "bogus2", http.DefaultClient)
+	require.Error(err)
+
+	time.Sleep(100 * time.Millisecond)
+	require.Equal(int32(1), atomic.LoadInt32(&fetches), "only one live fetch within the cooldown window")
+}
+
+func writeJWKS(t *testing.T, w http.ResponseWriter, kid string, pub *rsa.PublicKey) {
+	t.Helper()
+	jwk := JWK{
+		Kty: "RSA",
+		Kid: kid,
+		Alg: "RS256",
+		N:   base64.RawURLEncoding.EncodeToString(pub.N.Bytes()),
+		E:   base64.RawURLEncoding.EncodeToString([]byte{1, 0, 1}),
+	}
+	_ = json.NewEncoder(w).Encode(JWKSet{Keys: []JWK{jwk}})
+}