Failed-key throttle: reserve the attempt before authenticating; never…

app/auth.py

@@ -583,15 +583,20 @@ async def identity_middleware(request: Request, call_next):
             # A per-tenant API key (JARVIS_KEY_TENANTS): authenticated against the control schema, bound to the key's own tenant. Every way of failing is the one 401 a wrong ledger
             # key gets; an unreachable control schema is a 503 (never an admission); a valid key of a tenant that is not active is a 503 for that tenant only.
             client_addr = request.client.host if request.client else "unknown"
-            if key_auth.attempts_blocked(candidates, client_addr):
+            reserved = key_auth.reserve_attempt(candidates, client_addr)
+            if reserved is None:
                 return json_response(429, "Too many failed attempts", headers={"Retry-After": "60"}, path=path)
             try:
                 resolved = await run_in_threadpool(key_auth.authenticate, candidates)
             except key_auth.KeyAuthUnavailable:
+                key_auth.release_attempt(reserved)             # an outage is not a wrong guess
                 return json_response(503, "Ledger unavailable", path=path)
+            except BaseException:
+                key_auth.release_attempt(reserved)             # nor is a cancelled or crashed request
+                raise
             if resolved is None:
-                key_auth.record_failure(candidates, client_addr)
-                return json_response(401, "Invalid or missing API key", path=path)
+                return json_response(401, "Invalid or missing API key", path=path)      # the reservation stays counted: it was a failed guess
+            key_auth.release_attempt(reserved)
             if resolved.tenant_status != "active":
                 return json_response(503, "tenant_unavailable", path=path)
             if READ_SCOPE not in resolved.scopes:

app/key_auth.py

@@ -58,20 +58,30 @@ def _failure_subjects(candidates: Iterable[str], client: str) -> list[str]:
     return out or [f"keyfail:{client}:-"]
 
 
-def attempts_blocked(candidates: Iterable[str], client: str) -> bool:
-    """True when this address has already failed `failure_limit()` times against one of these key ids this minute: the request is refused with 429 BEFORE any HMAC is computed or the
-    database consulted, so guessing one key's secret is throttled. Keyed by address as well as key id, so a stranger who sends junk naming your key id cannot lock you out of your own key."""
+def reserve_attempt(candidates: Iterable[str], client: str) -> list[tuple[str, float]] | None:
+    """Count this attempt against every (client address, key id) it names BEFORE any HMAC or database work, atomically. None means the address has already used its failures for one of
+    those key ids this minute (a 429, nothing computed). Otherwise the returned subjects must be `release_attempt`ed if the attempt turns out not to be a failed guess (success, or a
+    database outage), and are left counted when it fails. Counting up front means a burst of concurrent guesses cannot all be admitted by a check made against an empty bucket. Keyed by address
+    as well as key id, so a stranger who sends junk naming your key id cannot lock you out of your own key."""
     from app.public_security import key_failure_limiter
 
     limit = failure_limit()
-    return any(key_failure_limiter.blocked(subject, limit) for subject in _failure_subjects(candidates, client))
+    taken: list[tuple[str, float]] = []
+    for subject in _failure_subjects(candidates, client):
+        reservation = key_failure_limiter.reserve(subject, limit)
+        if reservation is not None:
+            taken.append((subject, reservation))
+        else:
+            release_attempt(taken)
+            return None
+    return taken
 
 
-def record_failure(candidates: Iterable[str], client: str) -> None:
+def release_attempt(reservations: Iterable[tuple[str, float]]) -> None:
     from app.public_security import key_failure_limiter
 
-    for subject in _failure_subjects(candidates, client):
-        key_failure_limiter.allow(subject, limit=10**9)      # count it; the decision to refuse is made by attempts_blocked
+    for subject, reservation in reservations:
+        key_failure_limiter.release(subject, reservation)
 
 
 class KeyAuthUnavailable(RuntimeError):

app/public_security.py

@@ -51,27 +51,47 @@ def allow(self, subject: str, limit: int | None = None) -> bool:
             window.append(now)
             return True
 
-    def blocked(self, subject: str, limit: int) -> bool:
-        """Is `subject` already at `limit` in the current window? (Does not count this call.)"""
+    def reserve(self, subject: str, limit: int) -> float | None:
+        """Atomically count one in-flight attempt against `subject` unless it is already at `limit`. The caller `release`s it if the attempt turns out not to be a failure. Counting BEFORE the
+        expensive work means a burst of concurrent attempts cannot all pass a check made against an empty bucket. Buckets at their limit are never evicted to make room (that would reset
+        the throttle); when the table is full of such buckets a new subject is refused instead. Returns the reservation (its timestamp) for `release`, or None when refused."""
         now = time.monotonic()
         with self._lock:
-            window = self._requests.get(subject)
-            if window is None:
-                return False
+            self._protect_at = limit
+            if len(self._requests) >= self.MAX_BUCKETS and subject not in self._requests:
+                self._evict(now)
+                if len(self._requests) >= self.MAX_BUCKETS:
+                    return None
+            window = self._requests[subject]
             while window and now - window[0] >= self._window:
                 window.popleft()
