backend/src/routes/dataset-proxy.ts

import { Router, type Request, type Response } from "express";
import { Readable } from "stream";
import { extractToken } from "../auth.js";
import {
  getDatasetId,
  resolveToken,
  setUserToken,
} from "../hf-storage.js";

/**
 * Authenticated reverse proxy for the backing HF dataset.
 *
 * Why this exists
 * ---------------
 * The dataset that backs the editor is **private by default** (see
 * `hf-storage.ts::ensureDatasetExists`), but the published article's
 * `<img>` tags and the og:image / pdf links still need to be
 * resolvable from a normal browser - including by anonymous viewers
 * of a public Space. We can't just hand out
 * `huggingface.co/datasets/.../resolve/...` URLs because the dataset
 * itself is gated. Instead, the editor server proxies those reads on
 * the viewer's behalf, attaching a token it has on hand.
 *
 * Token resolution cascade (most specific to most ambient):
 *   1. Cookie from the current request (if the viewer is signed in)
 *   2. Last-known cached user token (warmed by any prior signed-in
 *      request - see `hf-storage::setUserToken`)
 *   3. The server-side `HF_TOKEN` env (when an operator set one)
 *   4. No token at all (works only for public datasets)
 *
 * Path whitelist
 * --------------
 * Only `images/` and `published/` may flow through the proxy. We
 * explicitly never expose `articles/` (raw Y.js `.yjs` snapshots
 * including unpublished drafts) and never expose the dataset root
 * either. Anything outside the whitelist gets a 404 - leaking the
 * proxy's existence is fine, leaking the path layout is not.
 *
 * Streaming
 * ---------
 * Responses are piped straight from `fetch`'s WHATWG body to the
 * Express `res`, so we never buffer a full PDF or large image in
 * Node memory. The cost is a tiny shim through `Readable.fromWeb`.
 *
 * Cache-Control
 * -------------
 * - `images/*` filenames are UUID-based and content-addressed in
 *   practice (the editor never overwrites an existing upload), so
 *   they get `immutable, max-age=1y` - browsers will skip
 *   revalidation entirely.
 * - `published/*` is mutable (every republish overwrites it) so we
 *   cap at a short max-age and tell intermediaries to revalidate.
 */
export function createDatasetProxyRouter(): Router {
  const router = Router();

  router.get(/^\/d\/(.+)$/, async (req: Request, res: Response) => {
    const datasetId = getDatasetId();
    if (!datasetId) {
      res.status(503).json({ error: "Dataset persistence not configured" });
      return;
    }

    const relPath = (req.params as { 0?: string })[0] || "";
    if (!isPathAllowed(relPath)) {
      // 404, not 403: don't confirm whether the caller "almost" had
      // the right path. An attacker probing for /d/articles/foo.yjs
      // gets the same response as one probing for /d/garbage.
      res.status(404).json({ error: "Not found" });
      return;
    }

    // Promote a cookie token to the cache opportunistically: even if
    // this very request was authenticated, future anonymous requests
    // from the same Space session benefit from a warm cache.
    const cookieToken = extractToken(req.headers.cookie);
    if (cookieToken) setUserToken(cookieToken);

    const token = resolveToken(cookieToken);
    const upstreamUrl = `https://huggingface.co/datasets/${datasetId}/resolve/main/${encodePath(relPath)}`;

    let upstream: Response | globalThis.Response;
    try {
      upstream = await fetch(upstreamUrl, {
        headers: token ? { Authorization: `Bearer ${token}` } : {},
      });
    } catch (err) {
      console.error(`[dataset-proxy] fetch failed for ${relPath}:`, (err as Error).message);
      res.status(502).json({ error: "Upstream fetch failed" });
      return;
    }

    if (!upstream.ok || !upstream.body) {
      // Don't reflect upstream 401 directly: the viewer can't do
      // anything about a missing server-side token, and surfacing
      // 401 would trigger browser auth prompts in some setups.
      // Map any non-2xx to a generic 404 + log the real status for
      // operators reading the Space logs.
      console.warn(
        `[dataset-proxy] ${relPath}: upstream status ${upstream.status} (token=${token ? "yes" : "no"})`,
      );
      res.status(upstream.status === 404 ? 404 : 502).json({
        error: upstream.status === 404 ? "Not found" : "Upstream error",
      });
      return;
    }

    // Forward content-type + content-length when present so the
    // browser renders the right viewer and shows progress for PDFs.
    const ct = upstream.headers.get("content-type");
    if (ct) res.setHeader("Content-Type", ct);
    const cl = upstream.headers.get("content-length");
    if (cl) res.setHeader("Content-Length", cl);

    res.setHeader("Cache-Control", cacheControlFor(relPath));

    // Stream the WHATWG body straight to the Express response. The
    // `as any` is required because @types/node's `Readable.fromWeb`
    // doesn't accept the DOM `ReadableStream` type without a cast.
    const nodeStream = Readable.fromWeb(upstream.body as any);
    nodeStream.on("error", (err) => {
      console.error(`[dataset-proxy] stream error for ${relPath}:`, err.message);
      if (!res.headersSent) res.status(502).end();
      else res.end();
    });
    nodeStream.pipe(res);
  });

  return router;
}

const ALLOWED_PREFIXES = ["images/", "published/"];

function isPathAllowed(path: string): boolean {
  // Reject path traversal up-front. `decodeURIComponent` would let
  // an attacker hide `..` segments behind URL encoding, but Express
  // already decodes the wildcard param for us, so a plain string
  // check is sufficient.
  if (path.includes("..") || path.includes("\0") || path.startsWith("/")) {
    return false;
  }
  return ALLOWED_PREFIXES.some((prefix) => path.startsWith(prefix));
}

/**
 * Re-encode the path before sending it upstream so spaces / unicode
 * filenames don't trip HF's URL parser. We split on `/` so the
 * slashes stay literal.
 */
function encodePath(path: string): string {
  return path.split("/").map(encodeURIComponent).join("/");
}

function cacheControlFor(path: string): string {
  if (path.startsWith("images/")) {
    // UUID names + no overwrite => safe to mark immutable.
    return "public, max-age=31536000, immutable";
  }
  // Published assets get republished in place, so cap aggressively
  // but allow shared caches to serve while revalidating.
  return "public, max-age=300, stale-while-revalidate=60";
}