dohertj2/mxaccessgw
1
0
Fork 0
Code Issues 43 Pull Requests 1 Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
src/ZB.MOM.WW.MxGateway.Server/Security/Authentication/CachingApiKeyVerifier.cs
+171 -16
View File
@@ -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)
src/ZB.MOM.WW.MxGateway.Tests/Security/Authentication/CachingApiKeyVerifierTests.cs
+114
View File
@@ -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>.