From 599b671c0ac3a5f7cf2c5962bec83c4022aafa1d Mon Sep 17 00:00:00 2001 From: Pete Bacon Darwin Date: Tue, 27 Oct 2026 14:44:50 +0000 Subject: [PATCH] [vite-plugin] Route DevTools WebSocket upgrades explicitly (#16519) --- .changeset/witty-pandas-serve.md | 4 +- .../src/__tests__/websockets.spec.ts | 149 +++++++++++------- .../vite-plugin-cloudflare/src/websockets.ts | 42 ++--- 3 files changed, 107 insertions(+), 88 deletions(-) diff --git a/.changeset/witty-pandas-serve.md b/.changeset/witty-pandas-serve.md index 73c3607b4f..332f4574f5 100644 --- a/.changeset/witty-pandas-serve.md +++ b/.changeset/witty-pandas-serve.md @@ -2,6 +2,6 @@ "@cloudflare/vite-plugin": patch --- -fix(vite-plugin): don't destroy WebSockets owned by other upgrade listeners +Fix Vite DevTools WebSocket connections closing unexpectedly -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. +Leave upgrades at `/__devtools/__ws` to Vite DevTools instead of forwarding them to the Worker. Sandbox preview requests continue to reach the Worker. diff --git a/packages/vite-plugin-cloudflare/src/__tests__/websockets.spec.ts b/packages/vite-plugin-cloudflare/src/__tests__/websockets.spec.ts index 8eae789033..096ba57e58 100644 --- a/packages/vite-plugin-cloudflare/src/__tests__/websockets.spec.ts +++ b/packages/vite-plugin-cloudflare/src/__tests__/websockets.spec.ts @@ -1,3 +1,4 @@ +import { once } from "node:events"; import http from "node:http"; import net from "node:net"; import { DeferredPromise, Miniflare, Response } from "miniflare"; @@ -11,10 +12,9 @@ import { test, vi, } from "vitest"; -import { WebSocketServer } from "ws"; +import { WebSocket, 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; @@ -186,7 +186,8 @@ describe("handleWebSocket", () => { async function dispatchUpgrade( expect: ExpectStatic, mf: Miniflare, - headers: Record + headers: Record, + pathname = "/" ) { const deferred = new DeferredPromise(); const mockedDispatchFetch = vi @@ -199,7 +200,7 @@ describe("handleWebSocket", () => { .map(([k, v]) => `${k}: ${v}`) .join("\r\n"); socket.write( - "GET / HTTP/1.1\r\n" + + `GET ${pathname} HTTP/1.1\r\n` + `${headerLines}\r\n` + "Upgrade: websocket\r\n" + "Connection: Upgrade\r\n" + @@ -456,72 +457,114 @@ describe("handleWebSocket", () => { }); // https://github.com/cloudflare/workers-sdk/issues/15654 - test("does not destroy a WebSocket another upgrade listener has already claimed", async ({ - expect, - }) => { + test.for(["/__devtools/__ws", "/__devtools/__ws?token=testing"])( + "leaves %s to Vite DevTools without dispatching to the Worker", + async (pathname, { expect }) => { + const mf = await listen(); + const dispatchFetch = vi.spyOn(mf, "dispatchFetch"); + + const owner = new WebSocketServer({ noServer: true }); + onTestFinished(() => { + for (const client of owner.clients) { + client.terminate(); + } + owner.close(); + }); + httpServer.on("upgrade", (request, socket, head) => { + if (request.url === pathname) { + owner.handleUpgrade(request, socket, head, (ws) => { + ws.on("message", (message) => ws.send(message)); + }); + } + }); + + const socket = await connect(); + const client = new WebSocket(`ws://127.0.0.1:${port}${pathname}`, { + createConnection: () => socket, + }); + onTestFinished(() => client.terminate()); + await once(client, "open"); + const received = once(client, "message"); + client.send("DevTools connected"); + const [message] = await received; + expect(message.toString()).toBe("DevTools connected"); + expect(dispatchFetch).not.toHaveBeenCalled(); + } + ); + + test.for(["/__devtools/__ws", "/__devtools/__ws?token=testing"])( + "forwards Sandbox requests for %s to the Worker", + async (pathname, { expect }) => { + const mf = await listen(); + const host = `4567-my-sandbox-sup3rs3cr3t.localhost:${port}`; + const url = await dispatchUpgrade(expect, mf, { Host: host }, pathname); + expect(`${url}`).toBe(`http://${host}${pathname}`); + } + ); + + test.for([ + "/__devtools/__ws-extra", + "/__devtools/__ws/child", + "/__devtools/other", + ])("forwards %s to the Worker", async (pathname, { expect }) => { const mf = await listen(); - const deferred = new DeferredPromise(); - vi.spyOn(mf, "dispatchFetch").mockReturnValue(deferred); + const host = `127.0.0.1:${port}`; + const url = await dispatchUpgrade(expect, mf, { Host: host }, pathname); + expect(`${url}`).toBe(`http://${host}${pathname}`); + }); - // 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(); + test("upgrades a Worker WebSocket after an earlier pipelined HTTP response", async ({ + expect, + }) => { + const pageResponse = new DeferredPromise(); + httpServer.removeAllListeners("request"); + httpServer.on("request", (_request, response) => { + pageResponse.resolve(response); }); - 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 mf = await listen(); + const dispatchFetch = mf.dispatchFetch.bind(mf); + const dispatchStarted = new DeferredPromise(); + const resumeDispatch = new DeferredPromise(); + onTestFinished(() => resumeDispatch.resolve()); + vi.spyOn(mf, "dispatchFetch").mockImplementation(async (...args) => { + dispatchStarted.resolve(); + await resumeDispatch; + return dispatchFetch(...args); }); 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(); - } - }); + const upgraded = new DeferredPromise(); + socket.on("data", (chunk) => { + chunks.push(chunk); + if (Buffer.concat(chunks).toString("utf8").includes("HTTP/1.1 101")) { + upgraded.resolve(); + } }); - socket.write( - "GET /__devtools/__ws HTTP/1.1\r\n" + + "GET /page HTTP/1.1\r\n" + + `Host: 127.0.0.1:${port}\r\n\r\n` + + "GET /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)); + // Finish the earlier HTTP response after the upgrade handler starts, + // but before the Worker returns its WebSocket. + await dispatchStarted; + const response = await pageResponse; + const finished = once(response, "finish"); + response.end("OK"); + await finished; + resumeDispatch.resolve(); + await upgraded; - // 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"); + const raw = Buffer.concat(chunks).toString("utf8"); + expect(raw).toMatch(/^HTTP\/1\.1 200/); + expect(raw).toContain("OKHTTP/1.1 101 Switching Protocols"); socket.destroy(); }); diff --git a/packages/vite-plugin-cloudflare/src/websockets.ts b/packages/vite-plugin-cloudflare/src/websockets.ts index 3bac468554..9af04258eb 100644 --- a/packages/vite-plugin-cloudflare/src/websockets.ts +++ b/packages/vite-plugin-cloudflare/src/websockets.ts @@ -4,7 +4,6 @@ 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"; @@ -39,26 +38,12 @@ export function handleWebSocket( } ); - // 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( + httpServer.on( "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 @@ -72,10 +57,12 @@ export function handleWebSocket( const isViteRequest = request.headers["sec-websocket-protocol"]?.startsWith("vite"); + const isDevToolsRequest = url.pathname === "/__devtools/__ws"; const isSandboxRequest = hasSandboxOrigin(url.origin); - // Ignore Vite HMR WebSockets but forward on all sandbox requests. - if (isViteRequest && !isSandboxRequest) { + // Leave Vite HMR and DevTools upgrades to their own handlers, but + // forward all sandbox requests to the Worker. + if ((isViteRequest || isDevToolsRequest) && !isSandboxRequest) { return; } @@ -92,17 +79,11 @@ export function handleWebSocket( const workerWebSocket = response.webSocket; if (!workerWebSocket) { - // 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(); - } + socket.destroy(); return; } - // A listener claimed the socket while `dispatchFetch` was in - // flight, so a second upgrade would corrupt the connection. - if (isClaimed()) { + if (socket.destroyed) { return; } @@ -129,14 +110,9 @@ 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. Defer the teardown one tick so another listener can claim - // the socket first, mirroring the `!workerWebSocket` path above. + // instead. workerResponseHeaders.delete(request); - setImmediate(() => { - if (!isClaimed()) { - socket.destroy(); - } - }); + socket.destroy(); } } );