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

fix(SEC-31,SEC-32): make probe admission atomic and stop Reset clearing a shared fallback partition

Browse Source 
Two defects found in code review of the limiter rework.

Probe admission was check-then-act across two lock scopes: Check() read
"probe due" under lock(state), released it, then re-acquired to advance
NextProbeAtTicks. A burst of requests arriving together at an interval boundary
could therefore all observe the slot as due and all be admitted, handing the
verifier the very burst the interval exists to bound. The claim is now a single
critical section (TryConsumeProbe). The two layers are still claimed one at a
time — holding two per-state locks at once would need a global lock ordering to
stay deadlock-free — so a slot claimed on the composite partition is compensated
via ReleaseProbe when the aggregate then refuses, which otherwise silently spent
the partition's next slot and pushed the legitimate holder out by a full
interval.

Reset() removed whatever partition the caller resolved to, including the
address's shared fallback partition when the caller's key id had been collapsed
into it by the per-peer cap (or when the token was junk-shaped). That bucket also
carries failures contributed by other key ids from the same address, so one
successful authentication became a reset button for an in-progress spray. Reset
now clears only a partition the caller owns (effectiveKeyId == presented key id);
the shared bucket decays by window expiry instead, and the caller still recovers
through probe admission. The key's aggregate is cleared either way, as designed.

Also applied from the review: closure-free GetOrAdd overload on _partitions, and
a remarks paragraph acknowledging the best-effort O(n) eviction scan under
sustained overflow. Threading the resolved partition key from Check through to
RecordFailure/Reset was declined: Check resolves with mint:false and RecordFailure
with mint:true, and the two can legitimately differ when a concurrent caller fills
the per-peer cap in between — reusing Check's key would record into the wrong
partition and bypass the cap, which is not worth saving one string concat.

