[vite-plugin] Preserve WebSockets owned by other v1 listeners

.changeset/witty-pandas-serve.md

@@ -0,0 +1,7 @@
+---
+"@cloudflare/vite-plugin": patch
+---
+
+fix(vite-plugin): don't destroy WebSockets owned by other upgrade listeners
+
+Node dispatches `upgrade` to every registered listener. The dev-server handler destroyed the raw socket when the Worker returned no route, even if another listener (such as Vite DevTools at `/__devtools/__ws`) had already completed the handshake. It now leaves the socket alone once another listener has written a response.

packages/vite-plugin-cloudflare/src/__tests__/websockets.spec.ts

@@ -11,8 +11,10 @@ import {
 	test,
 	vi,
 } from "vitest";
+import { WebSocketServer } from "ws";
 import { handleWebSocket } from "../websockets";
 import type { AddressInfo } from "node:net";
+import type { Duplex } from "node:stream";
 
 describe("handleWebSocket", () => {
 	let httpServer: http.Server;
@@ -452,4 +454,96 @@ describe("handleWebSocket", () => {
 
 		expect(unhandled).not.toHaveBeenCalled();
 	});
+
+	// https://github.com/cloudflare/workers-sdk/issues/15654
+	test("does not destroy a WebSocket another upgrade listener has already claimed", async ({
+		expect,
+	}) => {
+		const mf = await listen();
+		const deferred = new DeferredPromise<Response>();
+		vi.spyOn(mf, "dispatchFetch").mockReturnValue(deferred);
+
+		// The Worker has no route for the DevTools socket, so the plugin's
+		// dispatch resolves without a WebSocket after the owner has upgraded.
+		const owner = new WebSocketServer({ noServer: true });
+		let serverSocket: Duplex | undefined;
+		onTestFinished(() => {
+			for (const client of owner.clients) {
+				client.terminate();
+			}
+			owner.close();
+		});
+		httpServer.on("upgrade", (request, socket, head) => {
+			if (request.url === "/__devtools/__ws") {
+				serverSocket = socket;
+				owner.handleUpgrade(request, socket, head, (ws) => {
+					owner.emit("connection", ws, request);
+				});
+			}
+		});
+
+		const socket = await connect();
+		const chunks: Buffer[] = [];
+		const upgraded = new Promise<void>((resolve) => {
+			socket.on("data", (chunk) => {
+				chunks.push(chunk);
+				if (Buffer.concat(chunks).toString("utf8").includes("HTTP/1.1 101")) {
+					resolve();
+				}
+			});
+		});
+
+		socket.write(
+			"GET /__devtools/__ws HTTP/1.1\r\n" +
+				`Host: 127.0.0.1:${port}\r\n` +
+				"Upgrade: websocket\r\n" +
+				"Connection: Upgrade\r\n" +
+				"Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n" +
+				"Sec-WebSocket-Version: 13\r\n\r\n"
+		);
+		await upgraded;
+
+		if (!serverSocket) {
+			throw new Error("expected the DevTools listener to claim the socket");
+		}
+		const destroy = vi.fn();
+		const originalDestroy = serverSocket.destroy.bind(serverSocket);
+		vi.spyOn(serverSocket, "destroy").mockImplementation((...args) => {
+			destroy();
+			return originalDestroy(...args);
+		});
+
+		deferred.resolve(new Response("not found", { status: 404 }));
+		// Let the plugin handler resume and make its teardown decision.
+		await new Promise<void>((resolve) => setImmediate(resolve));
+
+		// Assert on the server socket directly; the client may not observe the
+		// close until a later poll phase.
+		expect(destroy).not.toHaveBeenCalled();
+		expect(socket.destroyed).toBe(false);
+		expect(Buffer.concat(chunks).toString("utf8")).toContain("HTTP/1.1 101");
+		socket.destroy();
+	});
+
+	test("destroys an unclaimed upgrade when the Worker does not route it", async () => {
+		const mf = await listen();
+		vi.spyOn(mf, "dispatchFetch").mockResolvedValue(
+			new Response("not found", { status: 404 })
+		);
+
+		const socket = await connect();
+		const closed = new Promise<void>((resolve) =>
+			socket.on("close", () => resolve())
+		);
+		socket.write(
+			"GET /unrouted HTTP/1.1\r\n" +
+				`Host: 127.0.0.1:${port}\r\n` +
+				"Upgrade: websocket\r\n" +
+				"Connection: Upgrade\r\n" +
+				"Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n" +
+				"Sec-WebSocket-Version: 13\r\n\r\n"
+		);
+
+		await closed;
+	});
 });

packages/vite-plugin-cloudflare/src/websockets.ts

@@ -4,6 +4,7 @@ import { UNKNOWN_HOST } from "./shared";
 import { createHeadersForIncomingMessage, getForwardedProto } from "./utils";
 import type { Headers, Miniflare } from "miniflare";
 import type { IncomingMessage } from "node:http";
+import type { Socket } from "node:net";
 import type { Duplex } from "node:stream";
 import type * as vite from "vite";
 
@@ -38,12 +39,26 @@ export function handleWebSocket(
 		}
 	);
 
-	httpServer.on(
+	// Prepend so the write-count baseline below is captured before another
+	// `upgrade` listener can answer the same request. Node dispatches `upgrade`
+	// to every listener, not just the first.
+	httpServer.prependListener(
 		"upgrade",
 		async (request: IncomingMessage, socket: Duplex, head: Buffer) => {
 			// Socket errors crash Node.js if unhandled
 			socket.on("error", () => socket.destroy());
 
+			// Another listener may already have completed the handshake (for
+			// example Vite DevTools at `/__devtools/__ws`, which the Worker does
+			// not route). `bytesWritten` is cumulative across keep-alive
+			// connections, so compare against a baseline rather than zero.
+			// See https://github.com/cloudflare/workers-sdk/issues/15654
+			const bytesWrittenAtStart =
+				(socket as unknown as Socket).bytesWritten ?? 0;
+			const isClaimed = () =>
+				socket.destroyed ||
+				((socket as unknown as Socket).bytesWritten ?? 0) > bytesWrittenAtStart;
+
 			try {
 				const rawHost = request.headers.host ?? UNKNOWN_HOST;
 				// Honor `X-Forwarded-Proto` so that the upgrade URL reflects the
@@ -77,7 +92,17 @@ export function handleWebSocket(
 				const workerWebSocket = response.webSocket;
 
 				if (!workerWebSocket) {
-					socket.destroy();
+					// This listener does not own the socket. Destroy it only if no
+					// other listener has written a response in the meantime.
+					if (!isClaimed()) {
+						socket.destroy();
+					}
+					return;
+				}
+
+				// A listener claimed the socket while `dispatchFetch` was in
+				// flight, so a second upgrade would corrupt the connection.
+				if (isClaimed()) {
 					return;
 				}
 
@@ -104,9 +129,14 @@ export function handleWebSocket(
 				// This listener is `async`, so an uncaught rejection here escapes as
 				// an unhandled rejection — which terminates the Node.js process on
 				// modern versions and leaks the client socket. Tear the socket down
-				// instead, mirroring the `!workerWebSocket` path above.
+				// instead. Defer the teardown one tick so another listener can claim
+				// the socket first, mirroring the `!workerWebSocket` path above.
 				workerResponseHeaders.delete(request);
-				socket.destroy();
+				setImmediate(() => {
+					if (!isClaimed()) {
+						socket.destroy();
+					}
+				});
 			}
 		}
 	);