fix(dev): proxy websocket upgrades for devProxy routes with ws

Author: danielroe

src/dev/app.ts

@@ -2,6 +2,7 @@ import type { Nitro } from "nitro/types";
 import type { H3Event, HTTPHandler } from "h3";
 import { createProxyServer, type ProxyServerOptions } from "httpxy";
 import type { IncomingMessage, ServerResponse } from "node:http";
+import type { Socket } from "node:net";
 import { H3, toEventHandler, serveStatic, fromNodeHandler, HTTPError } from "h3";
 import { joinURL } from "ufo";
 import mime from "mime";
@@ -21,6 +22,8 @@ export class NitroDevApp {
   nitro: Nitro;
   fetch: (req: Request) => Response | Promise<Response>;
 
+  #wsProxies: { route: string; proxy: ReturnType<typeof createHTTPProxy> }[] = [];
+
   constructor(nitro: Nitro, catchAllHandler?: HTTPHandler) {
     this.nitro = nitro;
     const app = this.#createApp(catchAllHandler);
@@ -97,6 +100,9 @@ export class NitroDevApp {
       }
       const proxy = createHTTPProxy(opts);
       app.all(route, proxy.handleEvent);
+      if (opts.ws) {
+        this.#wsProxies.push({ route, proxy });
+      }
     }
 
     // Main handler
@@ -106,6 +112,29 @@ export class NitroDevApp {
 
     return app;
   }
+
+  /**
+   * Proxy a WebSocket upgrade request if it matches a `devProxy` rule with `ws` enabled.
+   *
+   * @returns `true` if the socket was handed to a proxy, `false` if the caller should handle it.
+   */
+  proxyUpgrade(req: IncomingMessage, socket: Socket, head: any): boolean {
+    if (this.#wsProxies.length === 0) {
+      return false;
+    }
+    const path = (req.url || "/").split("?")[0]!;
+    const match = this.#wsProxies.find(
+      ({ route }) => path === route || path.startsWith(route.endsWith("/") ? route : `${route}/`)
+    );
+    if (!match) {
+      return false;
+    }
+    match.proxy.proxy.ws(req, socket, {}, head).catch((error) => {
+      this.nitro.logger.error(`Failed to proxy WebSocket upgrade for \`${path}\`:`, error);
+      socket.destroy();
+    });
+    return true;
+  }
 }
 
 // TODO: upstream to h3/node

src/dev/server.ts

@@ -123,6 +123,9 @@ export class NitroDevServer extends NitroDevApp implements RunnerRPCHooks {
   // #region Public Methods
 
   async upgrade(req: IncomingMessage, socket: Socket, head: any) {
+    if (this.proxyUpgrade(req, socket, head)) {
+      return;
+    }
     if (!this.#manager.upgrade) {
       throw new HTTPError({
         status: 501,

test/unit/dev-proxy-ws.test.ts

@@ -0,0 +1,122 @@
+import { createServer, request, type Server } from "node:http";
+import type { AddressInfo, Socket } from "node:net";
+import type { Nitro } from "nitro/types";
+import { afterAll, beforeAll, describe, expect, it } from "vitest";
+import { NitroDevApp } from "../../src/dev/app.ts";
+
+const sockets = new Set<Socket>();
+
+function listen(server: Server): Promise<number> {
+  server.on("connection", (socket) => {
+    sockets.add(socket);
+    socket.on("close", () => sockets.delete(socket));
+  });
+  return new Promise((resolve, reject) => {
+    server.on("error", reject);
+    server.listen(0, "127.0.0.1", () => resolve((server.address() as AddressInfo).port));
+  });
+}
+
+/** Minimal upgrade-aware target: replies `101` then echoes raw frames. */
+function createUpgradeTarget(upgrades: string[]): Server {
+  const server = createServer((_req, res) => res.end("http"));
+  server.on("upgrade", (req, socket) => {
+    upgrades.push(req.url!);
+    socket.write(
+      "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n"
+    );
+    socket.on("data", (chunk) => socket.write(chunk));
+  });
+  return server;
+}
+
+function upgradeRequest(port: number, path: string) {
+  return new Promise<{ status: number; echo: string }>((resolve, reject) => {
+    const req = request({
+      port,
+      host: "127.0.0.1",
+      path,
+      headers: { Connection: "Upgrade", Upgrade: "websocket" },
+    });
+    const timeout = setTimeout(() => reject(new Error("upgrade timed out")), 5000);
+    req.on("error", reject);
+    req.on("response", (res) => {
+      clearTimeout(timeout);
+      reject(new Error(`unexpected HTTP response: ${res.statusCode}`));
+    });
+    req.on("upgrade", (res, socket) => {
+      socket.once("data", (chunk) => {
+        clearTimeout(timeout);
+        socket.destroy();
+        resolve({ status: res.statusCode!, echo: chunk.toString() });
+      });
+      socket.write("ping");
+    });
+    req.end();
+  });
+}
+
+describe("dev server devProxy websocket upgrades", () => {
+  const upgrades: string[] = [];
+  const target = createUpgradeTarget(upgrades);
+  let devServer: Server;
+  let devPort: number;
+  let workerUpgrades = 0;
+
+  beforeAll(async () => {
+    const targetPort = await listen(target);
+    const app = new NitroDevApp({
+      logger: console,
+      options: {
+        baseURL: "/",
+        devHandlers: [],
+        publicAssets: [],
+        devProxy: {
+          "/proxy/ws": { target: `http://127.0.0.1:${targetPort}`, ws: true },
+          "/proxy/http": { target: `http://127.0.0.1:${targetPort}` },
+        },
+      },
+    } as unknown as Nitro);
+    devServer = createServer(async (req, res) => {
+      const response = await app.fetch(new Request(`http://127.0.0.1${req.url}`));
+      res.end(await response.text());
+    });
+    devServer.on("upgrade", (req, socket, head) => {
+      if (!app.proxyUpgrade(req, socket as Socket, head)) {
+        workerUpgrades++;
+        socket.destroy();
+      }
+    });
+    devPort = await listen(devServer);
+  });
+
+  afterAll(async () => {
+    for (const socket of sockets) {
+      socket.destroy();
+    }
+    await Promise.all([
+      new Promise((resolve) => devServer.close(resolve)),
+      new Promise((resolve) => target.close(resolve)),
+    ]);
+  });
+
+  it("proxies upgrades for a matching rule with `ws` enabled", async () => {
+    const res = await upgradeRequest(devPort, "/proxy/ws");
+    expect(res.status).toBe(101);
+    expect(res.echo).toBe("ping");
+    expect(upgrades).toContain("/proxy/ws");
+  });
+
+  it("proxies upgrades for subpaths of a matching rule", async () => {
+    const res = await upgradeRequest(devPort, "/proxy/ws/nested?foo=1");
+    expect(res.status).toBe(101);
+    expect(upgrades).toContain("/proxy/ws/nested?foo=1");
+  });
+
+  it("forwards upgrades to the worker when no rule enables `ws`", async () => {
+    await expect(upgradeRequest(devPort, "/proxy/http")).rejects.toThrow();
+    await expect(upgradeRequest(devPort, "/other")).rejects.toThrow();
+    expect(workerUpgrades).toBe(2);
+    expect(upgrades).not.toContain("/proxy/http");
+  });
+});