fix(auth): single-flight API-key verification misses; prune key-id index on eviction

CachingApiKeyVerifier's key-id index was only pruned by Invalidate, so it grew with every token ever verified. A post-eviction callback now drops the slot (under the same lock as Set + index, so a late callback cannot unindex a re-cached entry). Concurrent misses for the same token and key generation now share one inner verification, so a TTL expiry under load no longer stampedes SQLite; keying the flight on the generation keeps the revoke guarantee for joiners.
2026-10-22 06:33:18 -04:00
parent 68d6b5fd6f
commit ed3a2c9c8b
3 changed files with 295 additions and 16 deletions
@@ -58,6 +58,19 @@ public interface IApiKeyCacheInvalidator
 /// window and it is intentionally kept short (<see cref="SecurityOptions.ApiKeyVerificationCacheSeconds"/>,
 /// default 15 s).
 /// </para>
+/// <para>
+/// Misses are single-flight: concurrent misses for the same token and the same key generation share
+/// one inner verification, so a TTL expiry under load costs one SQLite read rather than one per
+/// in-flight RPC. A miss that starts after an <see cref="Invalidate"/> never joins a flight that
+/// started before it (the flight is keyed on the generation too), so the revoke guarantee holds for
+/// joiners as well. The shared inner call is not bound to any one caller's cancellation token; each
+/// caller stops waiting on its own token without cancelling the others.
+/// </para>
+/// <para>
+/// The key-id index is pruned by a post-eviction callback whenever the cache drops an entry on its
+/// own (expiry, compaction), not only by <see cref="Invalidate"/>, so it stays proportional to the
+/// live cache rather than to every token ever verified.
+/// </para>
 /// </remarks>
 public sealed class CachingApiKeyVerifier : IApiKeyVerifier, IApiKeyCacheInvalidator
 {
@@ -69,9 +82,20 @@ public sealed class CachingApiKeyVerifier : IApiKeyVerifier, IApiKeyCacheInvalid
    private readonly TimeSpan _ttl;

    // keyId -> set of cache keys currently held for that id, so a revoke/rotate can evict every
-    // secret variant (e.g. an old secret still within TTL after a rotate).
-    private readonly ConcurrentDictionary<string, ConcurrentDictionary<string, byte>> _keyIdIndex =
-        new(StringComparer.Ordinal);
+    // secret variant (e.g. an old secret still within TTL after a rotate). Guarded by _indexSync,
+    // which also covers the cache Set that an index entry describes: the post-eviction callback
+    // runs asynchronously, so Set + index and check-cache + unindex must each be atomic with
+    // respect to the other, or a late callback could unindex a freshly re-cached entry and leave
+    // it out of reach of a later Invalidate. Every operation under the lock is O(1) and the lock is
+    // only taken on misses, evictions and invalidations — never on a cache hit.
+    private readonly object _indexSync = new();
+    private readonly Dictionary<string, HashSet<string>> _keyIdIndex = new(StringComparer.Ordinal);
+
+    // (cache key, key generation) -> the in-flight inner verification for that pair (single-flight).
+    private readonly ConcurrentDictionary<(string CacheKey, long Generation), Lazy<Task<ApiKeyVerification>>> _inflight =
+        new();
+
+    private readonly PostEvictionDelegate _onEvicted;

    // keyId -> monotonic generation, bumped by Invalidate before it evicts. A VerifyAsync snapshots
    // the generation before calling the inner verifier and only caches if it is unchanged after the
@@ -103,6 +127,37 @@ public sealed class CachingApiKeyVerifier : IApiKeyVerifier, IApiKeyCacheInvalid
        _inner = inner;
        _cache = cache;
        _ttl = ttl;
+        _onEvicted = OnCacheEntryEvicted;
+    }
+
+    /// <summary>Gets the number of cache keys currently held in the key-id index (test seam).</summary>
+    internal int IndexedCacheKeyCount
+    {
+        get
+        {
+            lock (_indexSync)
+            {
+                int count = 0;
+                foreach (HashSet<string> set in _keyIdIndex.Values)
+                {
+                    count += set.Count;
+                }
+
+                return count;
+            }
+        }
+    }
+
+    /// <summary>Gets the number of key ids currently present in the key-id index (test seam).</summary>
+    internal int IndexedKeyIdCount
+    {
+        get
+        {
+            lock (_indexSync)
+            {
+                return _keyIdIndex.Count;
+            }
+        }
    }

    /// <summary>
