1import * as NodeHttp from "node:http";
2import type { IncomingMessage } from "node:http";
3import type { Duplex } from "node:stream";
4import type * as vite from "vite";
5import { proxyRequestHeaders, resolveForwardedHost } from "./forwarded-host.ts";
6
7/**
8 * Handles 'upgrade' requests on the Vite HTTP server and forwards the
9 * WebSocket handshake to the local workerd address as a raw HTTP upgrade.
10 *
11 * Returns a cleanup function that removes the listener (used on server restart).
12 */
13export function handleWebSocket(
14 httpServer: vite.HttpServer,
15 address: string | URL,
16 proxySharedSecret: string,
17): () => void {
18 const upstreamBase = typeof address === "string" ? new URL(address) : address;
19
20 // Sockets hijacked by an `upgrade` are not reaped by `server.closeAllConnections()`,
21 // yet `server.close()` still waits on them — so a lingering proxied WebSocket blocks
22 // the HTTP server from closing on restart. Track live sockets and destroy them in the
23 // cleanup function to close deterministically.
24 const sockets = new Set<Duplex>();
25 const track = (socket: Duplex) => {
26 sockets.add(socket);
27 socket.on("close", () => sockets.delete(socket));
28 };
29
30 const onUpgrade = (
31 request: IncomingMessage,
32 socket: Duplex,
33 head: Buffer,
34 ) => {
35 // Unhandled socket errors crash Node.
36 socket.on("error", () => socket.destroy());
37
38 // The URL — and thus the Sandbox-origin check below — is built from the
39 // resolved (forwarded) host, not the raw `Host` header. This diverges from
40 // upstream, which keys the origin off `Host`; here it's intentional so a
41 // tunnel-fronted Sandbox preview still matches. Direct Sandbox hits carry no
42 // `X-Forwarded-Host`, so they fall back to `Host` and behave identically.
43 const rawHost = resolveForwardedHost(request.headers, "localhost");
44 const base = /^https?:\/\//i.test(rawHost) ? rawHost : `http://${rawHost}`;
45 const url = new URL(request.url ?? "/", base);
46
47 const isViteRequest =
48 request.headers["sec-websocket-protocol"]?.startsWith("vite") ?? false;
49 const isSandboxRequest = hasSandboxOrigin(url.origin);
50
51 // Vite handles its own HMR upgrades; forward Sandbox preview URLs anyway.
52 if (isViteRequest && !isSandboxRequest) {
53 return;
54 }
55
56 const target = new URL(url.pathname + url.search, upstreamBase);
57 const upstream = NodeHttp.request({
58 hostname: target.hostname,
59 port: target.port,
60 path: target.pathname + target.search,
61 method: request.method,
62 // Forward the client-facing host so the worker sees the URL the client
63 // requested rather than the local workerd address.
64 headers: proxyRequestHeaders(request, url, proxySharedSecret),
65 });
66
67 const cleanup = () => {
68 upstream.destroy();
69 socket.destroy();
70 };
71
72 upstream.on("error", cleanup);
73 socket.on("close", () => upstream.destroy());
74
75 upstream.on("response", (response) => {
76 // Worker did not accept the upgrade.
77 if (!socket.destroyed) {
78 socket.destroy();
79 }
80 response.resume();
81 });
82
83 upstream.on("upgrade", (upstreamRes, upstreamSocket, upstreamHead) => {
84 upstreamSocket.on("error", () => upstreamSocket.destroy());
85
86 if (socket.destroyed) {
87 upstreamSocket.destroy();
88 return;
89 }
90
91 track(socket);
92 track(upstreamSocket);
93
94 const statusLine = `HTTP/1.1 ${upstreamRes.statusCode ?? 101} ${
95 upstreamRes.statusMessage ?? "Switching Protocols"
96 }`;
97 const headerLines: Array<string> = [statusLine];
98 for (let i = 0; i < upstreamRes.rawHeaders.length; i += 2) {
99 headerLines.push(
100 `${upstreamRes.rawHeaders[i]}: ${upstreamRes.rawHeaders[i + 1]}`,
101 );
102 }
103 socket.write(`${headerLines.join("\r\n")}\r\n\r\n`);
104
105 if (upstreamHead.length > 0) {
106 socket.write(upstreamHead);
107 }
108 if (head.length > 0) {
109 upstreamSocket.write(head);
110 }
111
112 socket.pipe(upstreamSocket).pipe(socket);
113 });
114
115 // WebSocket upgrade requests carry no body, and any early client bytes are
116 // forwarded above via `upstreamSocket.write(head)`. Ending the request
117 // directly flushes the upstream handshake deterministically, rather than
118 // relying on the incoming `request` stream to emit `end` (which can be
119 // delayed by TCP timing) and avoids a second consumer of the client socket.
120 upstream.end();
121 };
122
123 httpServer.on("upgrade", onUpgrade);
124 return () => {
125 httpServer.off("upgrade", onUpgrade);
126 for (const socket of sockets) {
127 socket.destroy();
128 }
129 sockets.clear();
130 };
131}
132
133/**
134 * Matches the origin of a Sandbox SDK preview URL.
135 * See: https://developers.cloudflare.com/sandbox/concepts/preview-urls/
136 *
137 * Pattern: https?://<port(4+ digits)>-<id(no dots)>-<token>.localhost
138 *
139 * IMPORTANT: The token segment is [a-z0-9_]+ (no hyphens) to prevent ReDoS — two adjacent
140 * [^.]+ groups separated by - cause quadratic backtracking on hyphen-heavy input. Tokens
141 * are documented as letters/digits/underscores only.
142 */
143const SANDBOX_ORIGIN_REGEXP =
144 /^https?:\/\/\d{4,}-[^.]+-[a-z0-9_]+\.localhost(:\d+)?$/i;
145
146function hasSandboxOrigin(origin: string) {
147 return SANDBOX_ORIGIN_REGEXP.test(origin);
148}