back_end/security.py

"""Rate limiting, request size limits and HTTP security headers."""
import ipaddress
import json
import threading
import time
from collections import defaultdict, deque

from fastapi import HTTPException, Request

from back_end.config import FRAME_ANCESTORS, MAX_REQUEST_BODY_BYTES, TRUSTED_PROXY_HOPS


# ─────────────────────────── rate limiting ───────────────────────────

class RateLimiter:
    """In-memory sliding-window limiter keyed by client IP.

    Good enough for a single-container Space. If you ever run several replicas,
    move this to Redis or similar.
    """

    def __init__(self, max_requests: int, window_seconds: int = 60):
        self.max_requests = max_requests
        self.window = window_seconds
        self._hits: dict[str, deque] = defaultdict(deque)
        self._lock = threading.Lock()

    def check(self, key: str) -> None:
        now = time.monotonic()
        with self._lock:
            hits = self._hits[key]
            while hits and now - hits[0] > self.window:
                hits.popleft()
            if len(hits) >= self.max_requests:
                retry_after = max(1, int(self.window - (now - hits[0])))
                raise HTTPException(
                    status_code=429,
                    detail="Too many requests. Please slow down.",
                    headers={"Retry-After": str(retry_after)},
                )
            hits.append(now)
            # Opportunistic cleanup so the dict can't grow forever
            if len(self._hits) > 10_000:
                for k in [k for k, v in self._hits.items() if not v]:
                    del self._hits[k]


def client_ip(request: Request) -> str:
    """Best-effort real client IP behind the Hugging Face proxy.

    X-Forwarded-For is "client, proxy1, proxy2". Anything a client sends itself
    ends up on the LEFT, so the left-most entry is spoofable. We walk from the
    right, skip internal proxy hops (private addresses) and use the first public
    IP. Set TRUSTED_PROXY_HOPS if you know the exact number of proxies.
    """
    peer = request.client.host if request.client else "unknown"
    xff = request.headers.get("x-forwarded-for", "")
    hops = [h.strip() for h in xff.split(",") if h.strip()]
    if not hops:
        return peer
    if TRUSTED_PROXY_HOPS > 0:
        return hops[-min(TRUSTED_PROXY_HOPS, len(hops))]
    for host in reversed(hops):
        try:
            ip = ipaddress.ip_address(host)
        except ValueError:
            continue
        if ip.is_global:
            return str(ip)
    return peer


# ─────────────────────────── security headers ───────────────────────────

CONTENT_SECURITY_POLICY = "; ".join([
    "default-src 'self'",
    "script-src 'self' https://cdnjs.cloudflare.com",
    "style-src 'self' 'unsafe-inline' https://fonts.googleapis.com",
    "font-src 'self' https://fonts.gstatic.com",
    "img-src 'self' data:",
    "connect-src 'self'",
    "object-src 'none'",
    "base-uri 'none'",
    "form-action 'self'",
    f"frame-ancestors {FRAME_ANCESTORS}",
])

SECURITY_HEADERS = [
    (b"content-security-policy", CONTENT_SECURITY_POLICY.encode()),
    (b"x-content-type-options", b"nosniff"),
    (b"referrer-policy", b"no-referrer"),
    (b"permissions-policy", b"camera=(), microphone=(), geolocation=(), payment=(), usb=()"),
    (b"strict-transport-security", b"max-age=31536000; includeSubDomains"),
    (b"cross-origin-opener-policy", b"same-origin"),
    (b"cross-origin-resource-policy", b"same-origin"),
]


class SecurityHeadersMiddleware:
    """Adds security headers and enforces a request body size limit.

    Pure ASGI (not BaseHTTPMiddleware) so it doesn't interfere with streaming.
    """

    def __init__(self, app):
        self.app = app

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http":
            return await self.app(scope, receive, send)

        is_api = scope.get("path", "").startswith("/api/")

        if scope.get("method") in {"POST", "PUT", "PATCH", "DELETE"}:
            length = dict(scope.get("headers", [])).get(b"content-length")
            if length is not None:
                if not length.isdigit():
                    return await self._reject(send, 400, "Invalid Content-Length.")
                if int(length) > MAX_REQUEST_BODY_BYTES:
                    return await self._reject(send, 413, "Request too large.")
            else:
                receive = self._capped_receive(receive)

        async def send_with_headers(message):
            if message["type"] == "http.response.start":
                headers = [
                    (k, v) for k, v in message.get("headers", [])
                    if k.lower() != b"server"
                ]
                headers.extend(SECURITY_HEADERS)
                if is_api:
                    headers.append((b"cache-control", b"no-store"))
                message["headers"] = headers
            await send(message)

        await self.app(scope, receive, send_with_headers)

    @staticmethod
    def _capped_receive(receive):
        """For bodies without Content-Length (chunked): drop the request once it exceeds the cap."""
        total = 0

        async def wrapped():
            nonlocal total
            message = await receive()
            if message["type"] == "http.request":
                total += len(message.get("body", b""))
                if total > MAX_REQUEST_BODY_BYTES:
                    return {"type": "http.disconnect"}
            return message

        return wrapped

    @staticmethod
    async def _reject(send, status: int, detail: str):
        body = json.dumps({"detail": detail}).encode()
        await send({
            "type": "http.response.start",
            "status": status,
            "headers": [
                (b"content-type", b"application/json"),
                (b"content-length", str(len(body)).encode()),
                (b"cache-control", b"no-store"),
                *SECURITY_HEADERS,
            ],
        })
        await send({"type": "http.response.body", "body": body})