Blame for app.py

"""Pi Web direct previews: a thin streaming relay on its own site.

Pi Web (https://mike0021-pi-web.hf.space) frames this Space to show a user's own dev server
(a Vite app on their machine's port 5173, say) natively. This origin is shared by every user
and holds nothing of value: each browser only ever holds its own user's capability, in an
HttpOnly cookie. The gateway checks the capability and resolves it to that user's machine.

  GET /__pv/<cap>/<rest>   handshake: cookie __pv=<cap>, 302 to /<rest> (the cap leaves the URL)
  anything else            needs the cookie; forwarded to GATEWAY/_hf/preview-relay/<path>?<query>
                           with x-hfpi-preview-relay (shared secret) and x-preview-cap; the app's own
                           cookies travel in x-preview-cookie. Bodies stream both ways (SSE too).
                           WebSocket upgrades are not supported.
"""

import base64
import http.cookiejar
import json
import os
import re
import time
import urllib.parse

import httpx
from starlette.applications import Starlette
from starlette.background import BackgroundTask
from starlette.requests import Request
from starlette.responses import HTMLResponse, RedirectResponse, Response, StreamingResponse
from starlette.routing import Route

SECRET = os.environ.get("PREVIEW_RELAY_SECRET", "").strip()
GATEWAY = os.environ.get("PREVIEW_GATEWAY", "https://mike0021-pi-web.hf.space").strip().rstrip("/")
# Who may frame previews: Pi Web, and huggingface.co, which frames Pi Web on its Space page
# (frame-ancestors checks every ancestor, not only the parent).
FRAME_ANCESTORS = os.environ.get("PREVIEW_FRAME_ANCESTORS", f"{GATEWAY} https://huggingface.co").split()

COOKIE = "__pv"
RELAY_HEADER = "x-hfpi-preview-relay"
CAP_HEADER = "x-preview-cap"
COOKIE_HEADER = "x-preview-cookie"
MAX_TTL = 12 * 3600
CAP_RE = re.compile(r"[A-Za-z0-9_-]{8,400}\.[A-Za-z0-9_-]{43}")

HOP = {"connection", "keep-alive", "proxy-authenticate", "proxy-authorization", "te", "trailer", "trailers", "transfer-encoding", "upgrade"}
DROP_REQUEST = HOP | {"host", "content-length", "cookie", "authorization", RELAY_HEADER, CAP_HEADER, COOKIE_HEADER,
                      "forwarded", "x-real-ip", "x-ip-token", "x-request-id", "x-amzn-trace-id"}
DROP_RESPONSE = HOP | {"content-length", "x-frame-options", "referrer-policy", "set-cookie", "content-security-policy"}

# Shared by every user: it must never keep cookies (httpx's jar would replay one user's app
# cookies on everyone's later requests).
_NO_COOKIES = http.cookiejar.CookieJar(policy=http.cookiejar.DefaultCookiePolicy(allowed_domains=[]))
_client = httpx.AsyncClient(
    cookies=_NO_COOKIES,
    timeout=httpx.Timeout(connect=15, read=None, write=120, pool=60),
    limits=httpx.Limits(max_connections=1000, max_keepalive_connections=100),
)


def _page(status: int, title: str, text: str) -> HTMLResponse:
    html = (f'<!doctype html><meta charset="utf-8"><meta name="viewport" content="width=device-width"><title>{title}</title>'
            '<style>body{font:14px/1.5 system-ui,sans-serif;margin:0;display:grid;place-items:center;min-height:100vh;'
            'background:#111;color:#ddd}@media (prefers-color-scheme:light){body{background:#fff;color:#222}}'
            f'main{{max-width:420px;padding:24px}}h1{{font-size:16px;margin:0 0 8px}}</style><main><h1>{title}</h1><p>{text}</p></main>')
    return _harden(HTMLResponse(html, status_code=status, headers={"cache-control": "no-store"}))


def _harden(response: Response, policies: list[str] = ()) -> Response:
    # The app's own policies without frame-ancestors, plus ours (every CSP header is enforced).
    for policy in policies:
        kept = [d.strip() for d in policy.split(";") if d.strip() and not d.strip().lower().startswith("frame-ancestors")]
        if kept:
            response.headers.append("content-security-policy", "; ".join(kept))
    response.headers.append("content-security-policy", "frame-ancestors " + " ".join(FRAME_ANCESTORS))
    # The app's own requests keep their Origin (no-referrer made a form POST's Origin "null", which
    # CSRF checks refuse); nothing of the preview's URLs reaches other sites.
    response.headers["referrer-policy"] = "same-origin"
    # Every answer of this relay says so: a 5xx without it is Hugging Face's own, for a preview Space
    # that is not serving (the gateway's link check tells the user that, not the app's error).
    response.headers["x-pi-preview"] = "relay"
    return response


def _ttl(cap: str) -> int:
    """Cookie lifetime from the cap's (unverified) expiry: the gateway verifies it on every request."""
    try:
        body = cap.split(".")[0]
        expires = int(json.loads(base64.urlsafe_b64decode(body + "=" * (-len(body) % 4)))["e"])
    except (ValueError, KeyError, TypeError):
        return 0
    return max(0, min(MAX_TTL, expires - int(time.time())))


