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..cf50f77 100644 --- a/app/key_auth.py +++ b/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): diff --git a/app/public_security.py b/app/public_security.py index c6ed504..1a78ef1 100644 --- a/app/public_security.py +++ b/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: diff --git a/tests/test_rate_limit_credential.py b/tests/test_rate_limit_credential.py index 980e1af..ddc1180 100644 --- a/tests/test_rate_limit_credential.py +++ b/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