iam/oidc: wrap the jti replay cache in its own type (#12017)

* iam/oidc: wrap the jti replay cache in its own type

The replay-protection state lived as ten fields and nine methods spread
across OIDCProvider. jtiReplayCache now owns the store, count, bound,
sweep scheduling, and cleanup worker; the provider keeps a single jti
field and delegates configure, consume, stop, adopt, and merge.

* iam/oidc: lazily allocate the jti replay cache

A zero-value OIDCProvider left jti nil, so Initialize, single-use
validation, Shutdown, and the replay-state handoffs panicked on a nil
cache. Allocate the cache through a sync.Once-guarded getter instead.

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>

weed/iam/oidc/jti_cache.go

@@ -0,0 +1,410 @@
+package oidc
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"fmt"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/seaweedfs/seaweedfs/weed/glog"
+	"github.com/seaweedfs/seaweedfs/weed/iam/providers"
+)
+
+const (
+	// clockSkewTolerance matches the typical JWT library clock skew allowance (RFC 7519)
+	clockSkewTolerance = 5 * time.Minute
+	// defaultJTICacheSize is the default maximum number of JTI entries to cache
+	defaultJTICacheSize = 10000
+	// defaultCleanupInterval is how often to scan for expired JTI entries
+	defaultCleanupInterval = 5 * time.Minute
+	// eagerCleanupMinInterval rate limits opportunistic sweeps near capacity
+	eagerCleanupMinInterval = 5 * time.Second
+	// tokenFingerprintPrefix marks replay-cache keys derived from the token's
+	// own fingerprint, used when the token carries no usable jti claim
+	tokenFingerprintPrefix = "fp:"
+)
+
+// jtiEntry represents a cached JTI with metadata
+type jtiEntry struct {
+	subject   string
+	expiresAt time.Time
+	// uncounted marks an entry merged in past the capacity bound: it keeps
+	// rejecting replays but does not charge the count, so merged excess can
+	// neither reject new exchanges nor be double-counted when swept.
+	uncounted bool
+}
+
+// jtiReplayCache wraps the consumed-token store that backs JTI replay
+// protection. Each consumed token identifier is held until the token's exp
+// plus clock skew, and a second presentation of the same identifier is
+// rejected with ErrProviderTokenReplayed.
+//
+// The cache is bounded by maxSize and uses automatic TTL-based expiration:
+// a background worker sweeps expired entries, an eager sweep runs when the
+// cache nears capacity, and the consume path reclaims expired slots before
+// rejecting a fresh token for capacity.
+//
+// Thread Safety:
+// All methods are safe for concurrent use. The consume path uses atomic
+// LoadOrStore to prevent TOCTOU races.
+type jtiReplayCache struct {
+	name string // owning provider name, for log messages
+
+	enabled   atomic.Bool
+	maxSize   atomic.Int64
+	lastSweep atomic.Int64 // unix nano of the last sweep; rate limits eager cleanups
+
+	// storeMu guards the store/count pair: replay-state handoffs (adoptFrom,
+	// mergeFrom) repoint them under the write lock while the consume path
+	// holds a read lock, so a consume can never land in a store the handoff
+	// already drained and abandoned.
+	// It is always acquired before sweepMu, never while holding it.
+	storeMu sync.RWMutex
+	store   *sync.Map     // map[string]*jtiEntry - atomic storage for consumed token keys
+	count   *atomic.Int64 // entries in store; shared with the store
+
+	// mu guards the lifecycle of the cleanup worker: cleanupCancel and
+	// cleanupDone. The worker itself only touches store/count, so holding
+	// mu while waiting for it is safe.
+	mu            sync.Mutex
+	cleanupCancel context.CancelFunc
+	cleanupDone   chan struct{} // Closed when cleanup goroutine exits
+
+	// sweepMu serializes cache sweeps so at most one scan runs at a time.
+	sweepMu sync.Mutex
+}
+
+func newJTIReplayCache(name string) *jtiReplayCache {
+	return &jtiReplayCache{
+		name:  name,
+		store: &sync.Map{},
+		count: &atomic.Int64{},
+	}
+}
+
+// configure applies the replay-protection settings: a previous worker is
+// always stopped, whether the new configuration enables protection or
+// disables it. The consumed-token store survives re-configuration: resetting
+// it would let already-seen tokens be replayed against the refreshed provider.
+func (c *jtiReplayCache) configure(enabled bool, maxSize int64) {
+	c.mu.Lock()
+	c.stopWorkerLocked()
+	if enabled {
+		if maxSize <= 0 {
+			maxSize = defaultJTICacheSize
+		}
+		// the bound must land before the flag: a consume that sees enabled
+		// with a zero bound would reject the first token it registers
+		c.maxSize.Store(maxSize)
+		ctx, cancel := context.WithCancel(context.Background())
+		done := make(chan struct{})
+		c.cleanupCancel = cancel
+		c.cleanupDone = done
+		c.enabled.Store(true)
+		c.mu.Unlock()
+
+		go c.cleanupLoop(ctx, done, defaultCleanupInterval)
+		glog.V(2).Infof("OIDC provider %q: JTI replay protection enabled (max size: %d)", c.name, maxSize)
+	} else {
+		c.enabled.Store(false)
+		c.mu.Unlock()
+		glog.Warningf("OIDC provider %q: JTI replay protection is DISABLED", c.name)
+	}
+}
+
+// stopWorkerLocked cancels the current cleanup goroutine and waits briefly
+// for it to exit. Caller must hold mu.
+func (c *jtiReplayCache) stopWorkerLocked() {
+	if c.cleanupCancel == nil {
+		return
+	}
+	c.cleanupCancel()
+	select {
+	case <-c.cleanupDone:
+		glog.V(3).Infof("OIDC provider %q: JTI cleanup goroutine exited", c.name)
+	case <-time.After(2 * time.Second):
+		glog.Warningf("OIDC provider %q: Timeout waiting for JTI cleanup goroutine", c.name)
+	}
+	c.cleanupCancel = nil
+	c.cleanupDone = nil
+}
+
+// consume marks a token as used and rejects a second presentation of it.
+// jti is the token's jti claim (may be empty); tokenExpiresAt is its exp.
+// A no-op while protection is disabled.
+func (c *jtiReplayCache) consume(token, jti, subject string, tokenExpiresAt time.Time) error {
+	if !c.enabled.Load() {
+		return nil
+	}
+
+	key := jti
+	if key == "" {
+		// Providers that issue no jti cannot be tracked by claim, so the
+		// signed input keys the entry: a repeated presentation of the same
+		// claims is a replay even when the signature bytes differ (e.g.
+		// ECDSA s/N-s malleability).
+		signingInput := token
+		if i := strings.LastIndexByte(token, '.'); i > 0 {
+			signingInput = token[:i]
+		}
+		sum := sha256.Sum256([]byte(signingInput))
+		key = tokenFingerprintPrefix + hex.EncodeToString(sum[:])
+	}
+
+	if tokenExpiresAt.IsZero() {
+		// a cache slot is freed only by expiry, so a token without exp
+		// would hold it forever and eventually exhaust the bound
+		return fmt.Errorf("%w: missing exp claim", providers.ErrProviderMissingClaims)
+	}
+	expiresAt := tokenExpiresAt.Add(clockSkewTolerance)
+	if !time.Now().Before(expiresAt) {
+		return fmt.Errorf("%w: token expired", providers.ErrProviderTokenExpired)
+	}
+
+	entry := &jtiEntry{
+		subject:   subject,
+		expiresAt: expiresAt,
+	}
+
+	// The read lock pins the store/count pair for the whole test-and-set
+	// so a replay-state handoff cannot repoint them mid-check and strand
+	// this consume in an abandoned store.
+	c.storeMu.RLock()
+	store, count := c.store, c.count
+
+	// Atomic test-and-set detects replays before any capacity work runs,
+	// so a replayed token cannot trigger cleanup scans.
+	if _, loaded := store.LoadOrStore(key, entry); loaded {
+		c.storeMu.RUnlock()
+		glog.Warningf("OIDC provider %q: Token replay detected (jti: %s)",
+			c.name, jtiLogKey(key))
+		return fmt.Errorf("%w: jti %s has already been used",
+			providers.ErrProviderTokenReplayed, jtiLogKey(key))
+	}
+
+	// The count is incremented as part of the insert, so concurrent
+	// validations cannot all slip under the limit: whoever pushes the
+	// count over evicts their own entry and is rejected.
+	newCount := count.Add(1)
+	maxSize := c.maxSize.Load()
+	if newCount > maxSize {
+		// Over the bound: expired entries may still occupy slots, so
+		// reclaim them with a synchronized sweep before rejecting.
+		c.sweepStore(store, count)
+		if count.Load() > maxSize {
+			if _, loaded := store.LoadAndDelete(key); loaded {
+				count.Add(-1)
+			}
+			c.storeMu.RUnlock()
+			glog.Errorf("OIDC provider %q: JTI cache FULL (%d entries) - rejecting token. "+
+				"This may indicate a DoS attack or insufficient cache size. Consider increasing JTICacheMaxSize.",
+				c.name, maxSize)
+			return fmt.Errorf("JTI cache capacity exceeded")
+		}
+	} else if newCount >= maxSize*9/10 {
+		glog.Warningf("OIDC provider %q: JTI cache at %d%% capacity (%d/%d), triggering eager cleanup",
+			c.name, (newCount*100)/maxSize, newCount, maxSize)
+		go c.performEagerCleanup()
+	}
+	c.storeMu.RUnlock()
+
+	glog.V(4).Infof("OIDC provider %q: Registered JTI %s (expires at %v, cache size: %d)",
+		c.name, jtiLogKey(key), expiresAt, newCount)
+	return nil
+}
+
+// evictExpired removes entries whose expiry has passed and returns how
+// many it deleted. Caller must hold sweepMu.
+func (c *jtiReplayCache) evictExpired(now time.Time, store *sync.Map, count *atomic.Int64) int {
+	removed := 0
+	store.Range(func(key, value interface{}) bool {
+		expiresAt := value.(*jtiEntry).expiresAt
+		if !expiresAt.IsZero() && now.After(expiresAt) {
+			// LoadAndDelete reports whether this call actually removed the
+			// entry, so concurrent deletes cannot decrement the count twice.
+			if v, loaded := store.LoadAndDelete(key); loaded {
+				if !v.(*jtiEntry).uncounted {
+					count.Add(-1)
+				}
+				removed++
+			}
+		}
+		return true
+	})
+	return removed
+}
+
+// state returns the cache's current store/count pair.
+func (c *jtiReplayCache) state() (*sync.Map, *atomic.Int64) {
+	c.storeMu.RLock()
+	defer c.storeMu.RUnlock()
+	return c.store, c.count
+}
+
+// sweepStore runs one serialized expiry scan over the given store. It only
+// takes sweepMu, so callers already holding storeMu may use it.
+func (c *jtiReplayCache) sweepStore(store *sync.Map, count *atomic.Int64) {
+	c.sweepMu.Lock()
+	defer c.sweepMu.Unlock()
+	c.lastSweep.Store(time.Now().UnixNano())
+	if removed := c.evictExpired(time.Now(), store, count); removed > 0 {
+		glog.V(4).Infof("OIDC provider %q: Cleaned up %d expired JTI entries (remaining: %d)",
+			c.name, removed, count.Load())
+	}
+}
+
+// sweepExpired runs a scan that is guaranteed to complete before the
+// caller proceeds; the periodic worker uses it so the scan always covers
+// the store currently in effect, not one captured earlier.
+func (c *jtiReplayCache) sweepExpired() {
+	store, count := c.state()
+	c.sweepStore(store, count)
+}
+
+// performEagerCleanup performs an opportunistic cleanup of expired JTI
+// entries when the cache approaches capacity. At most one sweep runs at a
+// time and no more often than eagerCleanupMinInterval, so a burst of
+// requests near the limit cannot each turn into a full-cache scan.
+func (c *jtiReplayCache) performEagerCleanup() {
+	if !c.enabled.Load() {
+		return
+	}
+	if time.Since(time.Unix(0, c.lastSweep.Load())) < eagerCleanupMinInterval {
+		return
+	}
+	store, count := c.state()
+	if !c.sweepMu.TryLock() {
+		return
+	}
+	defer c.sweepMu.Unlock()
+	c.lastSweep.Store(time.Now().UnixNano())
+
+	beforeCount := count.Load()
+	if removed := c.evictExpired(time.Now(), store, count); removed > 0 {
+		glog.Infof("OIDC provider %q: Eager cleanup removed %d expired JTI entries (%d -> %d)",
+			c.name, removed, beforeCount, count.Load())
+	} else {
+		glog.Warningf("OIDC provider %q: Eager cleanup found no expired entries to remove. "+
+			"Cache may be filled with valid long-lived tokens. Consider increasing JTICacheMaxSize.",
+			c.name)
+	}
+}
+
+// cleanupLoop periodically removes expired JTI entries from the cache.
+// This prevents memory growth by removing JTIs after their tokens have
+// expired. Runs until ctx is cancelled; done is closed on exit.
+func (c *jtiReplayCache) cleanupLoop(ctx context.Context, done chan struct{}, interval time.Duration) {
+	defer close(done)
+
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+
+	glog.V(3).Infof("OIDC provider %q: JTI cleanup goroutine started (interval: %v)", c.name, interval)
+
+	for {
+		select {
+		case <-ctx.Done():
+			glog.V(3).Infof("OIDC provider %q: JTI cleanup goroutine stopping", c.name)
+			return
+		case <-ticker.C:
+			c.sweepExpired()
+		}
+	}
+}
+
+// stop gracefully stops the cleanup worker.
+// Should be called when the provider is being shut down to prevent goroutine
+// leaks.
+func (c *jtiReplayCache) stop(ctx context.Context) error {
+	c.mu.Lock()
+	cancel, done := c.cleanupCancel, c.cleanupDone
+	c.mu.Unlock()
+
+	if cancel == nil {
+		return nil
+	}
+	cancel()
+
+	select {
+	case <-done:
+		glog.V(2).Infof("OIDC provider %q: Cleanup goroutine shutdown complete", c.name)
+	case <-ctx.Done():
+		glog.Warningf("OIDC provider %q: Shutdown timeout waiting for cleanup goroutine: %v", c.name, ctx.Err())
+		return ctx.Err()
+	}
+
+	c.mu.Lock()
+	if c.cleanupDone == done {
+		c.cleanupCancel = nil
+		c.cleanupDone = nil
+	}
+	c.mu.Unlock()
+	return nil
+}
+
+// adoptFrom takes over the consumed-token store of another cache, so tokens
+// it already saw stay rejected across the rebuild. The caller must shut down
+// the replaced provider; only one cleanup worker may sweep the shared store
+// at a time. Call before the new provider serves validators; sweepMu
+// excludes its own cleanup worker during the swap.
+func (c *jtiReplayCache) adoptFrom(old *jtiReplayCache) {
+	if old == nil || old == c {
+		return
+	}
+	c.storeMu.Lock()
+	c.sweepMu.Lock()
+	c.store, c.count = old.store, old.count
+	c.sweepMu.Unlock()
+	c.storeMu.Unlock()
+}
+
+// mergeFrom hands this cache's store to another cache and then folds the
+// other's consumed-token entries in, so rejections the other recorded keep
+// holding when its records move to a shared provider. Repointing first means
+// tokens the other still consumes — in-flight exchanges or ones routed to it
+// before the provider map updates — land in the shared store instead of a
+// store this merge already drained. Unlike adoptFrom it swaps nothing of its
+// own, so it is safe on a cache that is already serving validators.
+func (c *jtiReplayCache) mergeFrom(other *jtiReplayCache) {
+	if other == nil || other == c {
+		return
+	}
+	c.storeMu.RLock()
+	defer c.storeMu.RUnlock()
+	store, count := c.store, c.count
+
+	other.storeMu.Lock()
+	oldStore := other.store
+	if oldStore == store {
+		other.storeMu.Unlock()
+		return
+	}
+	other.store, other.count = store, count
+	other.storeMu.Unlock()
+
+	c.sweepMu.Lock()
+	defer c.sweepMu.Unlock()
+	now := time.Now()
+	oldStore.Range(func(key, value interface{}) bool {
+		entry := value.(*jtiEntry)
+		if expiresAt := entry.expiresAt; !expiresAt.IsZero() && now.After(expiresAt) {
+			return true
+		}
+		// A consumed token is never dropped to fit the bound: merged entries
+		// stay but do not charge the count, so a store filled past the bound
+		// by the merge cannot reject new exchanges while it drains by expiry.
+		store.LoadOrStore(key, &jtiEntry{subject: entry.subject, expiresAt: entry.expiresAt, uncounted: true})
+		return true
+	})
+}
+
+// jtiLogKey redacts a token identifier for logs and error strings: a stable
+// fingerprint that still lets operators correlate repeated presentations.
+func jtiLogKey(jti string) string {
+	sum := sha256.Sum256([]byte(jti))
+	return hex.EncodeToString(sum[:8])
+}

