proxy_gate.py

"""Standalone reverse proxy / auth gate in front of the self-reflective API.

Two modes, chosen by LIVE_UPSTREAM_PORTS:

Single-upstream (legacy) mode -- LIVE_UPSTREAM_PORTS unset:

    client -> 127.0.0.1:8180 (this)  -> 127.0.0.1:8100 (uvicorn live_app:app)

    One shared drift state. ``POST /live/rotate`` is presenter-only: answered
    on the loopback gate, refused with 403 for anything that arrives through
    a tunnel, and it advances the shared table by running drift_cron.py.

Session mode -- LIVE_UPSTREAM_PORTS is a comma list of N ports:

    client -> :8180 (this) -> one of N upstreams, each pinned at one snapshot

    Each upstream instance is a full copy of the API permanently seeded at
    snapshot version k and never reloaded (see run_live.py). A visitor's
    entire drift state is therefore one small integer -- the ring position --
    kept here against a session cookie. ``POST /live/rotate`` becomes public
    and per-session: it moves that visitor's pointer to the next version
    (``?to=N`` jumps anywhere on the ring) and no other visitor notices.
    ``GET /live/status`` is answered from the visitor's instance and then
    rewritten: rotations become the visitor's own history, and a ``session``
    block reports their position on the ring.

Responsibilities in both modes (none of which touch the upstream app):
  * everything under /admin requires ``Authorization: Bearer $LIVE_ADMIN_TOKEN``,
    otherwise 401;
  * a sliding rate limit per client on ALL paths, otherwise 429;
  * the Authorization header is stripped before forwarding, so the admin token
    never reaches the upstream process.

Standard library only -- no extra dependency beyond what the server needs.

Config (all via environment, no secrets in this file):
    LIVE_ADMIN_TOKEN        required; the proxy refuses to start without it
    LIVE_PROXY_PORT         default 8180
    LIVE_PROXY_BIND         default 127.0.0.1
    LIVE_UPSTREAM_HOST      default 127.0.0.1
    LIVE_UPSTREAM_PORT      default 8100 (single-upstream mode)
    LIVE_UPSTREAM_PORTS     comma list; two or more ports switch session mode on
    LIVE_RATE_LIMIT_PER_MIN default 60
    LIVE_SESSION_TTL        default 14400 s idle before a session is dropped
    LIVE_SESSION_MAX        default 2000 concurrent sessions kept
    LIVE_SNAPSHOT_DIR       default ./snapshots (session-mode diffs)
    LIVE_DRIFT_TABLE        default incompatible_ingredients
    LIVE_DRIFT_FILE         default incompatible_combinations

Run:
    python proxy_gate.py
"""
from __future__ import annotations

import hmac
import http.client
import json
import os
import secrets
import subprocess
import sys
import threading
import time
from collections import defaultdict, deque
from datetime import datetime, timezone
from http.cookies import SimpleCookie
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.parse import parse_qs, urlparse

# The startup banner may print filesystem paths, which contain non-ASCII on
# some hosts; a cp1252 console would raise UnicodeEncodeError and abort.
for _stream in (sys.stdout, sys.stderr):
    try:
        _stream.reconfigure(encoding="utf-8", errors="replace")
    except (AttributeError, ValueError):  # pragma: no cover
        pass

RESERVED_PORT = 8000

ADMIN_TOKEN = os.environ.get("LIVE_ADMIN_TOKEN", "")
PROXY_BIND = os.environ.get("LIVE_PROXY_BIND", "127.0.0.1")
PROXY_PORT = int(os.environ.get("LIVE_PROXY_PORT", "8180"))
UPSTREAM_HOST = os.environ.get("LIVE_UPSTREAM_HOST", "127.0.0.1")
UPSTREAM_PORT = int(os.environ.get("LIVE_UPSTREAM_PORT", "8100"))
UPSTREAM_PORTS = [int(p) for p in
                  os.environ.get("LIVE_UPSTREAM_PORTS", "").replace(",", " ").split()]
