diff --git a/jwt.go b/jwt.go index d0ab097..b7588c7 100644 --- a/jwt.go +++ b/jwt.go @@ -365,54 +365,50 @@ func (j *JWT) Verify(issuerURL, expectedAudience string, skipReplayCheck ...bool if jtiOk && !shouldSkipReplay && jtiValue != "" { initReplayCache() - // Use sharded cache for replay detection - no global mutex needed - // This reduces lock contention by ~64x under high load - if shardedReplayCache != nil { - if shardedReplayCache.Exists(jtiValue) { + expFloat, ok := claims["exp"].(float64) + var expTime time.Time + if ok { + expTime = time.Unix(int64(expFloat), 0) + } else { + expTime = time.Now().Add(10 * time.Minute) + } + duration := time.Until(expTime) + if duration <= 0 { + // Expired token: no replay window worth recording, so nothing to do. + goto replayDone + } + + // Use sharded cache for replay detection - no global mutex needed. + // SetIfAbsent is atomic: it holds the per-shard lock across the + // existence check and insert, closing the double-accept race where + // two concurrent requests carrying the same fresh JTI could both + // observe it absent and both be accepted. Guard the singleton read + // with replayCacheMu because cleanupReplayCache can nil it under the + // write lock. + replayCacheMu.RLock() + sc := shardedReplayCache + if sc != nil { + if !sc.SetIfAbsent(jtiValue, true, duration) { + replayCacheMu.RUnlock() return fmt.Errorf("token replay detected (jti: %s)", jtiValue) } - - expFloat, ok := claims["exp"].(float64) - var expTime time.Time - if ok { - expTime = time.Unix(int64(expFloat), 0) - } else { - expTime = time.Now().Add(10 * time.Minute) - } - - duration := time.Until(expTime) - if duration > 0 { - shardedReplayCache.Set(jtiValue, true, duration) - } - } else { - // Fall back to legacy cache with mutex (should rarely happen) - replayCacheMu.RLock() - _, exists := replayCache.Get(jtiValue) replayCacheMu.RUnlock() - + } else { + // Fall back to legacy cache (should rarely happen). Hold the + // mutex across the get+set so the record is atomic there too. + replayCacheMu.Lock() + _, exists := replayCache.Get(jtiValue) + if !exists && replayCache != nil { + replayCache.Set(jtiValue, true, duration) + } + replayCacheMu.Unlock() if exists { return fmt.Errorf("token replay detected (jti: %s)", jtiValue) } - - expFloat, ok := claims["exp"].(float64) - var expTime time.Time - if ok { - expTime = time.Unix(int64(expFloat), 0) - } else { - expTime = time.Now().Add(10 * time.Minute) - } - - duration := time.Until(expTime) - if duration > 0 { - replayCacheMu.Lock() - if replayCache != nil { - replayCache.Set(jtiValue, true, duration) - } - replayCacheMu.Unlock() - } } } +replayDone: sub, ok := claims["sub"].(string) if !ok || sub == "" { return fmt.Errorf("missing or empty 'sub' claim") diff --git a/review_r86_test.go b/review_r86_test.go new file mode 100644 index 0000000..e389770 --- /dev/null +++ b/review_r86_test.go @@ -0,0 +1,71 @@ +package traefikoidc + +import ( + "fmt" + "sync" + "testing" + "time" +) + +// TestReplayCache_SetIfAbsentIsAtomic verifies that SetIfAbsent is an atomic +// check-and-set under the per-shard lock: when many goroutines race to +// record the same key, exactly one wins. This closes the double-accept +// race in JWT replay detection where the old Exists()+Set() check-then-act +// let two concurrent requests carrying the same fresh JTI both observe it +// absent and both be accepted. +func TestReplayCache_SetIfAbsentIsAtomic(t *testing.T) { + c := NewShardedCache(8, 10000) + const ( + key = "shared-jti-abc" + n = 200 + ttl = time.Hour + ) + + start := make(chan struct{}) + var wg sync.WaitGroup + var mu sync.Mutex + winners := 0 + + for i := 0; i < n; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + if c.SetIfAbsent(key, true, ttl) { + mu.Lock() + winners++ + mu.Unlock() + } + }() + } + close(start) + wg.Wait() + + if winners != 1 { + t.Fatalf("expected exactly one SetIfAbsent winner, got %d (double-accept race)", winners) + } + if !c.Exists(key) { + t.Fatal("key should be present after a single winner") + } + + // A later arrival with the same key must still be rejected. + if c.SetIfAbsent(key, true, ttl) { + t.Fatal("SetIfAbsent should return false for an already-present key") + } +} + +// TestReplayCache_SetIfAbsentReplacesExpired verifies that an entry that has +// expired is treated as absent (so a fresh record replaces it rather than +// being spuriously reported as a replay). +func TestReplayCache_SetIfAbsentReplacesExpired(t *testing.T) { + c := NewShardedCache(8, 10000) + key := fmt.Sprintf("exp-jti-%d", time.Now().UnixNano()) + + if !c.SetIfAbsent(key, true, time.Millisecond) { + t.Fatal("first insert should succeed") + } + time.Sleep(10 * time.Millisecond) // let the entry expire + if !c.SetIfAbsent(key, true, time.Hour) { + t.Fatal("second insert of an expired key should succeed (replaces it)") + } +} diff --git a/sharded_cache.go b/sharded_cache.go index 91309b2..093198b 100644 --- a/sharded_cache.go +++ b/sharded_cache.go @@ -122,6 +122,45 @@ func (c *ShardedCache) Set(key string, value interface{}, ttl time.Duration) { shard.mu.Unlock() } +// SetIfAbsent atomically inserts the item only if the key is not already +// present and unexpired. Returns true if inserted, false if the key already +// existed (including when a concurrent caller already inserted it). Unlike +// the separate Exists()+Set() check-then-act, this holds the shard lock +// across both operations, closing the double-accept race in replay +// detection where two concurrent requests carrying the same fresh JTI could +// both observe it absent and both be accepted. +func (c *ShardedCache) SetIfAbsent(key string, value interface{}, ttl time.Duration) bool { + shard := c.getShard(key) + + var expiresAt time.Time + if ttl > 0 { + expiresAt = time.Now().Add(ttl) + } + + shard.mu.Lock() + defer shard.mu.Unlock() + + if it, exists := shard.items[key]; exists { + // Present and not expired -> duplicate. + if it.expiresAt.IsZero() || !time.Now().After(it.expiresAt) { + return false + } + // Present but expired -> treat as absent and replace below. + } + + if len(shard.items) >= c.maxPerShard { + c.evictFromShardLocked(shard) + } + + shard.items[key] = &shardedCacheItem{ + value: value, + expiresAt: expiresAt, + seq: shard.nextSeq, + } + shard.nextSeq++ + return true +} + // Delete removes an item from the cache. func (c *ShardedCache) Delete(key string) { shard := c.getShard(key)