diff --git a/packages/vite-plugin-cloudflare/src/__tests__/websockets.spec.ts b/packages/vite-plugin-cloudflare/src/__tests__/websockets.spec.ts index 7cf0dbee02..b480f3ffec 100644 --- a/packages/vite-plugin-cloudflare/src/__tests__/websockets.spec.ts +++ b/packages/vite-plugin-cloudflare/src/__tests__/websockets.spec.ts @@ -545,4 +545,59 @@ describe("handleWebSocket", () => { await closed; }); + + // A pipelined HTTP response can finish while `dispatchFetch` is still in + // flight. Its bytes must not be mistaken for another upgrade listener's + // handshake, or a valid Worker WebSocket is silently dropped. + test("does not treat an in-flight HTTP response as a claimed upgrade", async ({ + expect, + }) => { + startMiniflare(`export default { + async fetch() { + await new Promise((resolve) => setTimeout(resolve, 100)); + const [client, server] = Object.values(new WebSocketPair()); + server.accept(); + return new Response(null, { status: 101, webSocket: client }); + } + }`); + + httpServer.removeAllListeners("request"); + httpServer.on("request", (request, response) => { + if (request.url === "/page") { + setTimeout(() => response.end("OK"), 20); + } else { + response.end("OK"); + } + }); + + await listen(); + + 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(); + } + }); + }); + + // Pipeline a normal request and an upgrade on one connection so the + // normal response is written while the upgrade is still pending. + socket.write( + "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; + expect(Buffer.concat(chunks).toString("utf8")).toContain("HTTP/1.1 101"); + socket.destroy(); + }, 5000); }); diff --git a/packages/vite-plugin-cloudflare/src/websockets.ts b/packages/vite-plugin-cloudflare/src/websockets.ts index b68d9c2885..89dede5eb5 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"; @@ -50,14 +49,12 @@ export function handleWebSocket( // 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. + // not route). `ws` attaches a `data` listener when it takes ownership, + // while an in-flight HTTP response on the same keep-alive connection + // only writes bytes, so listener count distinguishes the two. // 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; + socket.destroyed || socket.listenerCount("data") > 0; try { const rawHost = request.headers.host ?? UNKNOWN_HOST;