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()); + } + }); + } +}