From 0d9ae4e5242a7c186ffede54069854e125b1a4ab Mon Sep 17 00:00:00 2001 From: warheart1984-ctrl Date: Sat, 17 Oct 2026 01:39:11 -0400 Subject: [PATCH 1/2] Failed-key throttle: reserve the attempt before authenticating; never evict a throttled bucket (Codex on #112) Co-Authored-By: Claude Sonnet 5.5 --- app/auth.py | 11 +++++-- app/key_auth.py | 23 ++++++++++----- app/public_security.py | 34 ++++++++++++++++------ tests/test_rate_limit_credential.py | 45 +++++++++++++++++++++++++++++ 4 files changed, 94 insertions(+), 19 deletions(-) diff --git a/app/auth.py b/app/auth.py index 0b5eac0..b592824 100644 --- a/app/auth.py +++ b/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: diff --git a/app/key_auth.py b/app/key_auth.py index 3449581..f58246b 100644 --- a/app/key_auth.py +++ b/app/key_auth.py @@ -58,20 +58,29 @@ 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[str] | 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[str] = [] + for subject in _failure_subjects(candidates, client): + if key_failure_limiter.reserve(subject, limit): + taken.append(subject) + else: + release_attempt(taken) + return None + return taken -def record_failure(candidates: Iterable[str], client: str) -> None: +def release_attempt(subjects: Iterable[str]) -> 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 in subjects: + key_failure_limiter.release(subject) class KeyAuthUnavailable(RuntimeError): diff --git a/app/public_security.py b/app/public_security.py index c6ed504..88ca333 100644 --- a/app/public_security.py +++ b/app/public_security.py @@ -51,27 +51,43 @@ 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) -> bool: + """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.""" 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 False + window = self._requests[subject] while window and now - window[0] >= self._window: window.popleft() - return len(window) >= limit + if len(window) >= limit: + return False + window.append(now) + return True + + def release(self, subject: str) -> None: + with self._lock: + window = self._requests.get(subject) + if window: + window.pop() 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: diff --git a/tests/test_rate_limit_credential.py b/tests/test_rate_limit_credential.py index 980e1af..597f7f6 100644 --- a/tests/test_rate_limit_credential.py +++ b/tests/test_rate_limit_credential.py @@ -95,3 +95,48 @@ 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)) + + 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): + assert lim.reserve("s", 3) + lim.release("s") # a successful credential, or an outage: never a wrong guess + assert [lim.reserve("s", 3) 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) + assert not lim.reserve("target", 3) # 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 not lim.reserve("target", 3), "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 not full.reserve("new", 3) + assert not full.reserve("b0", 3) From 1fa4abf57df73f221fa4a3b0ae112a1641ef720b Mon Sep 17 00:00:00 2001 From: warheart1984-ctrl Date: Sat, 17 Oct 2026 08:43:46 -0400 Subject: [PATCH 2/2] Failed-key throttle: release the reservation that completed, not the newest (Codex on #112) Co-Authored-By: Claude Sonnet 5.5 --- app/key_auth.py | 15 ++++++------ app/public_security.py | 18 ++++++++------ tests/test_rate_limit_credential.py | 37 ++++++++++++++++++++++------- 3 files changed, 47 insertions(+), 23 deletions(-) diff --git a/app/key_auth.py b/app/key_auth.py index f58246b..cf50f77 100644 --- a/app/key_auth.py +++ b/app/key_auth.py @@ -58,7 +58,7 @@ def _failure_subjects(candidates: Iterable[str], client: str) -> list[str]: return out or [f"keyfail:{client}:-"] -def reserve_attempt(candidates: Iterable[str], client: str) -> list[str] | None: +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 @@ -66,21 +66,22 @@ def reserve_attempt(candidates: Iterable[str], client: str) -> list[str] | None: from app.public_security import key_failure_limiter limit = failure_limit() - taken: list[str] = [] + taken: list[tuple[str, float]] = [] for subject in _failure_subjects(candidates, client): - if key_failure_limiter.reserve(subject, limit): - taken.append(subject) + 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 release_attempt(subjects: Iterable[str]) -> None: +def release_attempt(reservations: Iterable[tuple[str, float]]) -> None: from app.public_security import key_failure_limiter - for subject in subjects: - key_failure_limiter.release(subject) + for subject, reservation in reservations: + key_failure_limiter.release(subject, reservation) class KeyAuthUnavailable(RuntimeError): diff --git a/app/public_security.py b/app/public_security.py index 88ca333..1a78ef1 100644 --- a/app/public_security.py +++ b/app/public_security.py @@ -51,30 +51,34 @@ def allow(self, subject: str, limit: int | None = None) -> bool: window.append(now) return True - def reserve(self, subject: str, limit: int) -> bool: + 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.""" + 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: 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 False + return None window = self._requests[subject] while window and now - window[0] >= self._window: window.popleft() if len(window) >= limit: - return False + return None window.append(now) - return True + return now - def release(self, subject: str) -> None: + 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: - window.pop() + 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]: diff --git a/tests/test_rate_limit_credential.py b/tests/test_rate_limit_credential.py index 597f7f6..ddc1180 100644 --- a/tests/test_rate_limit_credential.py +++ b/tests/test_rate_limit_credential.py @@ -106,7 +106,7 @@ def test_a_burst_of_concurrent_attempts_cannot_all_pass_a_check_made_against_an_ def attempt(): gate.wait() - admitted.append(lim.reserve("keyfail:1.2.3.4:abc", 5)) + 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] @@ -117,9 +117,8 @@ def attempt(): def test_a_released_reservation_is_not_counted_as_a_failure(): lim = McpRateLimiter() for _ in range(50): - assert lim.reserve("s", 3) - lim.release("s") # a successful credential, or an outage: never a wrong guess - assert [lim.reserve("s", 3) for _ in range(4)] == [True, True, True, False] + 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(): @@ -127,16 +126,36 @@ def test_a_throttled_bucket_survives_a_flood_of_other_subjects_and_new_ones_are_ lim = McpRateLimiter() lim.MAX_BUCKETS = 100 for _ in range(3): - assert lim.reserve("target", 3) - assert not lim.reserve("target", 3) # at its limit + 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 not lim.reserve("target", 3), "the flood evicted the throttled bucket" + 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 not full.reserve("new", 3) - assert not full.reserve("b0", 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