Spaces:
Running
Running
| 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; | |
| } | |