Spaces:
Paused
Paused
| """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), | |
| ]) | |