fix(SEC-38,SEC-39): preserve aggregate and fallback throttles

2026-08-25 11:00:29 -04:00
parent 780720165a
commit 1875426e4b
3 changed files with 72 additions and 10 deletions
@@ -101,6 +101,37 @@ public sealed class ApiKeyFailureLimiterTests
            limiter.Check(new ApiKeyThrottlePartition("ipv4:10.0.9.9:1", "other-key")));
    }

+    /// <summary>
+    ///     SEC-38: a valid key id still feeds and consults its cross-peer aggregate when each
+    ///     attacking peer has exhausted its 32 dedicated key-id partitions and therefore records
+    ///     the peer layer through the fallback bucket.
+    /// </summary>
+    [Fact]
+    public void AggregateLayer_TripsForValidKeyIdCollapsedToPeerFallback()
+    {
+        ManualTimeProvider clock = new(DateTimeOffset.UnixEpoch);
+        ApiKeyFailureLimiter limiter = CreateLimiter(
+            clock,
+            limit: 100,
+            aggregateLimit: 3,
+            maxPartitions: 4096);
+
+        for (int peerIndex = 0; peerIndex < 3; peerIndex++)
+        {
+            string peer = $"ipv4:10.0.0.{peerIndex + 1}:1";
+            for (int keyIndex = 0; keyIndex < ApiKeyFailureLimiter.MaxKeyIdPartitionsPerPeer; keyIndex++)
+            {
+                limiter.RecordFailure(new ApiKeyThrottlePartition(peer, $"prefill-{peerIndex}-{keyIndex}"));
+            }
+
+            limiter.RecordFailure(new ApiKeyThrottlePartition(peer, "victim"));
+        }
+
+        Assert.Equal(
+            ApiKeyThrottleDecision.ThrottledByAggregate,
+            limiter.Check(new ApiKeyThrottlePartition("ipv4:10.0.9.9:1", "victim")));
+    }
+
    /// <summary>
    /// A throttled partition is a valve, not a wall: exactly one request per probe interval is
    /// admitted to the real verifier, so the holder of the correct secret can always get through.
@@ -340,6 +371,29 @@ public sealed class ApiKeyFailureLimiterTests
        Assert.Equal(ApiKeyThrottleDecision.ThrottledByPeer, limiter.Check(overCap));
    }

+    /// <summary>
+    ///     SEC-39: a request with no effective key id does not own the peer's shared fallback
+    ///     partition. Resetting after such a request must leave other fallback failures tracked.
+    /// </summary>
+    [Fact]
+    public void Reset_WithNullEffectiveKeyId_DoesNotClearSharedFallbackPartition()
+    {
+        ManualTimeProvider clock = new(DateTimeOffset.UnixEpoch);
+        ApiKeyFailureLimiter limiter = CreateLimiter(clock, limit: 3, aggregateLimit: 100);
+        ApiKeyThrottlePartition fallback = new("ipv4:10.0.0.1:1", KeyId: null);
+
+        RecordFailures(limiter, fallback, 3);
+        Assert.Equal(
+            ApiKeyThrottleDecision.ThrottledByPeer,
+            limiter.Check(fallback, out ApiKeyFailureLimiter.PartitionResolution resolution));
+        Assert.Null(resolution.EffectiveKeyId);
+
+        limiter.Reset(fallback, resolution);
+
+        Assert.True(limiter.IsTracked(fallback));
+        Assert.Equal(ApiKeyThrottleDecision.ThrottledByPeer, limiter.Check(fallback));
+    }
+
    /// <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