app/http.py

"""What every page of both apps (the explorer and its admin) gets: security headers, a localhost-only rule when run
locally, same-origin state changes, and pages that revalidate so a deploy is never half-cached."""

from __future__ import annotations

import os
import threading
import time
from urllib.parse import urlsplit

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse

from . import config

LOCAL_HOSTS = {"localhost", "127.0.0.1", "[::1]", "::1"} | {h.strip() for h in os.environ.get("RLX_ALLOWED_HOSTS", "").split(",") if h.strip()}
CSP = ("default-src 'self'; script-src 'self' https://cdn.jsdelivr.net; style-src 'self' 'unsafe-inline' https://fonts.googleapis.com https://cdn.jsdelivr.net; "
       "font-src https://fonts.gstatic.com https://cdn.jsdelivr.net; img-src 'self' data: blob: https:; media-src 'self' data: blob: https:; connect-src 'self'; frame-src 'self' https://*.hf.space; object-src 'none'; "
       "base-uri 'none'; form-action 'self'; frame-ancestors 'self' https://huggingface.co https://*.hf.space")


def client_ip(request: Request) -> str:
    """The visitor's address: the last X-Forwarded-For entry is the one the Space's proxy added (earlier ones are
    whatever the client claimed)."""
    fwd = [p.strip() for p in request.headers.get("x-forwarded-for", "").split(",") if p.strip()] if config.TRUST_PROXY else []
    return fwd[-1] if fwd else (request.client.host if request.client else "?")


class Limiter:
    """At most `n` hits per `window` seconds for one key (an address, an account): a sliding window, in memory, bounded
    so a flood of distinct keys can't grow it without end. Per replica, which is what a Space runs."""

    def __init__(self, n: int, window: float, keys: int = 20000):
        self.n, self.window, self.keys = n, window, keys
        self._hits: dict[str, list[float]] = {}
        self._lock = threading.Lock()

    def hit(self, key: str) -> bool:
        """Count one hit; False when `key` is over its limit."""
        now = time.time()
        with self._lock:
            hits = [t for t in self._hits.get(key, ()) if now - t < self.window]
            if len(hits) >= self.n:
                self._hits[key] = hits
                return False
            hits.append(now)
            self._hits[key] = hits
            if len(self._hits) > self.keys:
                for k in [k for k, v in self._hits.items() if not v or now - v[-1] >= self.window] or list(self._hits)[: self.keys // 5]:
                    self._hits.pop(k, None)
            return True


def same_origin(request: Request) -> bool:
    """Compare complete origins; forwarding headers supplied by a client aren't authority."""
    origin = request.headers.get("origin")
    if not origin:
        return request.headers.get("sec-fetch-site") in (None, "same-origin", "none")

    def canonical(value: str):
        try:
            u = urlsplit(value)
            if u.scheme not in ("http", "https") or not u.hostname or u.username or u.password or u.path or u.query or u.fragment:
                return None
            return u.scheme, u.hostname.lower(), u.port or (443 if u.scheme == "https" else 80)
        except ValueError:
            return None

    expected = canonical(config.PUBLIC_URL) if config.PUBLIC_URL else canonical(f"{request.url.scheme}://{request.headers.get('host', '')}")
    return expected is not None and canonical(origin) == expected


def install(app: FastAPI) -> None:
    @app.middleware("http")
    async def _revalidate(request: Request, call_next):
        """Code and styles revalidate on every load (cheap: ETag -> 304), so a deploy is never half-cached."""
        resp = await call_next(request)
        if request.url.path.startswith(("/api/", "/mcp/", "/login")):
            # APIs can contain a visitor's private dataset, run, or account. A shared
            # proxy or the browser's cache must never reuse them for another visitor.
            resp.headers["Cache-Control"] = "private, no-store"
            resp.headers["X-Robots-Tag"] = "noindex"
        else:
            resp.headers.setdefault("Cache-Control", "no-cache")
        return resp

    @app.middleware("http")
    async def _security(request: Request, call_next):
        """Locally, only answer to localhost: a web page could otherwise DNS-rebind its own name to 127.0.0.1 and drive
        this server (and the machine's HF token) from the browser. Everywhere: the usual security headers."""
        if config.LOCAL_MODE:
            raw = request.headers.get("host") or ""
            host = raw.split("]")[0] + "]" if raw.startswith("[") else raw.rsplit(":", 1)[0]
            if host not in LOCAL_HOSTS:
                return JSONResponse({"detail": "host not allowed (set RLX_ALLOWED_HOSTS to serve under another name)"}, 403)
        resp = await call_next(request)
        resp.headers.setdefault("X-Content-Type-Options", "nosniff")
        resp.headers.setdefault("Referrer-Policy", "strict-origin-when-cross-origin")
        resp.headers.setdefault("Permissions-Policy", "camera=(), microphone=(), geolocation=(), payment=()")
        if not request.url.path.startswith("/api/"):
            resp.headers.setdefault("Content-Security-Policy", CSP)
        return resp


    @app.middleware("http")
    async def _same_origin_posts(request: Request, call_next):
        """The session cookie is SameSite=None (it has to work in the huggingface.co iframe), so a state-changing
        request must prove it came from this page: another site could otherwise start rollouts billed to you."""
        if request.method not in ("GET", "HEAD", "OPTIONS") and request.url.path.startswith("/api/"):
            if not same_origin(request):
                return JSONResponse({"detail": "cross-site request refused"}, 403)
        return await call_next(request)