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

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
@@ -125,6 +125,100 @@ public sealed class ApiKeyFailureLimiterTests
Assert.Equal(ApiKeyThrottleDecision.ProbeAdmitted, limiter.Check(partition));
}
/// <summary>
/// The probe slot is claimed atomically: when a crowd of requests arrives at the same interval
/// boundary exactly one is admitted and the rest are still refused. A check-then-act grant would
/// let every arrival observe "due" and hand the whole burst through to the verifier.
/// </summary>
[Fact]
public void ProbeAdmission_UnderConcurrentArrivals_GrantsExactlyOneSlot()
{
ManualTimeProvider clock = new(DateTimeOffset.UnixEpoch);
ApiKeyFailureLimiter limiter = CreateLimiter(clock, limit: 3, aggregateLimit: 5);
ApiKeyThrottlePartition partition = new("ipv4:10.0.0.1:1", "victim");
// Five failures on one partition trip both layers (limit 3, aggregate 5), so the concurrent
// arrivals contend for the composite probe slot AND the aggregate's.
RecordFailures(limiter, partition, 5);
clock.Advance(ProbeInterval);
// The unsafe window between reading "probe due" and advancing the slot is nanoseconds wide,
// so one burst can miss it by luck. Repeating the boundary makes the red reliable, while the
// atomic implementation must yield exactly one admission in every round.
const int arrivals = 8;
const int rounds = 200;
int totalAdmitted = 0;
for (int round = 0; round < rounds; round++)
{
// Re-arm: the failures keep both layers over their limits, and recording pushes the next
// probe one interval out, which the advance below then reaches.
RecordFailures(limiter, partition, 5);
clock.Advance(ProbeInterval);
ApiKeyThrottleDecision[] decisions = new ApiKeyThrottleDecision[arrivals];
using (Barrier startLine = new(arrivals))
{
Thread[] threads = new Thread[arrivals];
for (int index = 0; index < arrivals; index++)
{
int slot = index;
threads[slot] = new Thread(() =>
{
startLine.SignalAndWait();
decisions[slot] = limiter.Check(partition);
});
threads[slot].Start();
}
foreach (Thread thread in threads)
{
Assert.True(thread.Join(TimeSpan.FromSeconds(30)), "probe-contention thread did not finish");
}
}
int admitted = decisions.Count(decision => decision == ApiKeyThrottleDecision.ProbeAdmitted);
Assert.Equal(
arrivals - admitted,
decisions.Count(decision => decision is ApiKeyThrottleDecision.ThrottledByPeer
or ApiKeyThrottleDecision.ThrottledByAggregate));
totalAdmitted += admitted;
}
Assert.Equal(rounds, totalAdmitted);
}
/// <summary>
/// The layers claim their probe slots one at a time, so a slot claimed on the composite partition
/// must be returned when the aggregate then refuses. Otherwise a refused request would silently
/// spend the partition's next slot and push the legitimate holder out by a full interval.
/// </summary>
[Fact]
public void ProbeAdmission_WhenAggregateRefuses_ReturnsTheClaimedPeerSlot()
{
ManualTimeProvider clock = new(DateTimeOffset.UnixEpoch);
ApiKeyFailureLimiter limiter = CreateLimiter(clock, limit: 3, aggregateLimit: 5);
ApiKeyThrottlePartition holder = new("ipv4:10.0.0.1:1", "victim");
// Both layers trip together, arming both slots for t0 + interval.
RecordFailures(limiter, holder, 5);
// A failure from a second address re-arms only the aggregate (its own composite is far under
// the limit), so the aggregate's slot now opens one interval later than the partition's.
clock.Advance(TimeSpan.FromSeconds(3));
limiter.RecordFailure(new ApiKeyThrottlePartition("ipv4:10.0.0.2:1", "victim"));
// t0 + 5s: the partition's slot is due, the aggregate's is not — the request is refused and
// the partition's slot must be handed back rather than consumed.
clock.Advance(TimeSpan.FromSeconds(2));
Assert.Equal(ApiKeyThrottleDecision.ThrottledByAggregate, limiter.Check(holder));
// t0 + 8s: the aggregate's slot opens. The partition's slot was restored, so this passes; had
// it been consumed above it would not reopen until t0 + 10s and this would be ThrottledByPeer.
clock.Advance(TimeSpan.FromSeconds(3));
Assert.Equal(ApiKeyThrottleDecision.ProbeAdmitted, limiter.Check(holder));
}
/// <summary>A zero probe interval restores absolute blocking (documented as not recommended).</summary>
[Fact]
public void ProbeIntervalZero_BlocksAbsolutely()
@@ -165,6 +259,39 @@ public sealed class ApiKeyFailureLimiterTests
limiter.Check(new ApiKeyThrottlePartition("ipv4:10.0.9.9:1", "victim")));
}
/// <summary>
/// A success whose key id was squeezed into the address's shared fallback bucket must not clear
/// that bucket: it carries failures from other key ids at the same address, so clearing it would
/// make one successful authentication a reset button for an in-progress spray.
/// </summary>
[Fact]
public void Reset_WithOverCapKeyId_DoesNotClearSharedFallbackPartition()
{
ManualTimeProvider clock = new(DateTimeOffset.UnixEpoch);
ApiKeyFailureLimiter limiter = CreateLimiter(clock, limit: 3, aggregateLimit: 100, maxPartitions: 4096);
const string peer = "ipv4:10.0.0.1:1";
// Fill the per-peer key-id cap, then spray past it so the overflow lands on — and trips —
// the address's shared fallback partition.
for (int i = 0; i < ApiKeyFailureLimiter.MaxKeyIdPartitionsPerPeer; i++)
{
limiter.RecordFailure(new ApiKeyThrottlePartition(peer, $"key{i}"));
}
for (int i = 0; i < 5; i++)
{
limiter.RecordFailure(new ApiKeyThrottlePartition(peer, $"overflow{i}"));
}
ApiKeyThrottlePartition overCap = new(peer, "overflow0");
Assert.Equal(ApiKeyThrottleDecision.ThrottledByPeer, limiter.Check(overCap));
limiter.Reset(overCap);
Assert.True(limiter.IsTracked(overCap));
Assert.Equal(ApiKeyThrottleDecision.ThrottledByPeer, limiter.Check(overCap));
}
/// <summary>
/// SEC-32: a spray of unique junk partitions (junk tokens resolve to the sender's fallback
/// partition, one per address) must not evict a partition that is currently throttled — the