From 87f048cc4d4dd8999bdbe96c7695a10361d01fc1 Mon Sep 17 00:00:00 2001 From: latte Date: Wed, 14 Oct 2026 17:32:03 +0200 Subject: [PATCH] fix: quota pending sign-ins and registrations per address, keep young clients, write the store off the loop (SEC-09) --- docs/security.md | 2 +- src/aegis_gitea_mcp/oauth/provider.py | 54 ++++++++++ src/aegis_gitea_mcp/oauth/store.py | 6 +- src/aegis_gitea_mcp/request_context.py | 11 ++ src/aegis_gitea_mcp/transports.py | 5 +- tests/security/test_oauth_input_bounds.py | 120 +++++++++++++++++++++- tests/test_oauth_provider.py | 1 + 7 files changed, 194 insertions(+), 5 deletions(-) diff --git a/docs/security.md b/docs/security.md index cdef3ee..4bf0922 100644 --- a/docs/security.md +++ b/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 diff --git a/src/aegis_gitea_mcp/oauth/provider.py b/src/aegis_gitea_mcp/oauth/provider.py index b5baf3f..ef53f08 100644 --- a/src/aegis_gitea_mcp/oauth/provider.py +++ b/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})}" diff --git a/src/aegis_gitea_mcp/oauth/store.py b/src/aegis_gitea_mcp/oauth/store.py index 9be412a..bc4587d 100644 --- a/src/aegis_gitea_mcp/oauth/store.py +++ b/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 diff --git a/src/aegis_gitea_mcp/request_context.py b/src/aegis_gitea_mcp/request_context.py index d562478..603e88e 100644 --- a/src/aegis_gitea_mcp/request_context.py +++ b/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) diff --git a/src/aegis_gitea_mcp/transports.py b/src/aegis_gitea_mcp/transports.py index 0aedab1..9d5d97b 100644 --- a/src/aegis_gitea_mcp/transports.py +++ b/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( diff --git a/tests/security/test_oauth_input_bounds.py b/tests/security/test_oauth_input_bounds.py index b13d776..a5d0890 100644 --- a/tests/security/test_oauth_input_bounds.py +++ b/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 diff --git a/tests/test_oauth_provider.py b/tests/test_oauth_provider.py index cdd0905..d8748ad 100644 --- a/tests/test_oauth_provider.py +++ b/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"))