fix(SEC-31,SEC-32): identify a probe-slot reservation by version, not by timestamp

ReleaseProbe recognised its own reservation by comparing NextProbeAtTicks to
now + _probeIntervalTicks. RecordInto's rearm-on-trip writes that identical
expression, so a concurrent RecordFailure on the same WindowState whose `now`
lands on the claimer's tick — routine at ~1 ms clock resolution under load — was
mistaken for the caller's own claim. The release then stomped the legitimate
fresh re-arm back to the stale previousProbeAtTicks, which is already due, handing
the next arrival a free probe the re-arm had just closed.

WindowState gains a monotonic ProbeVersion bumped by every writer of
NextProbeAtTicks (TryConsumeProbe's claim and RecordInto's re-arm alike).
TryConsumeProbe returns the stamp it set as part of a ProbeClaim; ReleaseProbe
restores the previous value only while the state's version still equals that
stamp, checking and restoring in one lock(state) section and bumping the version
again on restore so no other stale release can match either.

Test: ProbeSlotRestore_DoesNotStompConcurrentRearmAtSameTick, with the clock held
still so the claim and the interleaved failure necessarily share a tick. Making it
deterministic needed a seam — the claim-to-release window is a few nanoseconds and
racing threads do not hit it (an earlier thread-based attempt passed against the
defective guard three runs out of three, and its end state was ordering-dependent
rather than correctness-dependent, so it was dropped rather than shipped as
theatre). The seam is an internal ProbeReleaseInterleaveHook, null in production,
costing one null check on the already-refused path. Verified as a genuine red
against the timestamp guard: Expected ThrottledByPeer, Actual ProbeAdmitted.
This commit is contained in:
Joseph Doherty
2026-08-07 06:10:14 -04:00
parent acebe18773
commit 5b681ee59b
4 changed files with 102 additions and 13 deletions
@@ -112,6 +112,18 @@ public sealed class ApiKeyFailureLimiter
/// <summary>Gets the number of tracked per-key-id aggregates. Test seam.</summary>
internal int TrackedAggregateCount => _aggregates.Count;
/// <summary>
/// Test seam invoked between a probe claim and its compensating release; <see langword="null"/>
/// in production, where it costs one null check on the already-refused path.
/// </summary>
/// <remarks>
/// The window this straddles is a few nanoseconds wide, so the interleaving it exists to cover —
/// a concurrent failure re-arming the same state between the two — cannot be produced reliably
/// by racing threads. Without this seam the version guard in <see cref="ReleaseProbe"/> would be
/// verifiable only by inspection.
/// </remarks>
internal Action? ProbeReleaseInterleaveHook { get; set; }
/// <summary>Decides whether an authentication attempt may reach the verifier.</summary>
/// <param name="partition">The throttle partition derived from the request.</param>
/// <returns>The admission decision for this attempt.</returns>
@@ -150,11 +162,11 @@ public sealed class ApiKeyFailureLimiter
// 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;
ProbeClaim peerClaim = default;
bool peerProbeClaimed = false;
if (peerOver)
{
if (!TryConsumeProbe(peerState!, now, out peerProbeRestore))
if (!TryConsumeProbe(peerState!, now, out peerClaim))
{
return ApiKeyThrottleDecision.ThrottledByPeer;
}
@@ -166,7 +178,8 @@ public sealed class ApiKeyFailureLimiter
{
if (peerProbeClaimed)
{
ReleaseProbe(peerState!, now, peerProbeRestore);
ProbeReleaseInterleaveHook?.Invoke();
ReleaseProbe(peerState!, peerClaim);
}
return ApiKeyThrottleDecision.ThrottledByAggregate;
@@ -264,10 +277,13 @@ public sealed class ApiKeyFailureLimiter
state.LastActivityTicks = now;
// Arm (or push out) the probe slot whenever the state is at or over its limit, so the
// attempt that trips the limit is not itself followed by an immediate free probe.
// attempt that trips the limit is not itself followed by an immediate free probe. This
// is a write of NextProbeAtTicks, so it bumps the version that identifies a probe claim
// — otherwise a release could mistake this re-arm for its own reservation.
if (limit > 0 && state.FailureTicks.Count >= limit)
{
state.NextProbeAtTicks = now + _probeIntervalTicks;
state.ProbeVersion++;
}
}
}
@@ -292,35 +308,47 @@ public sealed class ApiKeyFailureLimiter
/// 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)
private bool TryConsumeProbe(WindowState state, long now, out ProbeClaim claim)
{
lock (state)
{
previousProbeAtTicks = state.NextProbeAtTicks;
long previousProbeAtTicks = state.NextProbeAtTicks;
if (now < previousProbeAtTicks)
{
claim = default;
return false;
}
state.NextProbeAtTicks = now + _probeIntervalTicks;
state.LastActivityTicks = now;
claim = new ProbeClaim(previousProbeAtTicks, ++state.ProbeVersion);
return true;
}
}
/// <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.
/// reservation does not cost the next arrival its slot.
/// </summary>
private void ReleaseProbe(WindowState state, long now, long previousProbeAtTicks)
/// <remarks>
/// The claim is identified by the per-state version stamped when it was made, never by the
/// timestamp it wrote. Every writer of <c>NextProbeAtTicks</c> bumps that version, so a re-arm
/// from a concurrent failure — which writes the identical <c>now + interval</c> expression, and
/// at ~1 ms clock resolution routinely lands on the same tick — cannot be mistaken for the
/// caller's own reservation and stomped back to a stale, already-due value. Restoring bumps the
/// version again so no other stale release can match either.
/// </remarks>
private static void ReleaseProbe(WindowState state, ProbeClaim claim)
{
lock (state)
{
if (state.NextProbeAtTicks == now + _probeIntervalTicks)
if (state.ProbeVersion != claim.Version)
{
state.NextProbeAtTicks = previousProbeAtTicks;
return;
}
state.NextProbeAtTicks = claim.PreviousProbeAtTicks;
state.ProbeVersion++;
}
}
@@ -504,8 +532,21 @@ public sealed class ApiKeyFailureLimiter
public long LastActivityTicks;
public long NextProbeAtTicks;
/// <summary>
/// Monotonic stamp bumped by every writer of <see cref="NextProbeAtTicks"/> (probe claim and
/// failure re-arm alike). It is what lets a compensating release recognise its own
/// reservation without comparing timestamps, which collide whenever two writers share a
/// clock tick.
/// </summary>
public long ProbeVersion;
}
/// <summary>A probe slot reservation: what to restore, and the stamp proving it is still ours.</summary>
/// <param name="PreviousProbeAtTicks">The slot value replaced when the claim was made.</param>
/// <param name="Version">The <see cref="WindowState.ProbeVersion"/> stamped by this claim.</param>
private readonly record struct ProbeClaim(long PreviousProbeAtTicks, long Version);
private sealed class PeerKeyIds
{
/// <summary>Key ids this transport peer has minted a partition for.</summary>
@@ -219,6 +219,54 @@ public sealed class ApiKeyFailureLimiterTests
Assert.Equal(ApiKeyThrottleDecision.ProbeAdmitted, limiter.Check(holder));
}
/// <summary>
/// A compensating release must never undo a re-arm written by a concurrent failure on the same
/// partition. Both writers store the identical <c>now + interval</c> value when they share a
/// clock tick, so identifying the caller's own reservation by timestamp would let the release
/// stomp a fresh re-arm back to an already-due value and reopen the probe slot early. The clock
/// is deliberately held still here, which forces exactly that collision.
/// </summary>
[Fact]
public void ProbeSlotRestore_DoesNotStompConcurrentRearmAtSameTick()
{
ManualTimeProvider clock = new(DateTimeOffset.UnixEpoch);
ApiKeyFailureLimiter limiter = CreateLimiter(clock, limit: 3, aggregateLimit: 5);
ApiKeyThrottlePartition holder = new("ipv4:10.0.0.1:1", "victim");
ApiKeyThrottlePartition other = new("ipv4:10.0.0.2:1", "victim");
// Trip both layers, then push the aggregate's slot one interval past the partition's, so
// every Check below claims the partition's slot and is then refused by the aggregate — the
// claim-and-compensate path under test.
RecordFailures(limiter, holder, 5);
clock.Advance(TimeSpan.FromSeconds(3));
limiter.RecordFailure(other);
clock.Advance(TimeSpan.FromSeconds(2));
// Land a failure on the same partition inside the claim-to-release window — the interleaving
// a concurrent RecordFailure produces, forced here so the assertion is deterministic. It
// shares the frozen clock tick with the claim, so both write the identical slot value.
int interleaved = 0;
limiter.ProbeReleaseInterleaveHook = () =>
{
if (Interlocked.Exchange(ref interleaved, 1) == 0)
{
limiter.RecordFailure(holder);
}
};
Assert.Equal(ApiKeyThrottleDecision.ThrottledByAggregate, limiter.Check(holder));
Assert.Equal(1, interleaved);
limiter.ProbeReleaseInterleaveHook = null;
// Drop the aggregate so the next decision reflects the composite partition alone.
limiter.Reset(other);
// The interleaved failure pushed the slot one interval past the (still unadvanced) clock, so
// no probe may be due. Restoring over it would leave the already-due earlier value and hand
// the next arrival a free probe.
Assert.Equal(ApiKeyThrottleDecision.ThrottledByPeer, limiter.Check(holder));
}
/// <summary>A zero probe interval restores absolute blocking (documented as not recommended).</summary>
[Fact]
public void ProbeIntervalZero_BlocksAbsolutely()