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

docs/security.md

@@ -98,7 +98,7 @@ Denial of service
 | Large request bodies | SDK body limit on the MCP endpoint and SDK OAuth routes; upload cap checked before decoding | `server.py` (`max_request_body_size`), `gitea/encoding.py` (consent route: SEC-01) |
 | Large or hostile responses | 16 MiB text cap, `MAX_BINARY_BYTES`, 64 KiB error bodies, JSON depth 100, item and character limits | `gitea/executor.py` |
 | Request floods | Sliding window per IP and per token hash, bounded key sets, rate-limit events audited once per window | `rate_limit.py`, `transports.py` (`ForwardedHeadersMiddleware`) |
-| Unbounded OAuth state | Caps on clients (500), pending sign-ins (1024), codes (1024), sessions per login (50); store pruning | `oauth/provider.py`, `oauth/store.py` (flush by floods: SEC-09) |
+| Unbounded OAuth state | Caps on clients (500), pending sign-ins (1024), codes (1024), sessions per login (50); per client address at most 20 pending sign-ins and 20 registrations per hour; clients younger than the pending window are never evicted; store pruning; the store write (fsync) runs in a worker thread | `oauth/provider.py`, `oauth/store.py` |
 
 Elevation of privilege
 

src/aegis_gitea_mcp/oauth/provider.py

@@ -69,6 +69,7 @@ from aegis_gitea_mcp.oauth.store import (
     token_hash,
 )
 from aegis_gitea_mcp.oauth.upstream import GiteaUpstream, UpstreamError
+from aegis_gitea_mcp.request_context import get_client_address
 
 logger = logging.getLogger(__name__)
 
@@ -83,6 +84,10 @@ MAX_CLIENTS = 500
 MAX_CLIENT_METADATA_BYTES = 8192
 MAX_REDIRECT_URIS = 10
 MAX_SESSIONS_PER_LOGIN = 50
+#: Per client address: sign-ins waiting for consent or Gitea, and registrations per hour.
+MAX_PENDING_PER_SOURCE = 20
+MAX_REGISTRATIONS_PER_SOURCE = 20
+REGISTRATION_WINDOW = 3600
 CONSENT_PATH = "/oauth/consent"
 
 _LOOPBACK_REDIRECT_HOSTS = frozenset({"localhost", "127.0.0.1", "::1"})
@@ -173,6 +178,11 @@ class _ExpiringMap(Generic[V]):
             del self._entries[next(iter(self._entries))]
         self._entries[key] = (self._clock() + ttl, value)
 
+    def count_where(self, predicate: Callable[[V], bool]) -> int:
+        """Number of live entries whose value satisfies ``predicate``."""
+        self._purge()
+        return sum(1 for _, value in self._entries.values() if predicate(value))
+
     def get(self, key: str) -> V | None:
         entry = self._entries.get(key)
         if entry is None:
@@ -203,6 +213,8 @@ class PendingAuthorization:
     code_verifier: str = field(repr=False)
     csrf_hash: str | None = field(default=None, repr=False)
     binding_hash: str | None = field(default=None, repr=False)
+    #: Rate-limit bucket of the client address that started the request (quota key).
+    source: str | None = None
 
 
 class StoredAuthorizationCode(AuthorizationCode):