SESSION_MODE = len(UPSTREAM_PORTS) >= 2
RATE_LIMIT = int(os.environ.get("LIVE_RATE_LIMIT_PER_MIN", "60"))
WINDOW_SECONDS = 60.0
ADMIN_PREFIX = "/admin"
ROTATE_PATH = "/live/rotate"

HERE = Path(__file__).resolve().parent
DRIFT_SCRIPT = HERE / "drift_cron.py"
SEED_INDEX = int(os.environ.get("LIVE_SEED_INDEX", "3"))

SESSION_COOKIE = "live_sid"
COOKIE_SAMESITE = os.environ.get("LIVE_COOKIE_SAMESITE", "auto").strip() or "auto"
SESSION_TTL = float(os.environ.get("LIVE_SESSION_TTL", "14400"))
SESSION_MAX = int(os.environ.get("LIVE_SESSION_MAX", "2000"))
SNAPSHOT_DIR = Path(os.environ.get("LIVE_SNAPSHOT_DIR", HERE / "snapshots")).resolve()
TABLE_NAME = os.environ.get("LIVE_DRIFT_TABLE", "incompatible_ingredients")
FILE_STEM = os.environ.get("LIVE_DRIFT_FILE", "incompatible_combinations")

# In legacy mode a request that arrives through a tunnel is recognisable by the
# forwarding headers the relay stamps on it; the presenter check uses that.
FORWARDED_MARKERS = (
    "cf-connecting-ip", "cf-ray", "cf-ipcountry", "cf-visitor", "cf-warp-tag-id",
    "x-forwarded-for", "x-forwarded-host", "x-real-ip",
)
_rotate_lock = threading.Lock()

# Headers that are connection-specific and must not be relayed.
HOP_BY_HOP = {
    "connection", "keep-alive", "proxy-authenticate", "proxy-authorization",
    "te", "trailer", "trailers", "transfer-encoding", "upgrade", "host",
    "content-length",
}


class SlidingWindowLimiter:
    """~N requests per 60s per key, kept entirely in memory (single process)."""

    def __init__(self, limit: int, window: float = WINDOW_SECONDS) -> None:
        self.limit = limit
        self.window = window
        self._hits: dict[str, deque[float]] = defaultdict(deque)
        self._lock = threading.Lock()

    def check(self, key: str) -> tuple[bool, int]:
        """Return (allowed, retry_after_seconds)."""
        now = time.monotonic()
        with self._lock:
            q = self._hits[key]
            cutoff = now - self.window
            while q and q[0] <= cutoff:
                q.popleft()
            if len(q) >= self.limit:
                return False, max(1, int(q[0] + self.window - now) + 1)
            q.append(now)
            return True, 0


limiter = SlidingWindowLimiter(RATE_LIMIT)


# --------------------------------------------------------------------------
# Session mode: ring snapshots, per-transition diffs, and the session store
# --------------------------------------------------------------------------

def _load_ring() -> dict[int, dict]:
    """Read the snapshot ring once; {index: data dict}."""
    ring: dict[int, dict] = {}
    for path in sorted(SNAPSHOT_DIR.glob(FILE_STEM + "_v*.json")):
        try:
            idx = int(path.stem.rsplit("_v", 1)[1])
            blob = json.loads(path.read_text(encoding="utf-8"))
        except (ValueError, OSError, json.JSONDecodeError):
            continue
        ring[idx] = blob.get("data") or {}
    return ring


def _diff_entries(before: dict, after: dict) -> dict:
    """Same shape live_app.py and the upstream reload endpoint produce."""
    changed: dict = {}
    for key in set(before) | set(after):
        old, new = before.get(key), after.get(key)
        if old == new:
            continue
        if isinstance(old, list) or isinstance(new, list):
            old_rows, new_rows = old or [], new or []
            changed[key] = {
                "removed": [r for r in old_rows if r not in new_rows],
                "added": [r for r in new_rows if r not in old_rows],
            }
        else:
            changed[key] = {"old": old, "new": new}
    return changed


