190 lines
8.0 KiB
Python
190 lines
8.0 KiB
Python
"""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"
|
|
# #59: il login con PIN e' lo stesso attacco (4 cifre: ancora piu' facile da indovinare)
|
|
LOGIN_PIN_PATH = "/api/auth/login-pin"
|
|
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 in (cls.LOGIN_PATH, cls.LOGIN_PIN_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)
|