@@ -264,6 +276,8 @@ class GiteaOAuthProvider(
         self._codes: _ExpiringMap[StoredAuthorizationCode] = _ExpiringMap(MAX_CODES, clock)
         # Hashes of exchanged codes -> session id, to detect replays (RFC 6749 4.1.2).
         self._used_codes: _ExpiringMap[str] = _ExpiringMap(MAX_CODES, clock)
+        # Registration times per client address (in memory), for the registration quota.
+        self._registrations: dict[str, list[float]] = {}
 
     # Helpers
 
@@ -348,6 +362,7 @@ class GiteaOAuthProvider(
         info = client_info.model_dump(mode="json", exclude_none=True)
         if len(json.dumps(info)) > MAX_CLIENT_METADATA_BYTES:
             raise RegistrationError("invalid_client_metadata", "client metadata is too large")
+        self._check_registration_quota()
         if len(self._store.clients) >= MAX_CLIENTS and not self._evict_one_client():
             logger.warning("oauth client registration refused: client limit reached")
             raise RegistrationError(
@@ -360,11 +375,41 @@ class GiteaOAuthProvider(
         await self._store.save()
         logger.info("oauth client registered: %s", client_info.client_id)
 
+    def _check_registration_quota(self) -> None:
+        """Refuse a registration when this address registered too many clients recently."""
+        source = get_client_address()
+        if source is None:
+            return
+        now = self._clock()
+        recent = [t for t in self._registrations.get(source, []) if t > now - REGISTRATION_WINDOW]
+        if len(recent) >= MAX_REGISTRATIONS_PER_SOURCE:
+            self._registrations[source] = recent
+            logger.warning("oauth client registration refused: quota of one address reached")
+            raise RegistrationError(
+                "invalid_client_metadata", "too many registrations from this address; try later"
+            )
+        recent.append(now)
+        self._registrations[source] = recent
+        if len(self._registrations) > MAX_CLIENTS * 4:
+            for key in [
+                k
+                for k, times in self._registrations.items()
+                if times[-1] <= now - REGISTRATION_WINDOW
+            ]:
+                del self._registrations[key]
+
     def _evict_one_client(self) -> bool:
+        """Evict the least recently used client without sessions, never a young one.
+
+        A client younger than :data:`PENDING_TTL` may still be on its way to its first
+        sign-in; evicting it would let a registration flood cancel those sign-ins.
+        """
+        now = self._clock()
         candidates = [
             (record.last_used_at, client_id)
             for client_id, record in self._store.clients.items()
             if not self._store.client_has_sessions(client_id)
+            and record.registered_at + PENDING_TTL <= now
         ]
         if not candidates:
             return False
@@ -392,6 +437,14 @@ class GiteaOAuthProvider(
             is not None
         ):
             raise AuthorizeError("invalid_request", "redirect URI is not allowed")
+        source = get_client_address()
+        if (
+            source is not None
+            and self._requests.count_where(lambda item: item.source == source)
+            >= MAX_PENDING_PER_SOURCE
+        ):
+            logger.warning("oauth authorize refused: pending sign-in quota of one address reached")
+            raise AuthorizeError("temporarily_unavailable", "too many pending sign-ins")
         request_id = secrets.token_urlsafe(32)
         pending = PendingAuthorization(
             client_id=client.client_id,
@@ -403,6 +456,7 @@ class GiteaOAuthProvider(
             scopes=list(params.scopes or [MCP_SCOPE]),
             resource=self.resource,
             code_verifier=secrets.token_urlsafe(64),
+            source=source,
         )
         self._requests.put(token_hash(request_id), pending, PENDING_TTL)
         return f"{self.base_url}{CONSENT_PATH}?{urlencode({'request': request_id})}"

src/aegis_gitea_mcp/oauth/store.py

@@ -40,6 +40,7 @@ from dataclasses import asdict, dataclass, field
 from pathlib import Path
 from typing import Any, Literal
 
+import anyio.to_thread
 from cryptography.fernet import Fernet, InvalidToken
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.kdf.hkdf import HKDF
@@ -272,7 +273,10 @@ class OAuthStore:
         """Prune expired entries and persist the state atomically (encrypted, mode 0600)."""
         async with self._lock:
             self.prune()
-            self._write(self._fernet.encrypt(self._document()))
+            data = self._fernet.encrypt(self._document())
+            # The write fsyncs: keep that off the event loop (the bytes are already built,
+            # so the thread touches no shared state).
+            await anyio.to_thread.run_sync(self._write, data)
 
     # Maintenance
 

src/aegis_gitea_mcp/request_context.py

@@ -7,6 +7,7 @@ from contextvars import ContextVar
 _REQUEST_ID: ContextVar[str] = ContextVar("request_id", default="-")
 _GITEA_USER_TOKEN: ContextVar[str | None] = ContextVar("gitea_user_token", default=None)
 _GITEA_USER_LOGIN: ContextVar[str | None] = ContextVar("gitea_user_login", default=None)
+_CLIENT_ADDRESS: ContextVar[str | None] = ContextVar("client_address", default=None)
 _GITEA_USER_SCOPES: ContextVar[tuple[str, ...]] = ContextVar("gitea_user_scopes", default=())
 
 
@@ -20,6 +21,16 @@ def get_request_id() -> str:
     return _REQUEST_ID.get()
 
 
+def set_client_address(address: str | None) -> None:
+    """Store the rate-limit bucket of the client address (set by the HTTP middleware)."""
+    _CLIENT_ADDRESS.set(address)
+
+
+def get_client_address() -> str | None:
+    """Rate-limit bucket of the current request's client address (None outside HTTP)."""
+    return _CLIENT_ADDRESS.get()
+
+
 def set_gitea_user_token(token: str) -> None:
     """Store the per-request Gitea OAuth user token in context-local state."""
     _GITEA_USER_TOKEN.set(token)

src/aegis_gitea_mcp/transports.py

@@ -44,7 +44,8 @@ from aegis_gitea_mcp.auth import (
 from aegis_gitea_mcp.config import Settings, is_loopback_host
 from aegis_gitea_mcp.gitea import GiteaExecutor
 from aegis_gitea_mcp.oauth import OAuthSetupError
-from aegis_gitea_mcp.rate_limit import RequestRateLimiter
+from aegis_gitea_mcp.rate_limit import RequestRateLimiter, client_bucket
+from aegis_gitea_mcp.request_context import set_client_address
 from aegis_gitea_mcp.server import (
     SERVER_NAME,
     AegisServer,
@@ -181,6 +182,8 @@ class RateLimitMiddleware:
                 break
         decision = self._limiter.check(client_ip, authorization)
         if decision.allowed:
+            # Lets the OAuth provider quota pending sign-ins and registrations per address.
+            set_client_address(client_bucket(client_ip))
             await self.app(scope, receive, send)
             return
         response = JSONResponse(

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

tests/test_oauth_provider.py

@@ -142,6 +142,7 @@ async def test_client_cap_evicts_the_least_recently_used_client_without_sessions
     await provider.register_client(client_info("newer"))
     oauth.clock.advance(5)
     await provider.get_client("old")  # used recently: "newer" is now least recently used
+    oauth.clock.advance(provider_module.PENDING_TTL)  # none of them is young any more
 
     await provider.register_client(client_info("newest"))