def rewrite_set_cookie(value: str) -> str | None:
    """An app cookie as a third-party iframe can keep it: Secure, SameSite=None, Partitioned, no Domain."""
    parts = [p.strip() for p in value.split(";")]
    name = parts[0].split("=", 1)[0].strip()
    if not name or name == COOKIE:
        return None  # the app never sets (or clears) the preview's own cookie
    attrs = [p for p in parts[1:] if p and p.split("=", 1)[0].strip().lower() not in ("domain", "secure", "samesite", "partitioned")]
    return "; ".join([parts[0], *attrs, "Secure", "SameSite=None", "Partitioned"])


async def handshake(request: Request) -> Response:
    cap = request.path_params["cap"]
    if not CAP_RE.fullmatch(cap) or not (ttl := _ttl(cap)):
        return _page(400, "Preview link invalid", "This preview link is malformed or expired. Reopen the preview from Pi Web.")
    raw = (request.scope.get("raw_path") or request.url.path.encode()).decode("latin-1").split("?", 1)[0]
    rest = raw.split("/", 3)[3] if raw.count("/") >= 3 else ""  # "/__pv/<cap>/<rest>", as encoded
    target = "/" + rest.lstrip("/\\")  # root-relative, never "//host"
    if not _safe(target.lstrip("/")):
        target = "/"
    if request.url.query:
        target += "?" + request.url.query
    response = RedirectResponse(target, status_code=302, headers={"cache-control": "no-store"})
    response.headers.append("set-cookie", f"{COOKIE}={cap}; Path=/; Max-Age={ttl}; Secure; HttpOnly; SameSite=None; Partitioned")
    return _harden(response)


def _safe(path: str) -> bool:
    for form in (path, urllib.parse.unquote(path), urllib.parse.unquote(urllib.parse.unquote(path))):
        if any(seg in (".", "..") for seg in form.replace("\\", "/").split("/")):
            return False
    return bool(re.fullmatch(r"[\x21-\x7e]*", path))


async def relay(request: Request) -> Response:
    cap = request.cookies.get(COOKIE)
    if not cap or not CAP_RE.fullmatch(cap):
        return _page(401, "Preview closed", "Reopen the preview from Pi Web.")
    if not SECRET:
        return _page(503, "Preview not configured", "This preview relay has no secret yet.")
    raw = (request.scope.get("raw_path") or request.url.path.encode()).decode("latin-1").split("?", 1)[0]
    path = raw.lstrip("/")
    if not _safe(path):
        return _page(400, "Bad path", "Dot segments are not allowed in preview paths.")
    url = f"{GATEWAY}/_hf/preview-relay/{path}"
    if request.url.query:
        url += "?" + request.url.query

    headers = [(k, v) for k, v in request.headers.items()
               if k.lower() not in DROP_REQUEST and not k.lower().startswith("x-forwarded-")]
    if "accept-encoding" not in request.headers:
        headers.append(("accept-encoding", "identity"))
    app_cookies = [c.strip() for c in request.headers.get("cookie", "").split(";")
                   if c.strip() and c.strip().split("=", 1)[0].strip() != COOKIE]
    if app_cookies:
        headers.append((COOKIE_HEADER, "; ".join(app_cookies)))
    headers += [(RELAY_HEADER, SECRET), (CAP_HEADER, cap)]

    has_body = request.method not in ("GET", "HEAD", "OPTIONS") or "content-length" in request.headers or "transfer-encoding" in request.headers
    try:
        upstream = await _client.send(
            _client.build_request(request.method, url, headers=headers, content=request.stream() if has_body else None), stream=True)
    except httpx.HTTPError as e:
        return _page(502, "Pi Web unreachable", f"The preview relay could not reach Pi Web ({e.__class__.__name__}).")

    response = StreamingResponse(upstream.aiter_raw(), status_code=upstream.status_code, background=BackgroundTask(upstream.aclose))
    for k, v in upstream.headers.multi_items():
        k_low = k.lower()
        if k_low == "set-cookie":
            if (cookie := rewrite_set_cookie(v)) is not None:
                response.headers.append("set-cookie", cookie)
        elif k_low not in DROP_RESPONSE:
            response.headers.append(k, v)
    if upstream.headers.get("content-type", "").startswith("text/event-stream"):
        response.headers["x-accel-buffering"] = "no"
    # One origin for every preview: cached responses must not cross caps (another port, another user).
    response.headers.append("vary", "Cookie")
    return _harden(response, upstream.headers.get_list("content-security-policy"))


async def healthz(request: Request) -> Response:
    return Response('{"ok":true}', media_type="application/json")


METHODS = ["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"]
app = Starlette(routes=[
    Route("/__pv/healthz", healthz),
    Route("/__pv/{cap}", handshake, methods=["GET", "HEAD"]),
    Route("/__pv/{cap}/{rest:path}", handshake, methods=["GET", "HEAD"]),
    Route("/{path:path}", relay, methods=METHODS),
])