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})}"