"""Rate limiting middleware for FastAPI. Implements in-memory sliding window rate limiting per client IP. Configurable limits for login and endpoints that carry credentials, and for everything else. """ import re import time from collections import defaultdict from typing import Callable from fastapi import Request, Response from starlette.middleware.base import BaseHTTPMiddleware from starlette.responses import JSONResponse from src.backend.config import settings class RateLimitMiddleware(BaseHTTPMiddleware): """Middleware that enforces per-IP rate limits using a sliding window. - Login bucket: `/api/auth/login`, the paths that carry a capoturno's credentials (`/api/measurements/*/authorise`, `/api/production-runs/{id}/pause|resume|close`), and everything under `/api/setup/` (the setup password). Limited to `rate_limit_login` req/min. Guessing a supervisor's or the setup password is the same attack as guessing a login password, just aimed at a different endpoint. - All other endpoints: limited to `rate_limit_general` req/min. Returns HTTP 429 with Retry-After header when limit exceeded. Both configured limits are divided across worker processes before being applied (see `_effective_limit`): each bucket lives in this process's memory, so with `--workers N` a single configured limit is enforced N times over across the cluster unless it is divided down first. """ LOGIN_PATH = "/api/auth/login" WINDOW_SECONDS = 60 # Paths that carry a supervisor's (capoturno) credentials in the request # body. Matched by pattern, not equality, because the id in the middle # varies per request. _SUPERVISOR_CREDENTIAL_PATTERNS = ( re.compile(r"^/api/measurements/\d+/authorise$"), re.compile(r"^/api/production-runs/\d+/(?:pause|resume|close)$"), ) def __init__(self, app) -> None: super().__init__(app) # {ip: [timestamp, ...]} per bucket self._login_requests: dict[str, list[float]] = defaultdict(list) self._general_requests: dict[str, list[float]] = defaultdict(list) self._request_count = 0 # Counter for triggering eviction @staticmethod def _client_ip(request: Request) -> str: """Resolve the originating client IP behind `trusted_proxy_count` proxies. Nginx (and, behind it, the Flask client) each add one hop to X-Forwarded-For. The *leftmost* entries are whatever the original caller claimed and are trivially spoofed; trust only extends to the hops actually appended by our own trusted proxies, counted from the *right*. With `trusted_proxy_count = N`, the client IP is the Nth element from the end of XFF (the one added by the outermost trusted proxy). If XFF has fewer elements than that, or is absent, fall back to `request.client.host` - except that `X-Real-IP` is tried first when XFF is simply missing and at least one proxy is trusted. With `trusted_proxy_count = 0`, no proxy header is honored at all (neither XFF nor X-Real-IP) and only `request.client.host` is used. """ count = settings.trusted_proxy_count if count > 0: xff = request.headers.get("x-forwarded-for") if xff: hops = [h.strip() for h in xff.split(",") if h.strip()] if len(hops) >= count: return hops[-count] # Fewer hops than trusted proxies: the chain is short of what # was configured, so fall back rather than trust an # under-populated header. else: real = request.headers.get("x-real-ip") if real: return real.strip() return request.client.host if request.client else "unknown" @classmethod def _is_login_bucket_request(cls, method: str, path: str) -> bool: """Whether this request counts against the login bucket rather than general. `/api/setup/*` is checked regardless of method: `GET /api/setup/status` now accepts an optional password (see setup.py), so guessing it there is the same attack as guessing it on any of the POST endpoints. """ if path.startswith("/api/setup/"): return True if method != "POST": return False if path == cls.LOGIN_PATH: return True return any(p.match(path) for p in cls._SUPERVISOR_CREDENTIAL_PATTERNS) @staticmethod def _effective_limit(configured: int) -> int: """Divide a configured per-minute limit across worker processes. Rounds up and never returns less than 1, so a low configured limit with high concurrency still allows at least one request per worker per window rather than locking everyone out. """ concurrency = max(settings.web_concurrency, 1) return max(1, -(-configured // concurrency)) def _clean_window(self, timestamps: list[float], now: float) -> list[float]: """Remove timestamps outside the current sliding window.""" cutoff = now - self.WINDOW_SECONDS return [t for t in timestamps if t > cutoff] def _evict_stale_ips(self, bucket: dict[str, list[float]], now: float) -> None: """Remove IP entries with no timestamps in the current window (memory leak prevention).""" cutoff = now - self.WINDOW_SECONDS stale_ips = [ip for ip, timestamps in bucket.items() if not timestamps or max(timestamps) <= cutoff] for ip in stale_ips: del bucket[ip] def _check_rate_limit( self, bucket: dict[str, list[float]], client_ip: str, limit: int, now: float, ) -> tuple[bool, int]: """Check if a request is within the rate limit. Returns: Tuple of (allowed, retry_after_seconds). """ bucket[client_ip] = self._clean_window(bucket[client_ip], now) if len(bucket[client_ip]) >= limit: # Calculate seconds until the oldest request falls out of window oldest = bucket[client_ip][0] retry_after = int(oldest + self.WINDOW_SECONDS - now) + 1 return False, max(retry_after, 1) bucket[client_ip].append(now) return True, 0 async def dispatch(self, request: Request, call_next: Callable) -> Response: client_ip = self._client_ip(request) now = time.time() path = request.url.path # Periodic eviction: every 100 requests, remove stale IP buckets self._request_count += 1 if self._request_count % 100 == 0: self._evict_stale_ips(self._login_requests, now) self._evict_stale_ips(self._general_requests, now) # Check login-bucket rate limit (login itself, supervisor-credential # paths, and the setup password) if self._is_login_bucket_request(request.method, path): allowed, retry_after = self._check_rate_limit( self._login_requests, client_ip, self._effective_limit(settings.rate_limit_login), now, ) if not allowed: return JSONResponse( status_code=429, content={"detail": "Too many login attempts. Please try again later."}, headers={"Retry-After": str(retry_after)}, ) # Check general rate limit allowed, retry_after = self._check_rate_limit( self._general_requests, client_ip, self._effective_limit(settings.rate_limit_general), now, ) if not allowed: return JSONResponse( status_code=429, content={"detail": "Too many requests. Please try again later."}, headers={"Retry-After": str(retry_after)}, ) return await call_next(request)