Blame for src/colab-worker.ts

import { timingSafeEqual } from "node:crypto";
import { Router, type Request } from "express";
import { z } from "zod";

const registrationSchema = z.object({ active: z.boolean().default(true), url: z.string().url().optional() }).strict();
const HEARTBEAT_TTL_MS = 75_000;

let worker: { url: string; lastSeenAt: number } | null = null;

function secretMatches(expected: string | undefined, received: string | undefined) {
  if (!expected || !received) return false;
  const expectedBytes = Buffer.from(expected);
  const receivedBytes = Buffer.from(received);
  return expectedBytes.length === receivedBytes.length && timingSafeEqual(expectedBytes, receivedBytes);
}

function validTunnelUrl(value: string) {
  try {
    const url = new URL(value);
    return url.protocol === "https:" && url.hostname.endsWith(".trycloudflare.com") &&
      !url.username && !url.password && url.pathname === "/" && !url.search && !url.hash;
  } catch {
    return false;
  }
}

export function colabWorkerRouter(analysisSecret: string, workerSecret?: string) {
  const router = Router();

  // The Colab notebook refreshes this lease while its tunnel is alive.
  router.post("/register", (req, res) => {
    if (!secretMatches(workerSecret, req.header("x-sporalize-colab-secret"))) {
      return res.status(401).json({ error: "Invalid worker registration secret." });
    }
    const parsed = registrationSchema.safeParse(req.body);
    if (!parsed.success) return res.status(400).json({ error: "Invalid worker registration." });
    if (!parsed.data.active) {
      worker = null;
      return res.json({ active: false });
    }
    if (!parsed.data.url || !validTunnelUrl(parsed.data.url)) {
      return res.status(400).json({ error: "Worker URL must be an HTTPS Cloudflare quick tunnel URL." });
    }
    worker = { url: parsed.data.url.replace(/\/$/, ""), lastSeenAt: Date.now() };
    return res.json({ active: true, expiresInSeconds: HEARTBEAT_TTL_MS / 1000 });
  });

  // Only the private Python service can ask Node for the current worker lease.
  router.get("/current", (req: Request, res) => {
    if (!secretMatches(analysisSecret, req.header("x-sporalize-proxy-secret"))) {
      return res.status(401).json({ error: "Unauthorized." });
    }
    if (!worker || Date.now() - worker.lastSeenAt > HEARTBEAT_TTL_MS) {
      worker = null;
      return res.json({ active: false });
    }
    return res.json({ active: true, url: worker.url });
  });

  return router;
}