@@ -126,11 +181,50 @@ public sealed class CachingApiKeyVerifier : IApiKeyVerifier, IApiKeyCacheInvalid
        }

        // Snapshot the key's generation BEFORE the inner verify so a concurrent Invalidate that runs
-        // while we are in the library is detected below and its repopulation discarded.
+        // while we are in the library is detected and its repopulation discarded. The snapshot is
+        // also half of the single-flight key, so a miss that starts after an Invalidate cannot join
+        // a flight that started before it.
        string? tokenKeyId = TryParseKeyId(authorizationHeader);
        long generationAtStart = tokenKeyId is null ? 0 : ReadGeneration(tokenKeyId);
+        (string CacheKey, long Generation) flightKey = (cacheKey, generationAtStart);

-        ApiKeyVerification result = await _inner.VerifyAsync(authorizationHeader, ct).ConfigureAwait(false);
+        Lazy<Task<ApiKeyVerification>> created = new(
+            () => VerifyMissAsync(authorizationHeader, cacheKey, tokenKeyId, generationAtStart));
+        Lazy<Task<ApiKeyVerification>> flight = _inflight.GetOrAdd(flightKey, created);
+        Task<ApiKeyVerification> verification = flight.Value;
+        if (ReferenceEquals(flight, created))
+        {
+            // The flight's creator retires it once it completes (success, failure or fault), so the
+            // next miss for this token starts fresh. Removing by (key, value) pair never retires a
+            // newer flight that replaced this one.
+            _ = verification.ContinueWith(
+                static (_, state) =>
+                {
+                    var (owner, key, completed) =
+                        ((CachingApiKeyVerifier, (string, long), Lazy<Task<ApiKeyVerification>>))state!;
+                    owner._inflight.TryRemove(
+                        new KeyValuePair<(string CacheKey, long Generation), Lazy<Task<ApiKeyVerification>>>(key, completed));
+                },
+                (this, flightKey, flight),
+                CancellationToken.None,
+                TaskContinuationOptions.ExecuteSynchronously,
+                TaskScheduler.Default);
+        }
+
+        return await verification.WaitAsync(ct).ConfigureAwait(false);
+    }
+
+    // The shared body of one single-flight miss. Not bound to any caller's token: a caller that
+    // gives up must not cancel the verification other callers are waiting on.
+    private async Task<ApiKeyVerification> VerifyMissAsync(
+        string authorizationHeader,
+        string cacheKey,
+        string? tokenKeyId,
+        long generationAtStart)
+    {
+        ApiKeyVerification result = await _inner
+            .VerifyAsync(authorizationHeader, CancellationToken.None)
+            .ConfigureAwait(false);

        if (result.Succeeded && result.Identity is not null)
        {
@@ -158,11 +252,21 @@ public sealed class CachingApiKeyVerifier : IApiKeyVerifier, IApiKeyCacheInvalid
        }

        string keyId = result.Identity!.KeyId;
-        _cache.Set(cacheKey, result, new MemoryCacheEntryOptions
+        MemoryCacheEntryOptions entryOptions = new()
        {
            AbsoluteExpirationRelativeToNow = _ttl,
+        };
+        entryOptions.PostEvictionCallbacks.Add(new PostEvictionCallbackRegistration
+        {
+            EvictionCallback = _onEvicted,
+            State = keyId,
        });
-        IndexCacheKey(keyId, cacheKey);
+
+        lock (_indexSync)
+        {
+            _cache.Set(cacheKey, result, entryOptions);
+            IndexCacheKey(keyId, cacheKey);
+        }

        // Set-then-recheck: an Invalidate that landed between the pre-Set check and IndexCacheKey
        // would have missed this entry (it was not yet indexed). If the generation has moved, evict
@@ -170,13 +274,38 @@ public sealed class CachingApiKeyVerifier : IApiKeyVerifier, IApiKeyCacheInvalid
        if (tokenKeyId is not null && ReadGeneration(tokenKeyId) != generationAtStart)
        {
            _cache.Remove(cacheKey);
-            if (_keyIdIndex.TryGetValue(keyId, out ConcurrentDictionary<string, byte>? set))
+            lock (_indexSync)
            {
-                set.TryRemove(cacheKey, out _);
+                UnindexCacheKey(keyId, cacheKey);
            }
        }
    }

+    // Post-eviction callback: drops the evicted entry's key-id index slot so the index tracks the
+    // live cache. MemoryCache raises it asynchronously, so by the time it runs the same cache key
+    // may already hold a newer entry (re-verified after the expiry); that entry owns the slot, and
+    // the check below runs under the lock that also covers Set + index, so it cannot interleave
+    // with a re-cache.
+    private void OnCacheEntryEvicted(object key, object? value, EvictionReason reason, object? state)
+    {
+        if (reason == EvictionReason.Replaced || key is not string cacheKey || state is not string keyId)
+        {
+            // Replaced: a newer entry under the same key was Set in this entry's place and is
+            // indexed under the same slot already.
+            return;
+        }
+
+        lock (_indexSync)
+        {
+            if (_cache.TryGetValue(cacheKey, out _))
+            {
+                return;
+            }
+
+            UnindexCacheKey(keyId, cacheKey);
+        }
+    }
+
    /// <inheritdoc />
    public void Invalidate(string keyId)
    {
@@ -189,12 +318,22 @@ public sealed class CachingApiKeyVerifier : IApiKeyVerifier, IApiKeyCacheInvalid
        // generation refuses to cache (or self-evicts) its now-stale repopulation.
        _generations.AddOrUpdate(keyId, 1, static (_, current) => current + 1);

-        if (_keyIdIndex.TryRemove(keyId, out ConcurrentDictionary<string, byte>? cacheKeys))
+        string[] cacheKeys;
+        lock (_indexSync)
        {
-            foreach (string cacheKey in cacheKeys.Keys)
+            if (!_keyIdIndex.Remove(keyId, out HashSet<string>? indexed))
            {
-                _cache.Remove(cacheKey);
+                return;
            }
+
+            cacheKeys = [.. indexed];
+        }
+
+        // Evicted outside the lock: the eviction callbacks these raise take it themselves (and find
+        // the slot already gone).
+        foreach (string cacheKey in cacheKeys)
+        {
+            _cache.Remove(cacheKey);
        }
    }

@@ -257,12 +396,28 @@ public sealed class CachingApiKeyVerifier : IApiKeyVerifier, IApiKeyCacheInvalid
        return new string(afterPrefix[..separator]);
    }

+    // Caller holds _indexSync.
    private void IndexCacheKey(string keyId, string cacheKey)
    {
-        ConcurrentDictionary<string, byte> set = _keyIdIndex.GetOrAdd(
-            keyId,
-            static _ => new ConcurrentDictionary<string, byte>(StringComparer.Ordinal));
-        set.TryAdd(cacheKey, 0);
+        if (!_keyIdIndex.TryGetValue(keyId, out HashSet<string>? set))
+        {
+            set = new HashSet<string>(StringComparer.Ordinal);
+            _keyIdIndex[keyId] = set;
+        }
+
+        set.Add(cacheKey);
+    }
+
+    // Caller holds _indexSync. Drops the key id's slot entirely once its last cache key goes, so
+    // the index does not accumulate empty sets for keys that are no longer cached.
+    private void UnindexCacheKey(string keyId, string cacheKey)
+    {
+        if (_keyIdIndex.TryGetValue(keyId, out HashSet<string>? set)
+            && set.Remove(cacheKey)
+            && set.Count == 0)
+        {
+            _keyIdIndex.Remove(keyId);
+        }
    }

    private static bool TryComputeCacheKey(string? authorizationHeader, out string cacheKey)