diff --git a/src/cmcp_gateway/mcp/server.py b/src/cmcp_gateway/mcp/server.py index 9d71f8fc..24d64be4 100644 --- a/src/cmcp_gateway/mcp/server.py +++ b/src/cmcp_gateway/mcp/server.py @@ -10,10 +10,13 @@ from __future__ import annotations +import asyncio import hmac import json import logging +import time import uuid +from collections import defaultdict from typing import TYPE_CHECKING, Any from agent_os.stateless import StatelessKernel @@ -53,6 +56,48 @@ async def _unhandled_error_handler(request: Request, exc: Exception) -> Response ) +class _RateLimitMiddleware(BaseHTTPMiddleware): + """NET-002: per-IP rate limit for unauthenticated endpoints (/health). + + Uses a sliding-window counter: at most `requests_per_minute` requests + from a single IP address within any 60-second window. + """ + + def __init__( + self, + app: Any, + *, + paths: frozenset[str], + requests_per_minute: int = 60, + ) -> None: + super().__init__(app) + self._paths = paths + self._limit = requests_per_minute + self._window = 60.0 + self._counts: dict[str, list[float]] = defaultdict(list) + self._lock = asyncio.Lock() + + async def dispatch(self, request: Request, call_next: Any) -> Response: + if request.url.path not in self._paths: + return await call_next(request) + ip = request.client[0] if request.client else "unknown" + now = time.monotonic() + async with self._lock: + cutoff = now - self._window + hits = self._counts[ip] + # Prune timestamps outside the window + while hits and hits[0] <= cutoff: + hits.pop(0) + if len(hits) >= self._limit: + return JSONResponse( + {"error": "Too Many Requests", "error_code": "RATE_LIMITED"}, + status_code=429, + headers={"Retry-After": "60"}, + ) + hits.append(now) + return await call_next(request) + + class _BearerAuthMiddleware(BaseHTTPMiddleware): """AUTH-001 (CRITICAL): validate Authorization: Bearer on all protected endpoints.""" @@ -108,7 +153,14 @@ def __init__( self._max_request_bytes = max_request_bytes self._audit = audit_chain self._kernel = StatelessKernel() - middleware = ( + # NET-002: rate-limit unauthenticated /health before auth middleware runs. + # Starlette applies middleware outermost-first (first in list = first to run). + rate_limit = Middleware( + _RateLimitMiddleware, + paths=frozenset(_AUTH_EXEMPT_PATHS), + requests_per_minute=60, + ) + middleware = [rate_limit] + ( [Middleware(_BearerAuthMiddleware, bearer_token=bearer_token)] if bearer_token is not None else [] diff --git a/tests/unit/test_mcp_server_auth.py b/tests/unit/test_mcp_server_auth.py index 77e4033a..70a586db 100644 --- a/tests/unit/test_mcp_server_auth.py +++ b/tests/unit/test_mcp_server_auth.py @@ -125,6 +125,92 @@ def test_content_length_check_rejects_before_body_read(): assert resp.status_code == 413 +# ── NET-002: /health rate limit ─────────────────────────────────────────────── + +def _make_server_with_low_rate_limit(requests_per_minute: int = 3) -> "MCPServer": + """Create a server with a very low rate limit for testing.""" + from cmcp_gateway.mcp.server import _RateLimitMiddleware + from starlette.middleware import Middleware + + proxy = MagicMock() + proxy._catalog = MagicMock() + proxy._catalog.entries = {} + with patch("cmcp_gateway.mcp.server.StatelessKernel"): + server = MCPServer(proxy, bearer_token=None) + + # Replace rate-limit middleware with a tighter one for this test + from starlette.applications import Starlette + from starlette.routing import Route + + server.app = Starlette( + routes=server.app.routes, + middleware=[ + Middleware( + _RateLimitMiddleware, + paths=frozenset({"/health"}), + requests_per_minute=requests_per_minute, + ) + ], + exception_handlers={}, + ) + return server + + +def test_health_allows_requests_within_limit(): + """NET-002: requests within rate limit return 200.""" + server = _make_server_with_low_rate_limit(requests_per_minute=5) + client = TestClient(server.app, raise_server_exceptions=False) + for _ in range(3): + resp = client.get("/health") + assert resp.status_code == 200 + + +def test_health_rate_limit_returns_429_when_exceeded(): + """NET-002: exceeding rate limit returns 429 with Retry-After header.""" + server = _make_server_with_low_rate_limit(requests_per_minute=2) + client = TestClient(server.app, raise_server_exceptions=False) + + # First two should pass + assert client.get("/health").status_code == 200 + assert client.get("/health").status_code == 200 + # Third exceeds limit + resp = client.get("/health") + assert resp.status_code == 429 + assert "Retry-After" in resp.headers + body = resp.json() + assert body["error_code"] == "RATE_LIMITED" + + +def test_rate_limit_middleware_paths_only(): + """NET-002: rate limit applies only to configured paths, not all endpoints.""" + from cmcp_gateway.mcp.server import _RateLimitMiddleware + from starlette.middleware import Middleware + from starlette.applications import Starlette + + proxy = MagicMock() + proxy._catalog = MagicMock() + proxy._catalog.entries = {} + with patch("cmcp_gateway.mcp.server.StatelessKernel"): + server = MCPServer(proxy, bearer_token=None) + + # Rate-limit ONLY /nonexistent (so /health is unaffected) + server.app = Starlette( + routes=server.app.routes, + middleware=[ + Middleware( + _RateLimitMiddleware, + paths=frozenset({"/nonexistent"}), + requests_per_minute=1, + ) + ], + exception_handlers={}, + ) + client = TestClient(server.app, raise_server_exceptions=False) + for _ in range(5): + resp = client.get("/health") + assert resp.status_code == 200 + + # ── INJECT-002: sanitize method in error responses ──────────────────────────── def test_unknown_method_non_ascii_is_replaced():