feat: allow pages dev to proxy websockets (#2212)

* feat: allow pages dev to proxy websockets

* add tests for websocket requests

* add new line at EOF

Co-authored-by: Jacob M-G Evans <27247160+JacobMGEvans@users.noreply.github.com>

* improve test logic

* rename ws-app to pages-ws-app

* Update packages/wrangler/src/miniflare-cli/assets.ts

Co-authored-by: Greg Brimble <gbrimble@cloudflare.com>

* style: formatted files in pages-ws-app

* chore: update changeset to minor

* undo unintended changes to lockfile

Reverted the file back to c725ce0

* rerun npm install

Co-authored-by: Skye <46286597+Skye-31@users.noreply.github.com>
Co-authored-by: Jacob M-G Evans <27247160+JacobMGEvans@users.noreply.github.com>
Co-authored-by: Greg Brimble <gbrimble@cloudflare.com>

.changeset/stupid-ligers-walk.md

@@ -0,0 +1,5 @@
+---
+"wrangler": minor
+---
+
+feat: Allow pages dev to proxy websocket requests

fixtures/pages-ws-app/package.json

@@ -0,0 +1,36 @@
+{
+	"name": "pages-ws-app",
+	"version": "0.0.0",
+	"private": true,
+	"sideEffects": false,
+	"main": "server/index.js",
+	"scripts": {
+		"build": "esbuild --bundle --platform=node server/index.ts --outfile=dist/index.js",
+		"dev": "npx wrangler pages dev --port 8790 --proxy 8791 -- npm run server",
+		"server": "node dist/index.js",
+		"test": "npm run build && npx jest --forceExit",
+		"test:ci": "npm run build && npx jest --forceExit"
+	},
+	"jest": {
+		"restoreMocks": true,
+		"testRegex": ".*.(test|spec)\\.[jt]sx?$",
+		"testTimeout": 30000,
+		"transform": {
+			"^.+\\.c?(t|j)sx?$": [
+				"esbuild-jest",
+				{
+					"sourcemap": true
+				}
+			]
+		},
+		"transformIgnorePatterns": [
+			"node_modules/(?!find-up|locate-path|p-locate|p-limit|yocto-queue|path-exists|execa|strip-final-newline|npm-run-path|path-key|onetime|mimic-fn|human-signals|is-stream)"
+		]
+	},
+	"devDependencies": {
+		"ws": "^8.8.0"
+	},
+	"engines": {
+		"node": ">=14"
+	}
+}

fixtures/pages-ws-app/server/index.ts

@@ -0,0 +1,24 @@
+import { createServer } from "http";
+import { WebSocketServer } from "ws";
+
+const server = createServer();
+const wsServer = new WebSocketServer({ noServer: true });
+
+wsServer.on("connection", function connection(ws) {
+	ws.on("message", (msg) => console.log(msg));
+});
+
+server.on("upgrade", (request, socket, head) => {
+	wsServer.handleUpgrade(request, socket, head, (ws) => {
+		wsServer.emit("connection", ws, request);
+	});
+});
+
+server.on("request", (_request, res) => {
+	res.writeHead(200, { "Content-Type": "text/plain" });
+	res.writeHead(200, { "X-Proxied": "true" });
+	res.write("Hello, world!");
+	res.end();
+});
+
+server.listen(8791);

fixtures/pages-ws-app/tests/index.test.ts

@@ -0,0 +1,72 @@
+import { spawn } from "child_process";
+import * as path from "path";
+import { upgradingFetch } from "@miniflare/web-sockets";
+import type { ChildProcess } from "child_process";
+import type { Response } from "miniflare";
+import type { RequestInit } from "undici";
+
+const waitUntilReady = async (
+	url: string,
+	requestInit?: RequestInit
+): Promise<Response> => {
+	let response: Response | undefined = undefined;
+
+	while (response === undefined) {
+		await new Promise((resolvePromise) => setTimeout(resolvePromise, 500));
+
+		try {
+			response = await upgradingFetch(url, requestInit);
+			if (response.status === 502) response = undefined;
+		} catch {}
+	}
+
+	return response as Response;
+};
+
+const isWindows = process.platform === "win32";
+
+describe("Pages Functions", () => {
+	let wranglerProcess: ChildProcess;
+
+	beforeAll(() => {
+		wranglerProcess = spawn("npm", ["run", "dev"], {
+			shell: isWindows,
+			cwd: path.resolve(__dirname, "../"),
+			env: { BROWSER: "none", ...process.env },
+		});
+		wranglerProcess.stdout?.on("data", (chunk) => {
+			console.log(chunk.toString());
+		});
+		wranglerProcess.stderr?.on("data", (chunk) => {
+			console.log(chunk.toString());
+		});
+	});
+
+	afterAll(async () => {
+		await new Promise((resolve, reject) => {
+			wranglerProcess.once("exit", (code) => {
+				if (!code) {
+					resolve(code);
+				} else {
+					reject(code);
+				}
+			});
+			wranglerProcess.kill("SIGTERM");
+		});
+	});
+
+	it("understands normal fetches", async () => {
+		const response = await waitUntilReady("http://localhost:8790/");
+		expect(response.headers.get("x-proxied")).toBe("true");
+		const text = await response.text();
+		expect(text).toContain("Hello, world!");
+	});
+
+	it("understands websocket fetches", async () => {
+		const response = await waitUntilReady("http://localhost:8790/ws", {
+			headers: { Upgrade: "websocket" },
+		});
+		expect(response.status).toBe(101);
+		expect(response.webSocket).toBeDefined();
+	});
+});

fixtures/pages-ws-app/tsconfig.json

@@ -0,0 +1,10 @@
+{
+	"include": ["server"],
+	"compilerOptions": {
+		"target": "ES2020",
+		"module": "CommonJS",
+		"lib": ["ES2020"],
+		"types": ["@cloudflare/workers-types"],
+		"moduleResolution": "node"
+	}
+}

packages/wrangler/scripts/deps.ts

@@ -13,6 +13,7 @@ export const EXTERNAL_DEPENDENCIES = [
 	"@miniflare/core",
 	"@miniflare/durable-objects",
 	"@miniflare/tre", // TODO: remove once Miniflare 3 moved in miniflare package
+	"@miniflare/web-sockets",
 	// todo - bundle miniflare too
 	"selfsigned",
 	"source-map",

packages/wrangler/src/miniflare-cli/assets.ts

@@ -3,11 +3,11 @@ import { join } from "node:path";
 import { createMetadataObject } from "@cloudflare/pages-shared/metadata-generator/createMetadataObject";
 import { parseHeaders } from "@cloudflare/pages-shared/metadata-generator/parseHeaders";
 import { parseRedirects } from "@cloudflare/pages-shared/metadata-generator/parseRedirects";
-import { fetch as miniflareFetch } from "@miniflare/core";
 import {
 	Response as MiniflareResponse,
 	Request as MiniflareRequest,
 } from "@miniflare/core";
+import { upgradingFetch as miniflareFetch } from "@miniflare/web-sockets";
 import { watch } from "chokidar";
 import { getType } from "mime";
 import { hashFile } from "../pages/hash";
@@ -40,7 +40,12 @@ export default async function generateASSETSBinding(options: Options) {
 			try {
 				const url = new URL(miniflareRequest.url);
 				url.host = `localhost:${options.proxyPort}`;
-				return await miniflareFetch(url, miniflareRequest);
+				const proxyRequest = new MiniflareRequest(url, miniflareRequest);
+				if (proxyRequest.headers.get("Upgrade") === "websocket") {
+					proxyRequest.headers.delete("Sec-WebSocket-Accept");
+					proxyRequest.headers.delete("Sec-WebSocket-Key");
+				}
+				return await miniflareFetch(proxyRequest);
 			} catch (thrown) {
 				options.log.error(new Error(`Could not proxy request: ${thrown}`));