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:
@@ -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>.
|
||||
|
||||
Reference in New Issue
Block a user