_RING_DATA = _load_ring() if SESSION_MODE else {}
_DIFF_CACHE: dict[tuple[int, int], dict] = {}


def _ring_diff(a: int, b: int) -> dict:
    if (a, b) not in _DIFF_CACHE:
        _DIFF_CACHE[(a, b)] = _diff_entries(
            _RING_DATA.get(a) or {}, _RING_DATA.get(b) or {})
    return _DIFF_CACHE[(a, b)]


def _iso_now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="seconds").replace("+00:00", "Z")


class SessionStore:
    """sid -> {pos, hist, seen}. TTL-swept, size-capped, thread-safe."""

    def __init__(self, size: int, ttl: float, cap: int) -> None:
        self.size = size
        self.ttl = ttl
        self.cap = cap
        self._d: dict[str, dict] = {}
        self._lock = threading.Lock()
        self._last_sweep = 0.0

    def touch(self, sid: str | None) -> tuple[str, dict, bool]:
        """Return (sid, session, created). Unknown or absent sid makes a new one."""
        now = time.monotonic()
        with self._lock:
            if now - self._last_sweep > 60:
                self._last_sweep = now
                dead = [k for k, v in self._d.items() if now - v["seen"] > self.ttl]
                for k in dead:
                    del self._d[k]
            if sid and sid in self._d:
                sess = self._d[sid]
                sess["seen"] = now
                return sid, sess, False
            if len(self._d) >= self.cap:
                oldest = min(self._d, key=lambda k: self._d[k]["seen"])
                del self._d[oldest]
            sid = secrets.token_urlsafe(16)
            sess = {"pos": 1, "hist": [], "seen": now}
            self._d[sid] = sess
            return sid, sess, True

    def count(self) -> int:
        with self._lock:
            return len(self._d)


SESSIONS = SessionStore(len(UPSTREAM_PORTS) or 1, SESSION_TTL, SESSION_MAX)
_VISITOR_CALLS = {"n": 0}
_visitor_lock = threading.Lock()


def _token_ok(header_value: str | None) -> bool:
    if not header_value:
        return False
    parts = header_value.split(None, 1)
    if len(parts) != 2 or parts[0].lower() != "bearer":
        return False
    return hmac.compare_digest(parts[1].strip(), ADMIN_TOKEN)


