R86: atomic SetIfAbsent for replay cache; close double-accept and singleton-read race

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")

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)")
+	}
+}

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)