1import { handleWebSocket } from "../websockets.ts";
2import * as NodeHttp from "node:http";
3import type { AddressInfo } from "node:net";
4import * as NodeNet from "node:net";
5import { Readable } from "node:stream";
6import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
7import { WebSocket, WebSocketServer } from "ws";
8
9const DEBUG = !!process.env.WS_PROXY_DEBUG;
10const dbg = (...args: Array<unknown>) => {
11 if (DEBUG) {
12 console.error(`[ws-dbg ${Date.now() % 100000}]`, ...args);
13 }
14};
15
16interface UpstreamCall {
17 url: string;
18 method: string | undefined;
19 headers: NodeHttp.IncomingHttpHeaders;
20}
21
22interface Harness {
23 clientServer: NodeHttp.Server;
24 clientPort: number;
25 upstreamServer: NodeHttp.Server;
26 upstreamWss: WebSocketServer;
27 upstreamPort: number;
28 upstreamCalls: Array<UpstreamCall>;
29 upstreamConnections: Array<WebSocket>;
30 cleanup: () => void;
31}
32
33const listen = (server: NodeHttp.Server) =>
34 new Promise<number>((resolve) => {
35 server.listen(0, "127.0.0.1", () => {
36 resolve((server.address() as AddressInfo).port);
37 });
38 });
39
40const close = (server: NodeHttp.Server) =>
41 new Promise<void>((resolve, reject) => {
42 server.close((err) => (err ? reject(err) : resolve()));
43 server.closeAllConnections();
44 });
45
46const setup = async (): Promise<Harness> => {
47 const upstreamCalls: Array<UpstreamCall> = [];
48 const upstreamConnections: Array<WebSocket> = [];
49
50 const upstreamServer = NodeHttp.createServer((_req, res) => {
51 res.statusCode = 200;
52 res.end("upstream-http");
53 });
54 const upstreamWss = new WebSocketServer({ noServer: true });
55 upstreamServer.on("upgrade", (req, socket, head) => {
56 dbg("upstream received upgrade", req.method, req.url);
57 upstreamCalls.push({
58 url: req.url ?? "",
59 method: req.method,
60 headers: { ...req.headers },
61 });
62 upstreamWss.handleUpgrade(req, socket, head, (ws) => {
63 dbg("upstream handshake completed, emitting connection");
64 upstreamConnections.push(ws);
65 upstreamWss.emit("connection", ws, req);
66 });
67 });
68
69 const clientServer = NodeHttp.createServer((_req, res) => res.end("OK"));
70 const upstreamPort = await listen(upstreamServer);
71 const removeListener = handleWebSocket(
72 clientServer,
73 `http://127.0.0.1:${upstreamPort}`,
74 "websocket-secret",
75 );
76 const clientPort = await listen(clientServer);
77
78 return {
79 clientServer,
80 clientPort,
81 upstreamServer,
82 upstreamWss,
83 upstreamPort,
84 upstreamCalls,
85 upstreamConnections,
86 cleanup: () => {
87 removeListener();
88 for (const client of upstreamWss.clients) {
89 client.terminate();
90 }
91 upstreamWss.close();
92 },
93 };
94};
95
96const makeFakeRequest = (overrides: {
97 host: string;
98 url?: string;
99 protocol?: string;
100}): NodeHttp.IncomingMessage => {
101 const readable = Readable.from([]);
102 return Object.assign(readable, {
103 url: overrides.url ?? "/",
104 method: "GET",
105 socket: {},
106 headers: {
107 host: overrides.host,
108 upgrade: "websocket",
109 connection: "Upgrade",
110 "sec-websocket-key": "dGhlIHNhbXBsZSBub25jZQ==",
111 "sec-websocket-version": "13",
112 ...(overrides.protocol
113 ? { "sec-websocket-protocol": overrides.protocol }
114 : {}),
115 },
116 }) as unknown as NodeHttp.IncomingMessage;
117};
118
119const teardown = async (harness: Harness) => {
120 harness.cleanup();
121 for (const ws of harness.upstreamConnections) {
122 ws.terminate();
123 }
124 harness.clientServer.closeAllConnections();
125 harness.upstreamServer.closeAllConnections();
126 await close(harness.clientServer);
127 await close(harness.upstreamServer);
128};
129
130describe("handleWebSocket", () => {
131 let harness: Harness;
132
133 beforeEach(async () => {
134 harness = await setup();
135 });
136
137 afterEach(async () => {
138 await teardown(harness);
139 });
140
141 test("proxies messages bidirectionally", async () => {
142 harness.upstreamWss.on("connection", (ws) => {
143 ws.on("message", (data) => {
144 ws.send(`echo:${data.toString()}`);
145 });
146 ws.send("hello");
147 });
148
149 const client = new WebSocket(
150 `ws://127.0.0.1:${harness.clientPort}/path?x=1`,
151 );
152 const received: Array<string> = [];
153 client.on("upgrade", (res) =>
154 dbg("client received upgrade response", res.statusCode),
155 );
156 client.on("unexpected-response", (_req, res) =>
157 dbg("client unexpected-response", res.statusCode),
158 );
159 client.on("close", (code, reason) =>
160 dbg("client close", code, reason.toString()),
161 );
162 await new Promise<void>((resolve, reject) => {
163 client.on("open", () => {
164 dbg("client open, sending ping");
165 client.send("ping");
166 });
167 client.on("message", (data) => {
168 dbg("client message", data.toString());
169 received.push(data.toString());
170 if (received.length === 2) {
171 client.close();
172 resolve();
173 }
174 });
175 client.on("error", (err) => {
176 dbg("client error", err.message);
177 reject(err);
178 });
179 });
180
181 expect(received).toEqual(["hello", "echo:ping"]);
182 expect(harness.upstreamCalls).toHaveLength(1);
183 const call = harness.upstreamCalls[0]!;
184 expect(call.url).toBe("/path?x=1");
185 expect(call.method).toBe("GET");
186 expect(call.headers.host).toBe(`127.0.0.1:${harness.clientPort}`);
187 expect(call.headers.upgrade).toBe("websocket");
188 });
189
190 test("ignores vite HMR upgrades", async () => {
191 // Emit a synthetic upgrade event directly to avoid leaving an orphaned
192 // upgraded socket on the test HTTP server.
193 const fakeReq = makeFakeRequest({
194 host: `127.0.0.1:${harness.clientPort}`,
195 protocol: "vite-hmr",
196 });
197 const fakeSocket = new NodeNet.Socket();
198 const destroySpy = vi.spyOn(fakeSocket, "destroy");
199 const writeSpy = vi.spyOn(fakeSocket, "write");
200
201 harness.clientServer.emit("upgrade", fakeReq, fakeSocket, Buffer.alloc(0));
202
203 await new Promise((r) => setTimeout(r, 50));
204
205 expect(harness.upstreamCalls).toHaveLength(0);
206 expect(destroySpy).not.toHaveBeenCalled();
207 expect(writeSpy).not.toHaveBeenCalled();
208 });
209
210 test("forwards sandbox-preview-URL upgrades even with vite protocol", async () => {
211 harness.upstreamWss.on("connection", (ws) => {
212 ws.send("ack");
213 });
214
215 const sandboxHost = `4567-my-sandbox-sup3rs3cr3t.localhost:${harness.clientPort}`;
216 const client = new WebSocket(
217 `ws://127.0.0.1:${harness.clientPort}/`,
218 "vite-hmr",
219 {
220 headers: { host: sandboxHost },
221 },
222 );
223 client.on("open", () => dbg("sandbox client open"));
224 client.on("upgrade", (res) =>
225 dbg("sandbox client received upgrade response", res.statusCode),
226 );
227 client.on("unexpected-response", (_req, res) =>
228 dbg("sandbox client unexpected-response", res.statusCode),
229 );
230 client.on("close", (code, reason) =>
231 dbg("sandbox client close", code, reason.toString()),
232 );
233 const message = await new Promise<string>((resolve, reject) => {
234 client.on("message", (data) => {
235 dbg("sandbox client message", data.toString());
236 resolve(data.toString());
237 client.close();
238 });
239 client.on("error", (err) => {
240 dbg("sandbox client error", err.message);
241 reject(err);
242 });
243 });
244
245 expect(message).toBe("ack");
246 expect(harness.upstreamCalls).toHaveLength(1);
247 expect(harness.upstreamCalls[0]!.headers.host).toBe(sandboxHost);
248 expect(harness.upstreamCalls[0]!.headers["sec-websocket-protocol"]).toBe(
249 "vite-hmr",
250 );
251 });
252
253 test("prefers X-Forwarded-Host over Host when set by a tunnel", async () => {
254 harness.upstreamWss.on("connection", (ws) => {
255 ws.send("ack");
256 });
257
258 const client = new WebSocket(`ws://127.0.0.1:${harness.clientPort}/`, {
259 headers: { "x-forwarded-host": "example.ngrok.app" },
260 });
261 const message = await new Promise<string>((resolve, reject) => {
262 client.on("message", (data) => {
263 resolve(data.toString());
264 client.close();
265 });
266 client.on("error", reject);
267 });
268
269 expect(message).toBe("ack");
270 expect(harness.upstreamCalls).toHaveLength(1);
271 expect(harness.upstreamCalls[0]!.headers.host).toBe("example.ngrok.app");
272 });
273
274 // https://github.com/cloudflare/workers-sdk/issues/12047
275 test("survives client disconnect during upgrade", async () => {
276 // Upstream upgrade handler that never responds — the proxy is stuck waiting
277 // when the client resets the connection.
278 harness.upstreamServer.removeAllListeners("upgrade");
279 harness.upstreamServer.on("upgrade", () => {});
280
281 const fakeReq = makeFakeRequest({ host: "localhost" });
282 const fakeSocket = new NodeNet.Socket();
283 harness.clientServer.emit("upgrade", fakeReq, fakeSocket, Buffer.alloc(0));
284
285 fakeSocket.destroy(new Error("ECONNRESET"));
286
287 // The proxy and Vite server should still be responsive.
288 const response = await fetch(`http://127.0.0.1:${harness.clientPort}/`);
289 expect(response.ok).toBe(true);
290 });
291
292 test("destroys client socket when worker does not upgrade", async () => {
293 // Upstream responds with a normal HTTP 200 instead of upgrading.
294 harness.upstreamServer.removeAllListeners("upgrade");
295
296 const fakeReq = makeFakeRequest({ host: "localhost" });
297 const fakeSocket = new NodeNet.Socket();
298 harness.clientServer.emit("upgrade", fakeReq, fakeSocket, Buffer.alloc(0));
299
300 await new Promise<void>((resolve) => fakeSocket.once("close", resolve));
301 expect(fakeSocket.destroyed).toBe(true);
302 });
303
304 test("returns a cleanup function that removes the upgrade listener", async () => {
305 const server = NodeHttp.createServer();
306 await listen(server);
307 const remove = handleWebSocket(
308 server,
309 `http://127.0.0.1:${harness.upstreamPort}`,
310 "websocket-secret",
311 );
312 expect(server.listenerCount("upgrade")).toBe(1);
313 remove();
314 expect(server.listenerCount("upgrade")).toBe(0);
315 await close(server);
316 });
317});