class GateHandler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"
    server_version = "live-gate/2.0"
    sys_version = ""

    # ---------- helpers ----------

    def _client_ip(self) -> str:
        return self.client_address[0] if self.client_address else "unknown"

    def _rate_key(self) -> str:
        """Behind a platform proxy every socket is the proxy; key on the first
        X-Forwarded-For hop there, and on the socket IP when direct."""
        if SESSION_MODE:
            xff = self.headers.get("X-Forwarded-For")
            if xff:
                return xff.split(",")[0].strip() or self._client_ip()
        return self._client_ip()

    def _cookie_sid(self) -> str | None:
        raw = self.headers.get("Cookie")
        if not raw:
            return None
        jar = SimpleCookie()
        try:
            jar.load(raw)
        except Exception:
            return None
        morsel = jar.get(SESSION_COOKIE)
        return morsel.value if morsel else None

    def _send_json(self, status: int, payload: dict, extra: dict | None = None) -> None:
        body = json.dumps(payload).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        for k, v in (extra or {}).items():
            self.send_header(k, v)
        self.end_headers()
        if self.command != "HEAD":
            self.wfile.write(body)

    def _read_body(self) -> bytes:
        try:
            length = int(self.headers.get("Content-Length") or 0)
        except ValueError:
            length = 0
        return self.rfile.read(length) if length > 0 else b""

    def _cookie_samesite(self) -> str:
        """SameSite policy for this request.

        A page embedded in a cross-site iframe (the Space view on
        huggingface.co) only keeps its cookie under SameSite=None, which in
        turn requires Secure and therefore HTTPS. Plain-HTTP deployments must
        stay on Lax or the cookie is stored but never sent back. "auto" picks
        per request from the TLS evidence in the headers.
        """
        if COOKIE_SAMESITE.lower() != "auto":
            return COOKIE_SAMESITE
        proto = (self.headers.get("X-Forwarded-Proto") or "").split(",")[0].strip().lower()
        host = (self.headers.get("Host") or "").split(":")[0].strip().lower()
        if proto == "https" or host.endswith(".hf.space"):
            return "None"
        return "Lax"

    def _set_cookie_header(self, sid: str) -> tuple[str, str]:
        samesite = self._cookie_samesite()
        secure = "; Secure" if samesite.lower() == "none" else ""
        return ("Set-Cookie",
                f"{SESSION_COOKIE}={sid}; Path=/; Max-Age={int(SESSION_TTL)}; "
                f"SameSite={samesite}{secure}; HttpOnly")

    def _forward(self, port: int) -> tuple[int, list, bytes] | None:
        """Relay the current request to an upstream; None on connect failure."""
        body = self._read_body()
        fwd_headers = {
            k: v for k, v in self.headers.items()
            if k.lower() not in HOP_BY_HOP and k.lower() != "authorization"
        }
        fwd_headers["Host"] = f"{UPSTREAM_HOST}:{port}"
        prior = self.headers.get("X-Forwarded-For")
        ip = self._client_ip()
        fwd_headers["X-Forwarded-For"] = f"{prior}, {ip}" if prior else ip
        fwd_headers["X-Forwarded-Proto"] = "http"
        if body:
            fwd_headers["Content-Length"] = str(len(body))
        try:
            conn = http.client.HTTPConnection(UPSTREAM_HOST, port, timeout=30)
            conn.request(self.command, self.path, body=body or None, headers=fwd_headers)
            resp = conn.getresponse()
            payload = resp.read()
            status = resp.status
            out_headers = [
                (k, v) for k, v in resp.getheaders()
                if k.lower() not in HOP_BY_HOP
            ]
            conn.close()
        except OSError as exc:
            self._send_json(502, {"error": "bad_gateway",
                                  "detail": f"upstream {UPSTREAM_HOST}:{port}: {exc}"})
            return None
        return status, out_headers, payload

    def _relay(self, status: int, headers: list, payload: bytes,
               extra: list | None = None) -> None:
        self.send_response(status)
        for k, v in headers:
            self.send_header(k, v)
        for k, v in (extra or []):
            self.send_header(k, v)
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        if self.command != "HEAD":
            self.wfile.write(payload)

    # ---------- session mode ----------

    def _session_rotate(self, sess: dict, cookie: tuple[str, str] | None) -> None:
        size = len(UPSTREAM_PORTS)
        want = parse_qs(urlparse(self.path).query).get("to", [None])[0]
        with _rotate_lock:
            src = sess["pos"]
            if want and want.isdigit():
                dst = int(want)
                if not 1 <= dst <= size:
                    self._send_json(400, {"error": "bad_target",
                                          "detail": f"to must be 1..{size}"},
                                    dict([cookie]) if cookie else None)
                    return
            else:
                dst = src % size + 1
            if dst != src:
                sess["pos"] = dst
                sess["hist"].insert(0, {
                    "ts": _iso_now(), "table": TABLE_NAME,
                    "from": f"{src}.0.0", "to": f"{dst}.0.0",
                    "changed": _ring_diff(src, dst),
                })
                del sess["hist"][12:]
        self._send_json(200, {
            "ok": True,
            "version": f"{dst}.0.0",
            "position": dst,
            "detail": f"{src}.0.0 -> {dst}.0.0" if dst != src else "already there",
        }, dict([cookie]) if cookie else None)

    def _session_status(self, sess: dict, port: int,
                        cookie: tuple[str, str] | None) -> None:
        got = self._forward(port)
        if got is None:
            return
        status, headers, payload = got
        if status == 200:
            try:
                doc = json.loads(payload)
            except (ValueError, UnicodeDecodeError):
                doc = None
            if isinstance(doc, dict):
                doc["rotations"] = sess["hist"]
                doc["visitor_requests"] = _VISITOR_CALLS["n"]
                doc["session"] = {"position": sess["pos"], "size": len(UPSTREAM_PORTS)}
                doc["sessions_active"] = SESSIONS.count()
                drift = doc.get("drift")
                if isinstance(drift, dict):
                    drift["schedule"] = "on demand - each visitor drifts an isolated session"
                    drift["reload_pending"] = False
                payload = json.dumps(doc).encode("utf-8")
                headers = [(k, v) for k, v in headers
                           if k.lower() != "content-type"] + \
                          [("Content-Type", "application/json")]
        self._relay(status, headers, payload, [cookie] if cookie else None)

    def _handle_session(self) -> None:
        path = self.path.split("?", 1)[0]
        sid, sess, created = SESSIONS.touch(self._cookie_sid())
        cookie = self._set_cookie_header(sid) if created else None
        port = UPSTREAM_PORTS[sess["pos"] - 1]

        if path == ROTATE_PATH:
            if self.command in ("POST", "GET"):
                self._session_rotate(sess, cookie)
            else:
                self._send_json(405, {"error": "method_not_allowed"})
            return

        if path.startswith(ADMIN_PREFIX):
            if not _token_ok(self.headers.get("Authorization")):
                self._send_json(
                    401,
                    {"error": "unauthorized",
                     "detail": "admin routes require 'Authorization: Bearer <LIVE_ADMIN_TOKEN>'"},
                    {"WWW-Authenticate": 'Bearer realm="live-admin"'},
                )
                return

        if path.startswith("/api/") and "x-live-probe" not in self.headers:
            with _visitor_lock:
                _VISITOR_CALLS["n"] += 1

        if path == "/live/status":
            self._session_status(sess, port, cookie)
            return

        got = self._forward(port)
        if got is None:
            return
        status, headers, payload = got
        self._relay(status, headers, payload, [cookie] if cookie else None)

    # ---------- legacy presenter-only rotate ----------

    def _is_presenter(self) -> bool:
        if any(self.headers.get(h) for h in FORWARDED_MARKERS):
            return False
        host = (self.headers.get("Host") or "").rsplit(":", 1)[0].strip().lower()
        if host not in ("127.0.0.1", "localhost", "[::1]", "::1"):
            return False
        return self._client_ip() in ("127.0.0.1", "::1")

    def _upstream_major(self) -> int | None:
        try:
            conn = http.client.HTTPConnection(UPSTREAM_HOST, UPSTREAM_PORT, timeout=10)
            conn.request("GET", "/live/status")
            payload = json.loads(conn.getresponse().read())
            conn.close()
        except (OSError, ValueError):
            return None
        table = (payload.get("drift") or {}).get("table")
        version = (payload.get("table_versions") or {}).get(table)
        try:
            return int(str(version).split(".")[0])
        except (TypeError, ValueError):
            return None

    def _local_rotate(self) -> None:
        if not self._is_presenter():
            self._send_json(403, {"error": "forbidden",
                                  "detail": "rotate is served on the loopback gate only"})
            return
        want = parse_qs(urlparse(self.path).query).get("to", [None])[0]
        with _rotate_lock:
            argv = [sys.executable, str(DRIFT_SCRIPT)]
            if want and want.isdigit():
                argv += ["--to", want]
            else:
                major = self._upstream_major()
                if major is not None and major > SEED_INDEX:
                    argv += ["--to", str(SEED_INDEX)]
            try:
                proc = subprocess.run(argv, cwd=str(HERE), capture_output=True,
                                      text=True, timeout=90)
            except (OSError, subprocess.SubprocessError) as exc:
                self._send_json(500, {"error": "rotate_failed", "detail": str(exc)})
                return
            after = self._upstream_major()
        tail = [ln for ln in (proc.stdout or "").splitlines() if "->" in ln][-1:]
        self._send_json(200 if proc.returncode == 0 else 500, {
            "ok": proc.returncode == 0,
            "version": f"{after}.0.0" if after is not None else None,
            "detail": tail[0].strip() if tail else (proc.stderr or "").strip()[-200:],
        })

    def _handle_legacy(self) -> None:
        if self.path.split("?", 1)[0] == ROTATE_PATH:
            if self.command in ("POST", "GET"):
                self._local_rotate()
            else:
                self._send_json(405, {"error": "method_not_allowed"})
            return

        if self.path.split("?", 1)[0].startswith(ADMIN_PREFIX):
            if not _token_ok(self.headers.get("Authorization")):
                self._send_json(
                    401,
                    {"error": "unauthorized",
                     "detail": "admin routes require 'Authorization: Bearer <LIVE_ADMIN_TOKEN>'"},
                    {"WWW-Authenticate": 'Bearer realm="live-admin"'},
                )
                return

        got = self._forward(UPSTREAM_PORT)
        if got is None:
            return
        status, headers, payload = got
        self._relay(status, headers, payload)

    # ---------- core ----------

    def _handle(self) -> None:
        allowed, retry_after = limiter.check(self._rate_key())
        if not allowed:
            self._send_json(
                429,
                {"error": "rate_limited",
                 "detail": f"more than {RATE_LIMIT} requests/min from this address"},
                {"Retry-After": str(retry_after)},
            )
            return
        if self.command in ("GET", "HEAD") \
                and self.path.split("?", 1)[0] == "/" \
                and "text/html" in (self.headers.get("Accept") or ""):
            self.send_response(302)
            self.send_header("Location", "/live")
            self.send_header("Content-Length", "0")
            self.end_headers()
            return
        if SESSION_MODE:
            self._handle_session()
        else:
            self._handle_legacy()

    do_GET = do_POST = do_PUT = do_PATCH = do_DELETE = do_HEAD = do_OPTIONS = _handle

    def log_message(self, fmt: str, *args) -> None:  # quieter, one line per request
        sys.stderr.write(f"[gate] {self._client_ip()} {fmt % args}\n")


