fix: quota pending sign-ins and registrations per address, keep young clients, write the store off the loop (SEC-09)

tests/security/test_oauth_input_bounds.py

@@ -2,11 +2,28 @@
 
 from __future__ import annotations
 
-import pytest
+import threading
+from typing import Any
 
+import pytest
+from mcp.server.auth.provider import RegistrationError
+from starlette.applications import Starlette
+from starlette.testclient import TestClient
+
+from aegis_gitea_mcp.oauth import provider as provider_module
 from aegis_gitea_mcp.oauth.callback import MAX_CONSENT_BODY_BYTES
+from aegis_gitea_mcp.oauth.provider import MAX_PENDING_PER_SOURCE, MAX_REGISTRATIONS_PER_SOURCE
+from aegis_gitea_mcp.request_context import set_client_address
 from tests.security.helpers import asgi_request
-from tests.test_oauth_helpers import BASE, build_oauth_app
+from tests.test_oauth_helpers import (
+    BASE,
+    REDIRECT,
+    build_oauth_app,
+    oauth_settings,
+    pkce,
+    register,
+)
+from tests.test_oauth_provider import client_info
 
 CHUNK = 64 * 1024
 
@@ -87,3 +104,102 @@ async def test_small_urlencoded_consent_post_still_reaches_the_handler() -> None
     # Reaches the handler, which answers its fixed 400 page for an unknown request.
     assert result.status == 400
     assert b"invalid or has expired" in result.body
+
+
+# SEC-09: pre-sign-in state cannot be flushed by one address
+
+
+def client_at(app: Starlette, address: str) -> TestClient:
+    """A browser-like client whose requests come from ``address``."""
+    return TestClient(app, base_url=BASE, client=(address, 50000), follow_redirects=False)
+
+
+def authorize(client: TestClient, registration: dict[str, str]) -> Any:
+    return client.get(
+        "/authorize",
+        params={
+            "client_id": registration["client_id"],
+            "redirect_uri": REDIRECT,
+            "response_type": "code",
+            "code_challenge": pkce()[1],
+            "code_challenge_method": "S256",
+        },
+    )
+
+
+def test_a_flood_of_authorize_calls_from_one_address_does_not_evict_another_users_request() -> None:
+    app = build_oauth_app(oauth_settings(rate_limit_per_minute=1_000_000)).app
+    victim = client_at(app, "198.51.100.20")
+    consent_url = authorize(victim, register(victim).json()).headers["location"]
+    assert victim.get(consent_url).status_code == 200
+
+    attacker = client_at(app, "203.0.113.66")
+    registration = register(attacker).json()
+    accepted = refused = 0
+    for _ in range(1100):
+        location = authorize(attacker, registration).headers["location"]
+        if "/oauth/consent" in location:
+            accepted += 1
+        else:
+            refused += 1
+            assert "error=temporarily_unavailable" in location
+
+    assert accepted == MAX_PENDING_PER_SOURCE
+    assert refused == 1100 - MAX_PENDING_PER_SOURCE
+    assert victim.get(consent_url).status_code == 200  # still pending
+
+
+def test_registrations_per_address_are_limited() -> None:
+    app = build_oauth_app(oauth_settings(rate_limit_per_minute=1_000_000)).app
+    flooder = client_at(app, "203.0.113.66")
+
+    statuses = [register(flooder).status_code for _ in range(MAX_REGISTRATIONS_PER_SOURCE + 5)]
+
+    assert statuses[:MAX_REGISTRATIONS_PER_SOURCE] == [201] * MAX_REGISTRATIONS_PER_SOURCE
+    assert statuses[MAX_REGISTRATIONS_PER_SOURCE:] == [400] * 5
+    # Another address is unaffected.
+    assert register(client_at(app, "198.51.100.20")).status_code == 201
+
+
+async def test_a_young_client_survives_a_registration_flood(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(provider_module, "MAX_CLIENTS", 10)
+    oauth = build_oauth_app()
+    provider = oauth.provider
+    set_client_address("198.51.100.20")
+    await provider.register_client(client_info("legit"))
+    oauth.clock.advance(60)
+
+    refused = 0
+    for number in range(40):
+        set_client_address(f"203.0.113.{number}")
+        try:
+            await provider.register_client(client_info(f"flood-{number}"))
+        except RegistrationError:
+            refused += 1
+
+    assert "legit" in provider._store.clients
+    assert refused > 0 and len(provider._store.clients) == 10
+
+    # Once clients are older than the pending window, old ones can be evicted again.
+    oauth.clock.advance(provider_module.PENDING_TTL + 1)
+    set_client_address("198.51.100.99")
+    await provider.register_client(client_info("later"))
+    assert "later" in provider._store.clients
+
+
+async def test_the_store_write_runs_off_the_event_loop(monkeypatch: pytest.MonkeyPatch) -> None:
+    store = build_oauth_app().store
+    threads: list[str] = []
+    original = store._write
+
+    def recording(data: bytes) -> None:
+        threads.append(threading.current_thread().name)
+        original(data)
+
+    monkeypatch.setattr(store, "_write", recording)
+
+    await store.save()
+
+    assert threads and threads[0] != threading.main_thread().name