diff --git a/.changeset/witty-pandas-serve.md b/.changeset/witty-pandas-serve.md new file mode 100644 index 0000000000..73c3607b4f --- /dev/null +++ b/.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. diff --git a/packages/vite-plugin-cloudflare/src/__tests__/websockets.spec.ts b/packages/vite-plugin-cloudflare/src/__tests__/websockets.spec.ts index ca8aef6f6e..8eae789033 100644 --- a/packages/vite-plugin-cloudflare/src/__tests__/websockets.spec.ts +++ b/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(); + 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((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((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((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; + }); }); diff --git a/packages/vite-plugin-cloudflare/src/websockets.ts b/packages/vite-plugin-cloudflare/src/websockets.ts index b2ad52b1e1..3bac468554 100644 --- a/packages/vite-plugin-cloudflare/src/websockets.ts +++ b/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(); + } + }); } } );