fix(miniflare): support WebSocket proxying to workerd (#10273)

Author: edmundhung

.changeset/light-garlics-count.md

@@ -0,0 +1,9 @@
+---
+"miniflare": patch
+---
+
+fix: support WebSocket proxying to workerd
+
+The dev registry proxy server now correctly handles WebSocket upgrade requests and
+tunnels bidirectional frames between the workerd processes. Previously,
+handshakes would fail due to missing upgrade logic.

packages/miniflare/src/shared/dev-registry.worker.ts

@@ -103,7 +103,9 @@ class ProxyServer {
 				serviceProxyTarget.entrypoint
 			);
 
-			this.handleProxy(req, res, address);
+			this.handleProxy(req, res, address, (callback) =>
+				this.subscribe(serviceProxyTarget.service, callback)
+			);
 			return;
 		}
 
@@ -120,7 +122,9 @@ class ProxyServer {
 				return;
 			}
 
-			this.handleProxy(req, res, address);
+			this.handleProxy(req, res, address, (callback) =>
+				this.subscribe(doProxyTarget.scriptName, callback)
+			);
 			return;
 		}
 
@@ -242,7 +246,8 @@ class ProxyServer {
 	private handleProxy(
 		req: http.IncomingMessage,
 		res: http.ServerResponse,
-		target: ProxyAddress
+		target: ProxyAddress,
+		onTargetUpdated: (callback: () => void) => void
 	) {
 		const headers = { ...req.headers };
 		let path = target.path;
@@ -271,17 +276,15 @@ class ProxyServer {
 			path,
 			headers,
 		};
+		const upstream = http.request(options);
 
-		const upstream = http.request(options, (upRes) => {
+		upstream.on("response", (upRes) => {
 			// Relay status and headers back to the original client
 			res.writeHead(upRes.statusCode ?? 500, upRes.headers);
 			// Pipe the response body
 			upRes.pipe(res);
 		});
 
-		// Pipe the client request body to the upstream
-		req.pipe(upstream);
-
 		upstream.on("error", (error) => {
 			log.error(
 				new Error(
@@ -294,6 +297,50 @@ class ProxyServer {
 			if (!res.headersSent) res.writeHead(502);
 			res.end("Bad Gateway");
 		});
+
+		if (req.headers.upgrade?.toLowerCase() === "websocket") {
+			upstream.on("upgrade", (upRes, socket, head) => {
+				// For WebSocket upgrades, we need to respond directly on the original request socket
+				if (req.socket && req.socket.writable) {
+					// Build and write complete HTTP response header
+					const statusLine = `HTTP/1.1 ${upRes.statusCode ?? 101} ${upRes.statusMessage ?? "Switching Protocols"}\r\n`;
+					let headersString = "";
+					for (let i = 0; i < upRes.rawHeaders.length; i += 2) {
+						headersString += `${upRes.rawHeaders[i]}: ${upRes.rawHeaders[i + 1]}\r\n`;
+					}
+
+					req.socket.write(statusLine + headersString + "\r\n");
+
+					// Write any buffered data
+					if (head && head.length > 0) {
+						req.socket.write(head);
+					}
+
+					// Pipe bidirectional WebSocket data
+					socket.pipe(req.socket, { end: false });
+					req.socket.pipe(socket, { end: false });
+
+					// Handle connection cleanup
+					socket.on("error", () => req.socket.destroy());
+					req.socket.on("error", () => socket.destroy());
+					socket.on("close", () => req.socket.destroy());
+					req.socket.on("close", () => socket.destroy());
+
+					// Close the socket when the target address is updated
+					onTargetUpdated(() => {
+						socket.end();
+					});
+				} else {
+					socket.end();
+				}
+			});
+
+			// End the request to trigger the upgrade
+			upstream.end();
+		} else {
+			// Pipe the client request body to the upstream for regular HTTP requests
+			req.pipe(upstream);
+		}
 	}
 
 	/**

packages/miniflare/test/dev-registry.spec.ts

@@ -126,6 +126,99 @@ test("DevRegistry: fetch to module worker", async (t) => {
 	});
 });
 
+test("DevRegistry: WebSocket upgrade to module worker", async (t) => {
+	const unsafeDevRegistryPath = await useTmp(t);
+	const local = new Miniflare({
+		name: "local-worker",
+		unsafeDevRegistryPath,
+		serviceBindings: {
+			SERVICE: {
+				name: "remote-worker",
+			},
+		},
+		compatibilityFlags: ["experimental"],
+		modules: true,
+		script: `
+			export default {
+				async fetch(request, env, ctx) {
+					const wsResponse = await env.SERVICE.fetch(request.url, {
+						headers: { Upgrade: "websocket" }
+					});
+
+					if (wsResponse.webSocket) {
+						wsResponse.webSocket.accept();
+
+						// Test bidirectional communication
+						wsResponse.webSocket.send("ping");
+
+						const messagePromise = new Promise((resolve) => {
+							wsResponse.webSocket.addEventListener("message", (event) => {
+								resolve(event.data);
+							});
+						});
+
+						const response = await messagePromise;
+
+						return new Response(\`WebSocket communication successful: \${response}\`, {
+							status: 200,
+						});
+					}
+
+					return new Response("WebSocket upgrade failed", {
+						status: 500,
+					});
+				}
+			}
+		`,
+	});
+	t.teardown(() => local.dispose());
+
+	const remote = new Miniflare({
+		name: "remote-worker",
+		unsafeDevRegistryPath,
+		compatibilityFlags: ["experimental"],
+		modules: true,
+		script: `
+			export default {
+				async fetch(request, env, ctx) {
+					// Handle WebSocket upgrade requests
+					if (request.headers.get("Upgrade") === "websocket") {
+						const [server, client] = Object.values(new WebSocketPair());
+						server.accept();
+
+						server.addEventListener("message", (event) => {
+							// Echo back with a response to test bidirectional communication
+							if (event.data === "ping") {
+								server.send("pong");
+							}
+						});
+
+						return new Response(null, { status: 101, webSocket: client });
+					}
+
+					// This test only focuses on WebSocket, no HTTP handling needed
+					return new Response("Not a WebSocket request", { status: 400 });
+				}
+			}
+		`,
+		unsafeDirectSockets: [
+			{
+				entrypoint: undefined,
+				proxy: true,
+			},
+		],
+	});
+	t.teardown(() => remote.dispose());
+
+	await remote.ready;
+	await waitUntil(t, async (t) => {
+		const res = await local.dispatchFetch("http://example.com");
+		const result = await res.text();
+		t.is(result, "WebSocket communication successful: pong");
+		t.is(res.status, 200);
+	});
+});
+
 test("DevRegistry: RPC to default entrypoint", async (t) => {
 	const unsafeDevRegistryPath = await useTmp(t);
 	const local = new Miniflare({