From 4c5991f8e66986d82d7431b1ab7cd02036ebd4f9 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Thu, 13 Aug 2026 15:30:58 +0100 Subject: [PATCH] R54: refresh JWKS on unknown-kid to pick up key rotation --- cache_manager.go | 2 +- jwk.go | 76 +++++++++++++++++++- review_r54_jwks_rotation_test.go | 120 +++++++++++++++++++++++++++++++ 3 files changed, 195 insertions(+), 3 deletions(-) create mode 100644 review_r54_jwks_rotation_test.go diff --git a/cache_manager.go b/cache_manager.go index fe9cc0c..f4cb58b 100644 --- a/cache_manager.go +++ b/cache_manager.go @@ -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 diff --git a/jwk.go b/jwk.go index e69d0ff..b4192e3 100644 --- a/jwk.go +++ b/jwk.go @@ -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 } diff --git a/review_r54_jwks_rotation_test.go b/review_r54_jwks_rotation_test.go new file mode 100644 index 0000000..e3cc707 --- /dev/null +++ b/review_r54_jwks_rotation_test.go @@ -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}}) +}