weed/iam/oidc/oidc_provider.go

@@ -6,7 +6,6 @@ import (
 	"crypto/elliptic"
 	"crypto/rsa"
 	"crypto/sha1"
-	"crypto/sha256"
 	"crypto/tls"
 	"crypto/x509"
 	"encoding/base64"
@@ -20,7 +19,6 @@ import (
 	"path/filepath"
 	"strings"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/golang-jwt/jwt/v5"
@@ -29,30 +27,6 @@ import (
 	"github.com/seaweedfs/seaweedfs/weed/security"
 )
 
-const (
-	// clockSkewTolerance matches the typical JWT library clock skew allowance (RFC 7519)
-	clockSkewTolerance = 5 * time.Minute
-	// defaultJTICacheSize is the default maximum number of JTI entries to cache
-	defaultJTICacheSize = 10000
-	// defaultCleanupInterval is how often to scan for expired JTI entries
-	defaultCleanupInterval = 5 * time.Minute
-	// eagerCleanupMinInterval rate limits opportunistic sweeps near capacity
-	eagerCleanupMinInterval = 5 * time.Second
-	// tokenFingerprintPrefix marks replay-cache keys derived from the token's
-	// own fingerprint, used when the token carries no usable jti claim
-	tokenFingerprintPrefix = "fp:"
-)
-
-// jtiEntry represents a cached JTI with metadata
-type jtiEntry struct {
-	subject   string
-	expiresAt time.Time
-	// uncounted marks an entry merged in past the capacity bound: it keeps
-	// rejecting replays but does not charge jtiCount, so merged excess can
-	// neither reject new exchanges nor be double-counted when swept.
-	uncounted bool
-}
-
 // OIDCProvider implements OpenID Connect authentication with JTI replay protection.
 //
 // Replay Protection:
@@ -84,31 +58,24 @@ type OIDCProvider struct {
 	resolvedJWKSUri string
 	discoveryFailed bool
 
-	// JTI replay protection fields. jtiStore and jtiCount are allocated once
-	// in NewOIDCProvider and survive re-initialization, so a config refresh
-	// cannot reopen already-consumed tokens for replay.
-	jtiEnabled   atomic.Bool
-	jtiStore     *sync.Map     // map[string]*jtiEntry - atomic storage for consumed token keys
-	jtiCount     *atomic.Int64 // entries in jtiStore; shared with the store
-	jtiMaxSize   atomic.Int64
-	jtiLastSweep atomic.Int64 // unix nano of the last sweep; rate limits eager cleanups
+	// jti wraps the consumed-token cache behind the replay protection. It is
+	// allocated by NewOIDCProvider or lazily by replayCache, and survives
+	// re-initialization, so a config refresh cannot reopen already-consumed
+	// tokens for replay. jtiOnce makes a zero-value OIDCProvider safe.
+	jti     *jtiReplayCache
+	jtiOnce sync.Once
+}
 
-	// jtiStoreMu guards the jtiStore/jtiCount pair: replay-state handoffs
-	// (AdoptReplayStateFrom, MergeReplayStateFrom) repoint them under the
-	// write lock while the consume path holds a read lock, so a consume can
-	// never land in a store the handoff already drained and abandoned.
-	// It is always acquired before jtiSweepMu, never while holding it.
-	jtiStoreMu sync.RWMutex
-
-	// jtiMu guards the lifecycle of the cleanup worker: jtiCleanupCancel and
-	// jtiCleanupDone. The worker itself only touches jtiStore/jtiCount, so
-	// holding jtiMu while waiting for it is safe.
-	jtiMu            sync.Mutex
-	jtiCleanupCancel context.CancelFunc
-	jtiCleanupDone   chan struct{} // Closed when cleanup goroutine exits
-
-	// jtiSweepMu serializes cache sweeps so at most one scan runs at a time.
-	jtiSweepMu sync.Mutex
+// replayCache returns the consumed-token cache, allocating it for providers
+// built without NewOIDCProvider so the delegated operations never panic on a
+// nil cache.
+func (p *OIDCProvider) replayCache() *jtiReplayCache {
+	p.jtiOnce.Do(func() {
+		if p.jti == nil {
+			p.jti = newJTIReplayCache(p.name)
+		}
+	})
+	return p.jti
 }
 
 // OIDCConfig holds OIDC provider configuration
@@ -391,8 +358,7 @@ func NewOIDCProvider(name string) *OIDCProvider {
 	return &OIDCProvider{
 		name:       name,
 		httpClient: &http.Client{Timeout: 30 * time.Second, CheckRedirect: refuseDowngradeRedirect},
-		jtiStore:   &sync.Map{},
-		jtiCount:   &atomic.Int64{},
+		jti:        newJTIReplayCache(name),
 	}
 }
 
@@ -512,53 +478,12 @@ func (p *OIDCProvider) Initialize(config interface{}) error {
 	if oidcConfig.JTIReplayProtectionEnabled != nil {
 		jtiEnabled = *oidcConfig.JTIReplayProtectionEnabled
 	}
-
-	p.jtiMu.Lock()
-	p.stopJTICleanupWorkerLocked()
-	if jtiEnabled {
-		maxSize := int64(oidcConfig.JTICacheMaxSize)
-		if maxSize <= 0 {
-			maxSize = defaultJTICacheSize
-		}
-		// the bound must land before the flag: a consume that sees enabled
-		// with a zero bound would reject the first token it registers
-		p.jtiMaxSize.Store(maxSize)
-		ctx, cancel := context.WithCancel(context.Background())
-		done := make(chan struct{})
-		p.jtiCleanupCancel = cancel
-		p.jtiCleanupDone = done
-		p.jtiEnabled.Store(true)
-		p.jtiMu.Unlock()
-
-		go p.cleanupExpiredJTIs(ctx, done, defaultCleanupInterval)
-		glog.V(2).Infof("OIDC provider %q: JTI replay protection enabled (max size: %d)", p.name, maxSize)
-	} else {
-		p.jtiEnabled.Store(false)
-		p.jtiMu.Unlock()
-		glog.Warningf("OIDC provider %q: JTI replay protection is DISABLED", p.name)
-	}
+	p.replayCache().configure(jtiEnabled, int64(oidcConfig.JTICacheMaxSize))
 
 	// For testing, we'll skip the actual OIDC client initialization
 	return nil
 }
 
-// stopJTICleanupWorkerLocked cancels the current cleanup goroutine and waits
-// briefly for it to exit. Caller must hold jtiMu.
-func (p *OIDCProvider) stopJTICleanupWorkerLocked() {
-	if p.jtiCleanupCancel == nil {
-		return
-	}
-	p.jtiCleanupCancel()
-	select {
-	case <-p.jtiCleanupDone:
-		glog.V(3).Infof("OIDC provider %q: JTI cleanup goroutine exited", p.name)
-	case <-time.After(2 * time.Second):
-		glog.Warningf("OIDC provider %q: Timeout waiting for JTI cleanup goroutine", p.name)
-	}
-	p.jtiCleanupCancel = nil
-	p.jtiCleanupDone = nil
-}
-
 // validateConfig validates the OIDC configuration
 func (p *OIDCProvider) validateConfig(config *OIDCConfig) error {
 	if config.Issuer == "" {
@@ -935,85 +860,14 @@ func (p *OIDCProvider) validateToken(ctx context.Context, token string, singleUs
 		tokenClaims.Claims[key] = value
 	}
 
-	// Extract and validate JTI claim for replay protection. The check runs
-	// only for single-use callers: a bearer token is meant to be presented on
-	// every request, so validating it must not consume it.
-	if singleUse && p.jtiEnabled.Load() {
+	// Consume the token for replay protection. The check runs only for
+	// single-use callers: a bearer token is meant to be presented on every
+	// request, so validating it must not consume it.
+	if singleUse {
 		jti, _ := claims["jti"].(string)
-		key := jti
-		if key == "" {
-			// Providers that issue no jti cannot be tracked by claim, so the
-			// signed input keys the entry: a repeated presentation of the same
-			// claims is a replay even when the signature bytes differ (e.g.
-			// ECDSA s/N-s malleability).
-			signingInput := token
-			if i := strings.LastIndexByte(token, '.'); i > 0 {
-				signingInput = token[:i]
-			}
-			sum := sha256.Sum256([]byte(signingInput))
-			key = tokenFingerprintPrefix + hex.EncodeToString(sum[:])
+		if err := p.replayCache().consume(token, jti, subject, tokenClaims.ExpiresAt); err != nil {
+			return nil, err
 		}
-
-		expiresAt := tokenClaims.ExpiresAt
-		if expiresAt.IsZero() {
-			// a cache slot is freed only by expiry, so a token without exp
-			// would hold it forever and eventually exhaust the bound
-			return nil, fmt.Errorf("%w: missing exp claim", providers.ErrProviderMissingClaims)
-		}
-		expiresAt = expiresAt.Add(clockSkewTolerance)
-		if !time.Now().Before(expiresAt) {
-			return nil, fmt.Errorf("%w: token expired", providers.ErrProviderTokenExpired)
-		}
-
-		entry := &jtiEntry{
-			subject:   subject,
-			expiresAt: expiresAt,
-		}
-
-		// The read lock pins the store/count pair for the whole test-and-set
-		// so a replay-state handoff cannot repoint them mid-check and strand
-		// this consume in an abandoned store.
-		p.jtiStoreMu.RLock()
-		store, count := p.jtiStore, p.jtiCount
-
-		// Atomic test-and-set detects replays before any capacity work runs,
-		// so a replayed token cannot trigger cleanup scans.
-		if _, loaded := store.LoadOrStore(key, entry); loaded {
-			p.jtiStoreMu.RUnlock()
-			glog.Warningf("OIDC provider %q: Token replay detected (jti: %s)",
-				p.name, jtiLogKey(key))
-			return nil, fmt.Errorf("%w: jti %s has already been used",
-				providers.ErrProviderTokenReplayed, jtiLogKey(key))
-		}
-
-		// The count is incremented as part of the insert, so concurrent
-		// validations cannot all slip under the limit: whoever pushes the
-		// count over evicts their own entry and is rejected.
-		newCount := count.Add(1)
-		maxSize := p.jtiMaxSize.Load()
-		if newCount > maxSize {
-			// Over the bound: expired entries may still occupy slots, so
-			// reclaim them with a synchronized sweep before rejecting.
-			p.sweepJTIStore(store, count)
-			if count.Load() > maxSize {
-				if _, loaded := store.LoadAndDelete(key); loaded {
-					count.Add(-1)
-				}
-				p.jtiStoreMu.RUnlock()
-				glog.Errorf("OIDC provider %q: JTI cache FULL (%d entries) - rejecting token. "+
-					"This may indicate a DoS attack or insufficient cache size. Consider increasing JTICacheMaxSize.",
-					p.name, maxSize)
-				return nil, fmt.Errorf("JTI cache capacity exceeded")
-			}
-		} else if newCount >= maxSize*9/10 {
-			glog.Warningf("OIDC provider %q: JTI cache at %d%% capacity (%d/%d), triggering eager cleanup",
-				p.name, (newCount*100)/maxSize, newCount, maxSize)
-			go p.performEagerCleanup()
-		}
-		p.jtiStoreMu.RUnlock()
-
-		glog.V(4).Infof("OIDC provider %q: Registered JTI %s (expires at %v, cache size: %d)",
-			p.name, jtiLogKey(key), expiresAt, newCount)
 	}
 
 	return tokenClaims, nil
@@ -1445,148 +1299,22 @@ func (p *OIDCProvider) mapUserInfoToIdentity(userInfo map[string]interface{}) *p
 	return identity
 }
 
-// evictExpiredJTIs removes entries whose expiry has passed and returns how
-// many it deleted. Caller must hold jtiSweepMu.
-func (p *OIDCProvider) evictExpiredJTIs(now time.Time, store *sync.Map, count *atomic.Int64) int {
-	removed := 0
-	store.Range(func(key, value interface{}) bool {
-		expiresAt := value.(*jtiEntry).expiresAt
-		if !expiresAt.IsZero() && now.After(expiresAt) {
-			// LoadAndDelete reports whether this call actually removed the
-			// entry, so concurrent deletes cannot decrement the count twice.
-			if v, loaded := store.LoadAndDelete(key); loaded {
-				if !v.(*jtiEntry).uncounted {
-					count.Add(-1)
-				}
-				removed++
-			}
-		}
-		return true
-	})
-	return removed
-}
-
-// replayState returns the provider's current store/count pair.
-func (p *OIDCProvider) replayState() (*sync.Map, *atomic.Int64) {
-	p.jtiStoreMu.RLock()
-	defer p.jtiStoreMu.RUnlock()
-	return p.jtiStore, p.jtiCount
-}
-
-// sweepJTIStore runs one serialized expiry scan over the given store. It only
-// takes jtiSweepMu, so callers already holding jtiStoreMu may use it.
-func (p *OIDCProvider) sweepJTIStore(store *sync.Map, count *atomic.Int64) {
-	p.jtiSweepMu.Lock()
-	defer p.jtiSweepMu.Unlock()
-	p.jtiLastSweep.Store(time.Now().UnixNano())
-	if removed := p.evictExpiredJTIs(time.Now(), store, count); removed > 0 {
-		glog.V(4).Infof("OIDC provider %q: Cleaned up %d expired JTI entries (remaining: %d)",
-			p.name, removed, count.Load())
-	}
-}
-
-// sweepExpiredJTIs runs a scan that is guaranteed to complete before the
-// caller proceeds; the capacity-rejection path uses it so a fresh token is
-// not refused while expired entries still hold slots.
-func (p *OIDCProvider) sweepExpiredJTIs() {
-	store, count := p.replayState()
-	p.sweepJTIStore(store, count)
-}
-
-// performEagerCleanup performs an opportunistic cleanup of expired JTI
-// entries when the cache approaches capacity. At most one sweep runs at a
-// time and no more often than eagerCleanupMinInterval, so a burst of
-// requests near the limit cannot each turn into a full-cache scan.
-func (p *OIDCProvider) performEagerCleanup() {
-	if !p.jtiEnabled.Load() {
-		return
-	}
-	if time.Since(time.Unix(0, p.jtiLastSweep.Load())) < eagerCleanupMinInterval {
-		return
-	}
-	store, count := p.replayState()
-	if !p.jtiSweepMu.TryLock() {
-		return
-	}
-	defer p.jtiSweepMu.Unlock()
-	p.jtiLastSweep.Store(time.Now().UnixNano())
-
-	beforeCount := count.Load()
-	if removed := p.evictExpiredJTIs(time.Now(), store, count); removed > 0 {
-		glog.Infof("OIDC provider %q: Eager cleanup removed %d expired JTI entries (%d -> %d)",
-			p.name, removed, beforeCount, count.Load())
-	} else {
-		glog.Warningf("OIDC provider %q: Eager cleanup found no expired entries to remove. "+
-			"Cache may be filled with valid long-lived tokens. Consider increasing JTICacheMaxSize.",
-			p.name)
-	}
-}
-
-// cleanupExpiredJTIs periodically removes expired JTI entries from the cache.
-// This prevents memory growth by removing JTIs after their tokens have expired.
-// Runs in a background goroutine until ctx is cancelled; done is closed on exit.
-func (p *OIDCProvider) cleanupExpiredJTIs(ctx context.Context, done chan struct{}, interval time.Duration) {
-	defer close(done)
-
-	ticker := time.NewTicker(interval)
-	defer ticker.Stop()
-
-	glog.V(3).Infof("OIDC provider %q: JTI cleanup goroutine started (interval: %v)", p.name, interval)
-
-	for {
-		select {
-		case <-ctx.Done():
-			glog.V(3).Infof("OIDC provider %q: JTI cleanup goroutine stopping", p.name)
-			return
-		case <-ticker.C:
-			p.sweepExpiredJTIs()
-		}
-	}
-}
-
 // Shutdown gracefully stops the JTI cleanup goroutine.
 // Should be called when the provider is being shut down to prevent goroutine leaks.
 func (p *OIDCProvider) Shutdown(ctx context.Context) error {
-	p.jtiMu.Lock()
-	cancel, done := p.jtiCleanupCancel, p.jtiCleanupDone
-	p.jtiMu.Unlock()
-
-	if cancel == nil {
-		return nil
-	}
-	cancel()
-
-	select {
-	case <-done:
-		glog.V(2).Infof("OIDC provider %q: Cleanup goroutine shutdown complete", p.name)
-	case <-ctx.Done():
-		glog.Warningf("OIDC provider %q: Shutdown timeout waiting for cleanup goroutine: %v", p.name, ctx.Err())
-		return ctx.Err()
-	}
-
-	p.jtiMu.Lock()
-	if p.jtiCleanupDone == done {
-		p.jtiCleanupCancel = nil
-		p.jtiCleanupDone = nil
-	}
-	p.jtiMu.Unlock()
-	return nil
+	return p.replayCache().stop(ctx)
 }
 
 // AdoptReplayStateFrom takes over the consumed-token store of the provider it
 // replaces, so tokens it already saw stay rejected across the rebuild. The
 // caller must shut down the replaced provider; only one cleanup worker may
 // sweep the shared store at a time. Call before the new provider serves
-// validators; jtiSweepMu excludes its own cleanup worker during the swap.
+// validators.
 func (p *OIDCProvider) AdoptReplayStateFrom(old *OIDCProvider) {
-	if old == nil || old == p {
+	if old == nil || old == p || old.jti == nil {
 		return
 	}
-	p.jtiStoreMu.Lock()
-	p.jtiSweepMu.Lock()
-	p.jtiStore, p.jtiCount = old.jtiStore, old.jtiCount
-	p.jtiSweepMu.Unlock()
-	p.jtiStoreMu.Unlock()
+	p.replayCache().adoptFrom(old.jti)
 }
 
 // MergeReplayStateFrom hands this provider's store to another provider and
@@ -1601,38 +1329,5 @@ func (p *OIDCProvider) MergeReplayStateFrom(other *OIDCProvider) {
 	if other == nil || other == p {
 		return
 	}
-	p.jtiStoreMu.RLock()
-	defer p.jtiStoreMu.RUnlock()
-	store, count := p.jtiStore, p.jtiCount
-
-	other.jtiStoreMu.Lock()
-	oldStore := other.jtiStore
-	if oldStore == store {
-		other.jtiStoreMu.Unlock()
-		return
-	}
-	other.jtiStore, other.jtiCount = store, count
-	other.jtiStoreMu.Unlock()
-
-	p.jtiSweepMu.Lock()
-	defer p.jtiSweepMu.Unlock()
-	now := time.Now()
-	oldStore.Range(func(key, value interface{}) bool {
-		entry := value.(*jtiEntry)
-		if expiresAt := entry.expiresAt; !expiresAt.IsZero() && now.After(expiresAt) {
-			return true
-		}
-		// A consumed token is never dropped to fit the bound: merged entries
-		// stay but do not charge the count, so a store filled past the bound
-		// by the merge cannot reject new exchanges while it drains by expiry.
-		store.LoadOrStore(key, &jtiEntry{subject: entry.subject, expiresAt: entry.expiresAt, uncounted: true})
-		return true
-	})
-}
-
-// jtiLogKey redacts a token identifier for logs and error strings: a stable
-// fingerprint that still lets operators correlate repeated presentations.
-func jtiLogKey(jti string) string {
-	sum := sha256.Sum256([]byte(jti))
-	return hex.EncodeToString(sum[:8])
+	p.replayCache().mergeFrom(other.replayCache())
 }

weed/iam/oidc/oidc_provider_test.go

@@ -1067,10 +1067,10 @@ func TestOIDCProviderJTICleanup(t *testing.T) {
 	require.NoError(t, err)
 
 	// Check cache size is 1
-	assert.Equal(t, int64(1), provider.jtiCount.Load())
+	assert.Equal(t, int64(1), provider.jti.count.Load())
 
 	// Verify entry exists and check its expiration includes clock skew
-	val, ok := provider.jtiStore.Load("cleanup-test-jti")
+	val, ok := provider.jti.store.Load("cleanup-test-jti")
 	require.True(t, ok, "JTI entry should exist")
 	entry := val.(*jtiEntry)
 
@@ -1088,21 +1088,21 @@ func TestOIDCProviderJTICleanup(t *testing.T) {
 	// both the entry and the count must actually be cleared
 	cleanupCtx, cleanupCancel := context.WithCancel(context.Background())
 	cleanupDone := make(chan struct{})
-	go provider.cleanupExpiredJTIs(cleanupCtx, cleanupDone, 20*time.Millisecond)
+	go provider.jti.cleanupLoop(cleanupCtx, cleanupDone, 20*time.Millisecond)
 	defer func() {
 		cleanupCancel()
 		<-cleanupDone
 	}()
 
-	provider.jtiStore.Store("expired-cleanup-jti", &jtiEntry{
+	provider.jti.store.Store("expired-cleanup-jti", &jtiEntry{
 		subject:   "expired-user",
 		expiresAt: time.Now().Add(-time.Minute),
 	})
-	provider.jtiCount.Add(1)
+	provider.jti.count.Add(1)
 
 	require.Eventually(t, func() bool {
-		_, stillThere := provider.jtiStore.Load("expired-cleanup-jti")
-		return !stillThere && provider.jtiCount.Load() == 1 // only the valid entry remains
+		_, stillThere := provider.jti.store.Load("expired-cleanup-jti")
+		return !stillThere && provider.jti.count.Load() == 1 // only the valid entry remains
 	}, 2*time.Second, 10*time.Millisecond, "periodic cleanup should evict the expired entry")
 }
 
@@ -1124,7 +1124,7 @@ func TestOIDCProviderReinitialization(t *testing.T) {
 
 	// First initialization
 	require.NoError(t, provider.Initialize(config))
-	require.NotNil(t, provider.jtiCleanupCancel, "First cleanup cancel should be set")
+	require.NotNil(t, provider.jti.cleanupCancel, "First cleanup cancel should be set")
 
 	// Create and validate a token to ensure JTI protection is working
 	token1 := createTestJWT(t, privateKey, jwt.MapClaims{
@@ -1140,7 +1140,7 @@ func TestOIDCProviderReinitialization(t *testing.T) {
 
 	// Second initialization should not cause errors
 	require.NoError(t, provider.Initialize(config))
-	require.NotNil(t, provider.jtiCleanupCancel, "Second cleanup cancel should be set")
+	require.NotNil(t, provider.jti.cleanupCancel, "Second cleanup cancel should be set")
 
 	// The JTI cache survives re-initialization, so the consumed token is
 	// still detected as a replay rather than being usable again
@@ -1150,7 +1150,7 @@ func TestOIDCProviderReinitialization(t *testing.T) {
 
 	// Third initialization to ensure it works multiple times
 	require.NoError(t, provider.Initialize(config))
-	require.NotNil(t, provider.jtiCleanupCancel, "Third cleanup cancel should be set")
+	require.NotNil(t, provider.jti.cleanupCancel, "Third cleanup cancel should be set")
 
 	// Verify JTI replay protection still works after re-initialization
 	token2 := createTestJWT(t, privateKey, jwt.MapClaims{
@@ -1201,16 +1201,16 @@ func TestOIDCProviderJTIEagerCleanup(t *testing.T) {
 			subject:   fmt.Sprintf("expired-user-%d", i),
 			expiresAt: time.Now().Add(-time.Hour), // Already expired
 		}
-		provider.jtiStore.Store(jti, entry)
-		provider.jtiCount.Add(1)
+		provider.jti.store.Store(jti, entry)
+		provider.jti.count.Add(1)
 	}
 
-	assert.Equal(t, int64(5), provider.jtiCount.Load(), "Should have 5 expired entries")
+	assert.Equal(t, int64(5), provider.jti.count.Load(), "Should have 5 expired entries")
 
 	// Perform eager cleanup - should remove all expired entries
-	provider.performEagerCleanup()
+	provider.jti.performEagerCleanup()
 
-	assert.Equal(t, int64(0), provider.jtiCount.Load(),
+	assert.Equal(t, int64(0), provider.jti.count.Load(),
 		"All expired entries should be removed by eager cleanup")
 
 	// Now test that adding valid tokens works after cleanup
@@ -1228,7 +1228,7 @@ func TestOIDCProviderJTIEagerCleanup(t *testing.T) {
 	}
 
 	// Should have 9 valid entries now (90% of capacity, triggers eager cleanup warning)
-	assert.Equal(t, int64(9), provider.jtiCount.Load(), "Should have 9 valid entries")
+	assert.Equal(t, int64(9), provider.jti.count.Load(), "Should have 9 valid entries")
 
 	// 10th token should succeed (at capacity but not over)
 	token10 := createTestJWT(t, privateKey, jwt.MapClaims{
@@ -1282,11 +1282,11 @@ func TestOIDCProviderConcurrentCleanupNoNegativeCount(t *testing.T) {
 			subject:   fmt.Sprintf("user-%d", i),
 			expiresAt: time.Now().Add(-time.Hour), // Already expired
 		}
-		provider.jtiStore.Store(jti, entry)
-		provider.jtiCount.Add(1)
+		provider.jti.store.Store(jti, entry)
+		provider.jti.count.Add(1)
 	}
 
-	initialCount := provider.jtiCount.Load()
+	initialCount := provider.jti.count.Load()
 	assert.Equal(t, int64(50), initialCount, "Should have 50 expired entries")
 
 	// Run multiple concurrent cleanups
@@ -1297,14 +1297,14 @@ func TestOIDCProviderConcurrentCleanupNoNegativeCount(t *testing.T) {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			provider.performEagerCleanup()
+			provider.jti.performEagerCleanup()
 		}()
 	}
 
 	wg.Wait()
 
 	// Counter should be 0 (all entries removed), never negative
-	finalCount := provider.jtiCount.Load()
+	finalCount := provider.jti.count.Load()
 	assert.Equal(t, int64(0), finalCount,
 		"Counter should be 0 after concurrent cleanups, got %d", finalCount)
 	assert.GreaterOrEqual(t, finalCount, int64(0),
@@ -1312,7 +1312,7 @@ func TestOIDCProviderConcurrentCleanupNoNegativeCount(t *testing.T) {
 
 	// Verify store is actually empty
 	storeCount := 0
-	provider.jtiStore.Range(func(key, value interface{}) bool {
+	provider.jti.store.Range(func(key, value interface{}) bool {
 		storeCount++
 		return true
 	})