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.
This commit is contained in:
Joseph Doherty
2026-10-22 06:33:18 -04:00
parent 68d6b5fd6f
commit ed3a2c9c8b
3 changed files with 295 additions and 16 deletions
+10
View File
@@ -113,6 +113,16 @@ library:
library returns carries no expiry timestamp, so the cache cannot cap an entry at the
key's expiry (capping it needs the donor library to surface expiry on the
verification identity). The default 15 s TTL bounds both windows.
Misses are **single-flight**: concurrent misses for the same token share one inner
verification, so a TTL expiry under load costs one store read instead of one per
in-flight RPC. The flight is keyed on the token hash *and* the key's generation, so
a miss that starts after an `Invalidate` never joins a verification that started
before it and cannot receive the pre-revoke identity. The shared inner call is not
bound to any one caller's cancellation token; a caller that gives up stops waiting
without cancelling the others. The per-key index `Invalidate` uses to find every
cached secret variant is pruned by a post-eviction callback when the cache drops an
entry on its own (TTL expiry, compaction), so it tracks the live cache instead of
growing with every token ever verified.
- **`CoalescingMarkApiKeyStore`** wraps the library `IApiKeyStore` and forwards at
most one `MarkUsed` write per key per
`MxGateway:Security:ApiKeyLastUsedCoalesceSeconds` (default 60 s), so even under a
@@ -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)
@@ -129,6 +129,120 @@ public sealed class CachingApiKeyVerifierTests
Assert.Equal(2, inner.CallCount);
}
/// <summary>
/// Single-flight: concurrent misses for the same token (a TTL expiry under load) share one
/// inner verification instead of each issuing its own SQLite read.
/// </summary>
/// <returns>A task that represents the asynchronous operation.</returns>
[Fact]
public async Task VerifyAsync_ConcurrentMisses_ShareOneInnerVerification()
{
GatedVerifier inner = new(Success("operator01"));
using MemoryCache cache = NewCache();
CachingApiKeyVerifier verifier = new(inner, cache, TimeSpan.FromSeconds(30));
Task<ApiKeyVerification> first = verifier.VerifyAsync(Header, CancellationToken.None);
await inner.Entered;
Task<ApiKeyVerification>[] joiners = Enumerable
.Range(0, 16)
.Select(_ => verifier.VerifyAsync(Header, CancellationToken.None))
.ToArray();
inner.Release();
ApiKeyVerification[] results = await Task.WhenAll([first, .. joiners]);
Assert.All(results, result => Assert.True(result.Succeeded));
Assert.Equal(1, inner.CallCount);
}
/// <summary>
/// A miss that arrives after <see cref="IApiKeyCacheInvalidator.Invalidate"/> must not join a
/// verification that started before it: that flight may return the pre-revoke identity, so the
/// late caller starts its own flight against the inner verifier.
/// </summary>
/// <returns>A task that represents the asynchronous operation.</returns>
[Fact]
public async Task VerifyAsync_MissAfterInvalidate_DoesNotJoinPreInvalidateFlight()
{
GatedVerifier inner = new(Success("operator01"));
using MemoryCache cache = NewCache();
CachingApiKeyVerifier verifier = new(inner, cache, TimeSpan.FromSeconds(30));
Task<ApiKeyVerification> preRevoke = verifier.VerifyAsync(Header, CancellationToken.None);
await inner.Entered;
((IApiKeyCacheInvalidator)verifier).Invalidate("operator01");
// The gated first call is still parked; a joined call would park with it.
ApiKeyVerification postRevoke = await verifier
.VerifyAsync(Header, CancellationToken.None)
.WaitAsync(TimeSpan.FromSeconds(5));
Assert.True(postRevoke.Succeeded);
Assert.Equal(2, inner.CallCount);
inner.Release();
await preRevoke;
}
/// <summary>
/// A caller whose own token is cancelled stops waiting, but the shared verification it joined
/// is not cancelled for the other waiters.
/// </summary>
/// <returns>A task that represents the asynchronous operation.</returns>
[Fact]
public async Task VerifyAsync_CancelledWaiter_DoesNotCancelSharedVerification()
{
GatedVerifier inner = new(Success("operator01"));
using MemoryCache cache = NewCache();
CachingApiKeyVerifier verifier = new(inner, cache, TimeSpan.FromSeconds(30));
using CancellationTokenSource cancelled = new();
Task<ApiKeyVerification> survivor = verifier.VerifyAsync(Header, CancellationToken.None);
await inner.Entered;
Task<ApiKeyVerification> abandoned = verifier.VerifyAsync(Header, cancelled.Token);
await cancelled.CancelAsync();
await Assert.ThrowsAnyAsync<OperationCanceledException>(() => abandoned);
inner.Release();
Assert.True((await survivor).Succeeded);
Assert.Equal(1, inner.CallCount);
}
/// <summary>
/// The key-id index is pruned when the cache evicts an entry on its own (expiry, compaction),
/// not only on <see cref="IApiKeyCacheInvalidator.Invalidate"/> — otherwise every token ever
/// verified leaves an index entry behind for the life of the process.
/// </summary>
/// <returns>A task that represents the asynchronous operation.</returns>
[Fact]
public async Task CacheEviction_PrunesKeyIdIndex()
{
FakeVerifier inner = new(Success("operator01"));
using MemoryCache cache = NewCache();
CachingApiKeyVerifier verifier = new(inner, cache, TimeSpan.FromSeconds(30));
await verifier.VerifyAsync(Header, CancellationToken.None);
Assert.Equal(1, verifier.IndexedCacheKeyCount);
cache.Compact(1.0);
DateTime deadline = DateTime.UtcNow + TimeSpan.FromSeconds(5);
while (verifier.IndexedCacheKeyCount != 0 && DateTime.UtcNow < deadline)
{
await Task.Delay(10);
}
Assert.Equal(0, verifier.IndexedCacheKeyCount);
Assert.Equal(0, verifier.IndexedKeyIdCount);
// The evicted entry is gone, so the next request re-verifies, and Invalidate still reaches
// the repopulated entry through the index.
await verifier.VerifyAsync(Header, CancellationToken.None);
Assert.Equal(2, inner.CallCount);
((IApiKeyCacheInvalidator)verifier).Invalidate("operator01");
await verifier.VerifyAsync(Header, CancellationToken.None);
Assert.Equal(3, inner.CallCount);
}
/// <summary>
/// The store decorator coalesces repeated <c>MarkUsed</c> writes for the same key inside the
/// window down to a single forwarded write — the ≤1/min guarantee for <c>last_used_utc</c>.