From 2df0107de353055164a63cf3f4c0abd680703afc Mon Sep 17 00:00:00 2001 From: Patrick Ribbsaeter Date: Sat, 29 Aug 2026 15:07:16 +0200 Subject: [PATCH] fix(vite-plugin): deliver non-101 responses on rejected WebSocket upgrades Fix WebSocket upgrade handling in Vite plugin dev mode so non-101 HTTP responses (such as 401 Unauthorized or 403 Forbidden) returned by Workers are properly written to the raw socket with backpressure and stream cleanup before terminating, rather than abruptly destroying the socket. Fixes #15170 --- .../vite-plugin-ws-rejection-response.md | 5 + .../src/__tests__/websockets.spec.ts | 294 +++++++++++++++++- .../vite-plugin-cloudflare/src/websockets.ts | 190 ++++++++++- 3 files changed, 485 insertions(+), 4 deletions(-) create mode 100644 .changeset/vite-plugin-ws-rejection-response.md diff --git a/.changeset/vite-plugin-ws-rejection-response.md b/.changeset/vite-plugin-ws-rejection-response.md new file mode 100644 index 00000000000..1ef75e4a30e --- /dev/null +++ b/.changeset/vite-plugin-ws-rejection-response.md @@ -0,0 +1,5 @@ +--- +"@cloudflare/vite-plugin": patch +--- + +Fix WebSocket upgrade handling in dev mode so non-101 HTTP responses (e.g. 401 Unauthorized or 403 Forbidden) from Workers are delivered to the client instead of abruptly destroying the socket. diff --git a/packages/vite-plugin-cloudflare/src/__tests__/websockets.spec.ts b/packages/vite-plugin-cloudflare/src/__tests__/websockets.spec.ts index 2601abe4430..c38b436efb4 100644 --- a/packages/vite-plugin-cloudflare/src/__tests__/websockets.spec.ts +++ b/packages/vite-plugin-cloudflare/src/__tests__/websockets.spec.ts @@ -1,5 +1,6 @@ import http from "node:http"; import net from "node:net"; +import { PassThrough } from "node:stream"; import { DeferredPromise, Miniflare, Response } from "miniflare"; import { afterEach, @@ -11,7 +12,7 @@ import { test, vi, } from "vitest"; -import { handleWebSocket } from "../websockets"; +import { handleWebSocket, waitForSocketDrain } from "../websockets"; import type { AddressInfo } from "node:net"; describe("handleWebSocket", () => { @@ -67,6 +68,7 @@ describe("handleWebSocket", () => { */ async function connect() { const socket = net.connect(port, "127.0.0.1"); + socket.on("error", () => {}); openSockets.add(socket); socket.on("close", () => openSockets.delete(socket)); await new Promise((r) => socket.on("connect", r)); @@ -288,6 +290,7 @@ describe("handleWebSocket", () => { await listen(); const socket = await connect(); + socket.on("error", () => {}); const chunks: Buffer[] = []; socket.on("data", (chunk) => chunks.push(chunk)); @@ -418,4 +421,293 @@ describe("handleWebSocket", () => { expect(unhandled).not.toHaveBeenCalled(); }); + + // https://github.com/cloudflare/workers-sdk/issues/15170 + test("delivers non-101 HTTP response when Worker rejects WebSocket upgrade", async ({ + expect, + }) => { + const largePayload = "X".repeat(64 * 1024); + const expectedBody = "Unauthorized: upgrade rejected\n" + largePayload; + const mf = await listen(); + vi.spyOn(mf, "dispatchFetch").mockResolvedValue( + new Response(expectedBody, { + status: 401, + headers: { + "X-Custom-Auth": "denied", + "Content-Type": "text/plain", + }, + }) + ); + + const socket = await connect(); + + const chunks: Buffer[] = []; + const closed = new Promise((resolve) => + socket.on("close", () => resolve()) + ); + socket.on("end", () => { + socket.end(); + }); + socket.on("data", (chunk) => chunks.push(chunk)); + + socket.write( + "GET / 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; + + const raw = Buffer.concat(chunks).toString("utf8"); + expect(raw).toContain("HTTP/1.1 401"); + expect(raw).toContain("x-custom-auth: denied"); + expect(raw).toContain(expectedBody); + }); + + test("terminates response writer and cleans up resources when client disconnects during backpressure", async ({ + expect, + }) => { + const unhandled = vi.fn(); + process.on("unhandledRejection", unhandled); + onTestFinished(() => { + process.off("unhandledRejection", unhandled); + }); + const backpressureEntered = new DeferredPromise(); + httpServer.prependOnceListener("upgrade", (_request, serverSocket) => { + vi.spyOn(serverSocket, "write") + .mockImplementationOnce(() => true) + .mockImplementationOnce(() => { + backpressureEntered.resolve(); + return false; + }); + }); + const mf = await listen(); + + const chunk = new Uint8Array(1024).fill(65); + const stream = new ReadableStream({ + pull(controller) { + controller.enqueue(chunk); + }, + }); + + vi.spyOn(mf, "dispatchFetch").mockResolvedValue( + new Response(stream, { + status: 403, + headers: { + "Content-Type": "application/octet-stream", + }, + }) + ); + + const socket = await connect(); + socket.write( + "GET / 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 backpressureEntered; + expect(stream.locked).toBe(true); + socket.resetAndDestroy(); + + // Verify server-side response stream reader is released and response writer terminates + await vi.waitFor(() => { + expect(stream.locked).toBe(false); + }); + + // Verify server remains responsive and did not hang + const response = await fetch(`http://127.0.0.1:${port}`); + expect(response.ok).toBe(true); + expect(unhandled).not.toHaveBeenCalled(); + }); + + test("cancels response body when client disconnects before response body writing", async ({ + expect, + }) => { + const unhandled = vi.fn(); + process.on("unhandledRejection", unhandled); + onTestFinished(() => { + process.off("unhandledRejection", unhandled); + }); + + let cancelCalled = false; + const stream = new ReadableStream({ + pull() {}, + cancel() { + cancelCalled = true; + }, + }); + + const mf = await listen(); + const responseObj = new Response(stream, { + status: 403, + headers: { + "Content-Type": "text/plain", + }, + }); + vi.spyOn(mf, "dispatchFetch").mockResolvedValue(responseObj); + + const socket = await connect(); + socket.write( + "GET / 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 vi.waitFor(() => expect(mf.dispatchFetch).toHaveBeenCalled()); + socket.resetAndDestroy(); + + await vi.waitFor( + () => { + expect(cancelCalled).toBe(true); + expect(responseObj.body?.locked).toBe(false); + }, + { timeout: 1000 } + ); + + const response = await fetch(`http://127.0.0.1:${port}`); + expect(response.ok).toBe(true); + expect(unhandled).not.toHaveBeenCalled(); + }); + + test("cancels response body and unblocks handler when client disconnects while reader.read() is pending", async ({ + expect, + }) => { + const unhandled = vi.fn(); + process.on("unhandledRejection", unhandled); + onTestFinished(() => { + process.off("unhandledRejection", unhandled); + }); + + const pullStarted = new DeferredPromise(); + let cancelCalled = false; + let cancelResolve: (() => void) | undefined; + + const mf = await listen(); + vi.spyOn(mf, "dispatchFetch").mockImplementation(async () => { + const stream = new ReadableStream({ + pull() { + pullStarted.resolve(); + return new Promise((resolve) => { + cancelResolve = resolve; + }); + }, + cancel() { + cancelCalled = true; + cancelResolve?.(); + }, + }); + return new Response(stream, { + status: 403, + headers: { + "Content-Type": "text/plain", + }, + }); + }); + + const socket = await connect(); + socket.write( + "GET / 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" + ); + + // Wait until reader.read() is demonstrably pending + await pullStarted; + + // Client disconnects while reader.read() is pending + socket.resetAndDestroy(); + + // Verify response stream is actively cancelled without hanging + await vi.waitFor( + () => { + expect(cancelCalled).toBe(true); + }, + { timeout: 1000 } + ); + + const response = await fetch(`http://127.0.0.1:${port}`); + expect(response.ok).toBe(true); + expect(unhandled).not.toHaveBeenCalled(); + }); +}); + +describe("waitForSocketDrain", () => { + test("resolves when drain event is emitted and removes all listeners", async ({ + expect, + }) => { + const socket = new PassThrough(); + const promise = waitForSocketDrain(socket); + + expect(socket.listenerCount("drain")).toBe(1); + expect(socket.listenerCount("close")).toBe(1); + expect(socket.listenerCount("error")).toBe(1); + + socket.emit("drain"); + await promise; + + expect(socket.listenerCount("drain")).toBe(0); + expect(socket.listenerCount("close")).toBe(0); + expect(socket.listenerCount("error")).toBe(0); + }); + + test("resolves when close event is emitted and removes all listeners", async ({ + expect, + }) => { + const socket = new PassThrough(); + const promise = waitForSocketDrain(socket); + + expect(socket.listenerCount("drain")).toBe(1); + expect(socket.listenerCount("close")).toBe(1); + expect(socket.listenerCount("error")).toBe(1); + + socket.emit("close"); + await promise; + + expect(socket.listenerCount("drain")).toBe(0); + expect(socket.listenerCount("close")).toBe(0); + expect(socket.listenerCount("error")).toBe(0); + }); + + test("resolves when error event is emitted and removes all listeners", async ({ + expect, + }) => { + const socket = new PassThrough(); + const promise = waitForSocketDrain(socket); + + expect(socket.listenerCount("drain")).toBe(1); + expect(socket.listenerCount("close")).toBe(1); + expect(socket.listenerCount("error")).toBe(1); + + socket.emit("error", new Error("connection reset")); + await promise; + + expect(socket.listenerCount("drain")).toBe(0); + expect(socket.listenerCount("close")).toBe(0); + expect(socket.listenerCount("error")).toBe(0); + }); + + test("resolves immediately if socket is already destroyed", async ({ + expect, + }) => { + const socket = new PassThrough(); + socket.destroy(); + + await waitForSocketDrain(socket); + expect(socket.listenerCount("drain")).toBe(0); + expect(socket.listenerCount("close")).toBe(0); + expect(socket.listenerCount("error")).toBe(0); + }); }); diff --git a/packages/vite-plugin-cloudflare/src/websockets.ts b/packages/vite-plugin-cloudflare/src/websockets.ts index 8b9521b819c..f84a3dbc036 100644 --- a/packages/vite-plugin-cloudflare/src/websockets.ts +++ b/packages/vite-plugin-cloudflare/src/websockets.ts @@ -1,9 +1,10 @@ +import * as http from "node:http"; import { createHeaders } from "@remix-run/node-fetch-server"; import { CoreHeaders, coupleWebSocket } from "miniflare"; import { WebSocketServer } from "ws"; import { UNKNOWN_HOST } from "./shared"; import { getForwardedProto } from "./utils"; -import type { Headers, Miniflare } from "miniflare"; +import type { Headers, Miniflare, Response } from "miniflare"; import type { IncomingMessage } from "node:http"; import type { Duplex } from "node:stream"; import type * as vite from "vite"; @@ -43,7 +44,10 @@ export function handleWebSocket( "upgrade", async (request: IncomingMessage, socket: Duplex, head: Buffer) => { // Socket errors crash Node.js if unhandled - socket.on("error", () => socket.destroy()); + socket.on("error", () => {}); + socket.on("close", () => { + socket.on("error", () => {}); + }); try { const rawHost = request.headers.host ?? UNKNOWN_HOST; @@ -78,7 +82,21 @@ export function handleWebSocket( const workerWebSocket = response.webSocket; if (!workerWebSocket) { - socket.destroy(); + if (!socket.destroyed && !socket.closed) { + socket.resume(); + try { + await writeHttpResponse( + socket, + response as unknown as Response, + request.method + ); + } catch { + // Client disconnected or write failed + socket.destroy(); + } + } else if (response.body != null) { + void response.body.cancel().catch(() => {}); + } return; } @@ -170,3 +188,169 @@ const SANDBOX_ORIGIN_REGEXP = function hasSandboxOrigin(origin: string) { return SANDBOX_ORIGIN_REGEXP.test(origin); } + +/** + * Waits for a socket to drain when `write()` returns false. + * Resolves on 'drain', 'close', or 'error' to prevent hanging indefinitely + * if the remote peer disconnects while backpressure is being relieved. + */ +export function waitForSocketDrain(socket: Duplex): Promise { + if (socket.destroyed || socket.closed) { + return Promise.resolve(); + } + return new Promise((resolve) => { + const cleanup = () => { + socket.off("drain", onDrain); + socket.off("close", onClose); + socket.off("error", onError); + }; + const onDrain = () => { + cleanup(); + resolve(); + }; + const onClose = () => { + cleanup(); + resolve(); + }; + const onError = () => { + cleanup(); + resolve(); + }; + socket.once("drain", onDrain); + socket.once("close", onClose); + socket.once("error", onError); + }); +} + +/** + * Writes an HTTP response directly to a raw socket with backpressure support and orderly EOF. + * Used when a Worker returns a non-101 rejection response to a WebSocket upgrade request. + */ +async function writeHttpResponse( + socket: Duplex, + response: Response, + requestMethod?: string +) { + if (socket.destroyed || socket.closed) { + if (response.body != null) { + void response.body.cancel().catch(() => {}); + } + return; + } + + const status = response.status; + const statusText = + response.statusText || http.STATUS_CODES[status] || "Unknown"; + let headerString = `HTTP/1.1 ${status} ${statusText}\r\n`; + + let hasConnection = false; + const setCookies = + typeof response.headers.getSetCookie === "function" + ? response.headers.getSetCookie() + : []; + for (const cookie of setCookies) { + headerString += `Set-Cookie: ${cookie}\r\n`; + } + + for (const [name, value] of response.headers.entries()) { + const lower = name.toLowerCase(); + if (lower === "set-cookie") { + continue; + } + if (lower === "connection") { + hasConnection = true; + } + headerString += `${name}: ${value}\r\n`; + } + + if (!hasConnection) { + headerString += "Connection: close\r\n"; + } + headerString += "\r\n"; + + if (socket.destroyed || socket.closed) { + if (response.body != null) { + void response.body.cancel().catch(() => {}); + } + return; + } + + if (socket.write(headerString) === false) { + await waitForSocketDrain(socket); + } + + if (response.body != null) { + if (requestMethod === "HEAD" || socket.destroyed || socket.closed) { + void response.body.cancel().catch(() => {}); + } else { + const reader = response.body.getReader(); + let isSocketClosed = false; + let onSocketTerminal: (() => void) | undefined; + + const socketClosedPromise = new Promise<{ + done: true; + value: undefined; + }>((resolve) => { + onSocketTerminal = () => { + isSocketClosed = true; + void reader.cancel().catch(() => {}); + resolve({ done: true, value: undefined }); + }; + if (socket.destroyed || socket.closed) { + onSocketTerminal(); + } + }); + + if (onSocketTerminal && !socket.destroyed && !socket.closed) { + socket.once("close", onSocketTerminal); + socket.once("error", onSocketTerminal); + socket.once("end", onSocketTerminal); + } + + try { + while (!isSocketClosed && !socket.destroyed && !socket.closed) { + const { done, value } = await Promise.race([ + reader.read(), + socketClosedPromise, + ]); + if (done) { + break; + } + if (isSocketClosed || socket.destroyed || socket.closed) { + break; + } + if (value && socket.write(value) === false) { + await waitForSocketDrain(socket); + } + } + } catch { + // Reader was cancelled or read threw on disconnect/error + } finally { + if (onSocketTerminal) { + socket.off("close", onSocketTerminal); + socket.off("error", onSocketTerminal); + socket.off("end", onSocketTerminal); + } + if (isSocketClosed || socket.destroyed || socket.closed) { + try { + await reader.cancel(); + } catch {} + } + try { + reader.releaseLock(); + } catch {} + } + } + } + + if (!socket.destroyed && !socket.closed) { + socket.end(); + await new Promise((resolve) => { + if (socket.destroyed || socket.closed) { + resolve(); + } else { + socket.once("close", () => resolve()); + } + }); + } +}