"""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// handshake: cookie __pv=, 302 to / (the cap leaves the URL) anything else needs the cookie; forwarded to GATEWAY/_hf/preview-relay/? 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'{title}' '

{title}

{text}

') 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//", 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), ])