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

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(