From 8844d9823f561362025bf42a57c53ebe47bcd924 Mon Sep 17 00:00:00 2001 From: Chris Lu Date: Sat, 24 Oct 2026 07:03:42 -0700 Subject: [PATCH] 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 | 410 ++++++++++++++++++++++++++++ weed/iam/oidc/oidc_provider.go | 365 ++----------------------- weed/iam/oidc/oidc_provider_test.go | 44 +-- 3 files changed, 462 insertions(+), 357 deletions(-) create mode 100644 weed/iam/oidc/jti_cache.go diff --git a/weed/iam/oidc/jti_cache.go b/weed/iam/oidc/jti_cache.go new file mode 100644 index 0000000000..77dab0d044 --- /dev/null +++ b/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]) +} diff --git a/weed/iam/oidc/oidc_provider.go b/weed/iam/oidc/oidc_provider.go index c2ccf4ab8f..6b115041d5 100644 --- a/weed/iam/oidc/oidc_provider.go +++ b/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()) } diff --git a/weed/iam/oidc/oidc_provider_test.go b/weed/iam/oidc/oidc_provider_test.go index 3744ec3138..29f7e3e656 100644 --- a/weed/iam/oidc/oidc_provider_test.go +++ b/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 })