-            return len(window) >= limit
+            if len(window) >= limit:
+                return None
+            window.append(now)
+            return now
+
+    def release(self, subject: str, reservation: float) -> None:
+        """Remove THIS reservation, not the newest one: overlapping attempts finish in any order, and a finished success must not erase another attempt's failed guess."""
+        with self._lock:
+            window = self._requests.get(subject)
+            if window:
+                try:
+                    window.remove(reservation)
+                except ValueError:
+                    pass                                                   # already aged out of the window
 
     def _evict(self, now: float) -> None:
         for key in [k for k, w in self._requests.items() if not w or now - w[-1] >= self._window]:
             del self._requests[key]
-        if len(self._requests) >= self.MAX_BUCKETS:             # every bucket is live: drop the oldest-touched half rather than refuse to track anyone new
-            for key, _ in sorted(self._requests.items(), key=lambda kv: kv[1][-1])[: self.MAX_BUCKETS // 2]:
+        if len(self._requests) >= self.MAX_BUCKETS:             # every bucket is live: drop the oldest-touched half, but never one that is at its limit (a throttle must survive a flood of other subjects)
+            keep_at = getattr(self, "_protect_at", None)
+            droppable = [(k, w) for k, w in self._requests.items() if keep_at is None or len(w) < keep_at]
+            for key, _ in sorted(droppable, key=lambda kv: kv[1][-1])[: self.MAX_BUCKETS // 2]:
                 del self._requests[key]
 
 
 mcp_rate_limiter = McpRateLimiter()
-key_failure_limiter = McpRateLimiter()          # failed tenant-key attempts, keyed (client address, key id); see key_auth.attempts_blocked
+key_failure_limiter = McpRateLimiter()          # failed tenant-key attempts, keyed (client address, key id); see key_auth.reserve_attempt
 
 
 def rate_limit_subject(request: Request) -> str:

tests/test_rate_limit_credential.py

@@ -95,3 +95,67 @@ def test_more_than_one_worker_is_a_startup_refusal_unless_a_shared_limiter_is_de
         public_security.assert_config()
     monkeypatch.setenv("JARVIS_RATE_LIMIT_SHARED_BACKEND", "edge-waf")
     assert config_problems() == []
+
+
+def test_a_burst_of_concurrent_attempts_cannot_all_pass_a_check_made_against_an_empty_bucket():
+    """Codex on #112: the attempt is counted BEFORE the expensive verification, atomically, so at most `limit` of a concurrent burst are admitted."""
+    import threading
+
+    lim = McpRateLimiter()
+    admitted, gate = [], threading.Barrier(40)
+
+    def attempt():
+        gate.wait()
+        admitted.append(lim.reserve("keyfail:1.2.3.4:abc", 5) is not None)
+
+    threads = [threading.Thread(target=attempt) for _ in range(40)]
+    [t.start() for t in threads]
+    [t.join() for t in threads]
+    assert sum(admitted) == 5
+
+
+def test_a_released_reservation_is_not_counted_as_a_failure():
+    lim = McpRateLimiter()
+    for _ in range(50):
+        lim.release("s", lim.reserve("s", 3))                                                   # a successful credential, or an outage: never a wrong guess
+    assert [lim.reserve("s", 3) is not None for _ in range(4)] == [True, True, True, False]
+
+
+def test_a_throttled_bucket_survives_a_flood_of_other_subjects_and_new_ones_are_refused_when_full():
+    """Codex on #112: 10,000 distinct key ids must not evict the throttled target and reset its count."""
+    lim = McpRateLimiter()
+    lim.MAX_BUCKETS = 100
+    for _ in range(3):
+        assert lim.reserve("target", 3) is not None
+    assert lim.reserve("target", 3) is None                                # at its limit
+    for i in range(1000):
+        lim.reserve(f"junk-{i}", 3)                                        # the flood; each junk id has a single failed guess
+    assert lim.reserve("target", 3) is None, "the flood evicted the throttled bucket"
+    # and when every bucket is at its limit, a brand-new subject is refused rather than evicting one
+    full = McpRateLimiter()
+    full.MAX_BUCKETS = 5
+    for i in range(5):
+        for _ in range(3):
+            full.reserve(f"b{i}", 3)
+    assert full.reserve("new", 3) is None
+    assert full.reserve("b0", 3) is None
+
+
+def test_releasing_a_success_removes_its_own_reservation_not_a_newer_failures():
+    """Codex on #112 (round 2): attempts overlap and finish in any order. An old valid request finishing after a newer invalid one must leave the invalid one counted, at ITS time."""
+    import app.public_security as ps
+
+    clock = [100.0]
+    real = ps.time.monotonic
+    ps.time.monotonic = lambda: clock[0]
+    try:
+        lim = McpRateLimiter(window=60)
+        old_valid = lim.reserve("s", 2)
+        clock[0] = 150.0
+        newer_invalid = lim.reserve("s", 2)                                 # stays counted: it was a wrong guess
+        lim.release("s", old_valid)                                         # the older, valid one finishes last
+        assert list(lim._requests["s"]) == [newer_invalid]
+        clock[0] = 170.0                                                    # 70s after the old one, 20s after the failure: the failure must still count
+        assert lim.reserve("s", 2) is not None and lim.reserve("s", 2) is None
+    finally:
+        ps.time.monotonic = real