Tests (limiter suite 11 -> 14): ProbeAdmission_UnderConcurrentArrivals_
GrantsExactlyOneSlot (200 rounds x 8 barrier-released threads at the boundary),
ProbeAdmission_WhenAggregateRefuses_ReturnsTheClaimedPeerSlot, and
Reset_WithOverCapKeyId_DoesNotClearSharedFallbackPartition. The latter two were
confirmed as genuine reds against the unfixed code; the concurrency test is a
guard — it is deterministically green on the fixed structure but did not
reproduce the original nanosecond-wide window on its own.
This commit is contained in:
Joseph Doherty
2026-08-07 05:57:26 -04:00
parent df710e18a9
commit acebe18773
5 changed files with 196 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ZB.MOM.WW.MxGateway.Server/Security/Authorization/ApiKeyFailureLimiter.cs
+66 -23
View File
@@ -35,6 +35,13 @@ namespace ZB.MOM.WW.MxGateway.Server.Security.Authorization;
/// peers are removed on reset, so in steady state the map holds only partitions with recent failures
/// — the common success path is a lock-free dictionary miss.
/// </para>
/// <para>
/// Eviction is deliberately a best-effort full scan of the over-capacity map, run only on the
/// failure path and only once the cap is exceeded. Under a sustained overflow that scan is O(n) per
/// recorded failure, which is accepted: n is bounded by the cap, the work lands on attack traffic
/// rather than on authenticated calls, and an exact ordered structure would need a second index kept
/// consistent with the per-state locks for no security gain.
/// </para>
/// </remarks>
public sealed class ApiKeyFailureLimiter
{
@@ -140,26 +147,29 @@ public sealed class ApiKeyFailureLimiter
return peerOver ? ApiKeyThrottleDecision.ThrottledByPeer : ApiKeyThrottleDecision.ThrottledByAggregate;
}
// Both over-limit layers must have a slot before either is consumed, so a request cannot burn
// the peer's probe and then be refused by the aggregate.
if (peerOver && !IsProbeDue(peerState!, now))
{
return ApiKeyThrottleDecision.ThrottledByPeer;
}
if (aggregateOver && !IsProbeDue(aggregateState!, now))
{
return ApiKeyThrottleDecision.ThrottledByAggregate;
}
// Every over-limit layer must yield its probe slot for the request to pass. The slots are
// claimed one at a time (holding two per-state locks at once would need a global ordering to
// stay deadlock-free), so a claim is reserved and then compensated if a later layer refuses.
long peerProbeRestore = 0;
bool peerProbeClaimed = false;
if (peerOver)
{
ConsumeProbe(peerState!, now);
if (!TryConsumeProbe(peerState!, now, out peerProbeRestore))
{
return ApiKeyThrottleDecision.ThrottledByPeer;
}
peerProbeClaimed = true;
}
if (aggregateOver)
if (aggregateOver && !TryConsumeProbe(aggregateState!, now, out _))
{
ConsumeProbe(aggregateState!, now);
if (peerProbeClaimed)
{
ReleaseProbe(peerState!, now, peerProbeRestore);
}
return ApiKeyThrottleDecision.ThrottledByAggregate;
}
return ApiKeyThrottleDecision.ProbeAdmitted;
@@ -178,7 +188,10 @@ public sealed class ApiKeyFailureLimiter
long now = _clock.GetUtcNow().UtcTicks;
(string partitionKey, string? effectiveKeyId) = ResolvePartitionKey(peer, partition.KeyId, mint: true);
WindowState state = _partitions.GetOrAdd(partitionKey, _ => new WindowState(peer, effectiveKeyId));
WindowState state = _partitions.GetOrAdd(
partitionKey,
static (_, owner) => new WindowState(owner.Peer, owner.KeyId),
(Peer: peer, KeyId: effectiveKeyId));
RecordInto(state, now, _limit);
// Only a key id that earned its own partition feeds the aggregate: an id squeezed out by the
@@ -198,8 +211,17 @@ public sealed class ApiKeyFailureLimiter
public void Reset(ApiKeyThrottlePartition partition)
{
string peer = RequirePeer(partition);
(string partitionKey, _) = ResolvePartitionKey(peer, partition.KeyId, mint: false);
RemovePartition(partitionKey);
(string partitionKey, string? effectiveKeyId) = ResolvePartitionKey(peer, partition.KeyId, mint: false);
// Clear only a partition this caller actually owns. When its key id was squeezed into the
// address's shared fallback bucket by the per-peer cap, that bucket also holds failures
// contributed by other key ids — and by junk-shaped tokens — from the same address, so
// removing it would let one successful authentication wipe an in-progress spray. The shared
// bucket decays by window expiry instead; the caller still recovers through probe admission.
if (string.Equals(effectiveKeyId, partition.KeyId, StringComparison.Ordinal))
{
RemovePartition(partitionKey);
}
// The aggregate is cleared on the presented key id even when the composite partition
// collapsed to the fallback: a verified secret is proof the key is not under successful
@@ -264,20 +286,41 @@ public sealed class ApiKeyFailureLimiter
}
}
private bool IsProbeDue(WindowState state, long now)
/// <summary>
/// Atomically claims this interval's probe slot. Checking and advancing must happen in one
/// critical section: a due-check that released the lock before advancing would let every request
/// arriving at the interval boundary observe "due" and all be admitted, which is exactly the
/// unbounded-guessing burst the probe interval exists to prevent.
/// </summary>
private bool TryConsumeProbe(WindowState state, long now, out long previousProbeAtTicks)
{
lock (state)
{
return now >= state.NextProbeAtTicks;
previousProbeAtTicks = state.NextProbeAtTicks;
if (now < previousProbeAtTicks)
{
return false;
}
state.NextProbeAtTicks = now + _probeIntervalTicks;
state.LastActivityTicks = now;
return true;
}
}
private void ConsumeProbe(WindowState state, long now)
/// <summary>
/// Returns a probe slot claimed for a request that a later layer then refused, so the wasted
/// reservation does not cost the next arrival its slot. Only the caller's own reservation is
/// undone — a slot re-granted or re-armed in the meantime wins.
/// </summary>
private void ReleaseProbe(WindowState state, long now, long previousProbeAtTicks)
{
lock (state)
{
state.NextProbeAtTicks = now + _probeIntervalTicks;
state.LastActivityTicks = now;
if (state.NextProbeAtTicks == now + _probeIntervalTicks)
{
state.NextProbeAtTicks = previousProbeAtTicks;
}
}
}