def main() -> None:
    if not ADMIN_TOKEN:
        raise SystemExit(
            "LIVE_ADMIN_TOKEN is not set. The gate refuses to start without it "
            "(fail closed). See .env.example."
        )
    if PROXY_PORT == RESERVED_PORT or UPSTREAM_PORT == RESERVED_PORT \
            or RESERVED_PORT in UPSTREAM_PORTS:
        raise SystemExit("port 8000 is reserved on this host")

    srv = ThreadingHTTPServer((PROXY_BIND, PROXY_PORT), GateHandler)
    srv.daemon_threads = True
    print(f"[gate] listening on http://{PROXY_BIND}:{PROXY_PORT}")
    if SESSION_MODE:
        print(f"[gate] session mode: {len(UPSTREAM_PORTS)} pinned upstreams "
              f"{UPSTREAM_PORTS} on {UPSTREAM_HOST}")
        print(f"[gate] ring diffs from {SNAPSHOT_DIR} "
              f"({len(_RING_DATA)} snapshots of {FILE_STEM})")
    else:
        print(f"[gate] forwarding to http://{UPSTREAM_HOST}:{UPSTREAM_PORT}")
    print(f"[gate] admin prefix {ADMIN_PREFIX} requires a bearer token")
    print(f"[gate] rate limit {RATE_LIMIT} req/min per client")
    try:
        srv.serve_forever()
    except KeyboardInterrupt:
        print("\n[gate] shutting down")
    finally:
        srv.server_close()


if __name__ == "__main__":
    main()