From 5e114d867fe5600f191b00553009a11ffb6ac02a Mon Sep 17 00:00:00 2001
From: Ajay Ghanwat <ghanwat.ajay@workroot.in>
Date: Tue, 18 Aug 2026 19:16:54 +0530
Subject: [PATCH] feat(gateway): forward application identity on WebSocket
 upgrades

Adds gatewayWebSocketBackendHeaders so proxied upgrades carry application
identity while Bun keeps ownership of WebSocket framing.

Pre-existing working-tree change, committed as-is rather than authored
here.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 packages/dev-server/src/gateway.ts       | 38 +++++++++++++++++++++++-
 packages/dev-server/test/gateway.test.ts | 30 +++++++++++++++++++
 2 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/packages/dev-server/src/gateway.ts b/packages/dev-server/src/gateway.ts
index 98996169..969ff1d8 100644
--- a/packages/dev-server/src/gateway.ts
+++ b/packages/dev-server/src/gateway.ts
@@ -113,6 +113,7 @@ interface Target extends GatewayApp {
 interface WsBridge {
   origin: string;
   path: string;
+  headers: Record<string, string>;
   backend?: WebSocket;
   queue: Array<string | ArrayBuffer>;
   maxMessageBytes: number;
@@ -433,6 +434,30 @@ export function gatewayProxyHeaders(
   return headers;
 }
 
+/** Forward application identity to the child while leaving WebSocket framing to Bun. */
+export function gatewayWebSocketBackendHeaders(
+  req: Request,
+  url: URL,
+  ip: string,
+  forwardedHeaders: boolean,
+  backendOrigin: string,
+): Record<string, string> {
+  const headers = stripUntrustedInternalHeaders(
+    gatewayProxyHeaders(req, url, ip, forwardedHeaders),
+  );
+  headers.delete("host");
+  headers.delete("connection");
+  headers.delete("upgrade");
+  headers.delete("accept-encoding");
+  for (const name of [...headers.keys()]) {
+    if (name.startsWith("sec-websocket-")) headers.delete(name);
+  }
+  // The public origin was validated at the gateway edge. The child receives a
+  // new, trusted same-origin connection from its private gateway listener.
+  headers.set("origin", backendOrigin);
+  return Object.fromEntries(headers);
+}
+
 /** Remove headers that only a direct workspace-to-app request may supply. */
 export function stripUntrustedInternalHeaders(headers: Headers): Headers {
   const sanitized = new Headers(headers);
@@ -692,6 +717,13 @@ export async function startGateway(opts: GatewayOptions): Promise<RunningGateway
             data: {
               origin: target.origin,
               path: url.pathname + url.search,
+              headers: gatewayWebSocketBackendHeaders(
+                req,
+                url,
+                ip,
+                forwardedHeaders,
+                target.origin,
+              ),
               queue: [],
               maxMessageBytes: websocketSecurity.maxMessageBytes ?? 64 * 1024,
               maxQueuedMessages: websocketSecurity.maxQueuedMessages ?? 100,
@@ -753,7 +785,11 @@ export async function startGateway(opts: GatewayOptions): Promise<RunningGateway
       websocket: {
         open(ws) {
           const backendUrl = ws.data.origin.replace(/^http/, "ws") + ws.data.path;
-          const backend = new WebSocket(backendUrl);
+          const BackendWebSocket = WebSocket as unknown as new (
+            url: string,
+            options: Bun.WebSocketOptions,
+          ) => WebSocket;
+          const backend = new BackendWebSocket(backendUrl, { headers: ws.data.headers });
           ws.data.backend = backend;
           backend.addEventListener("open", () => {
             for (const m of ws.data.queue) backend.send(m);
diff --git a/packages/dev-server/test/gateway.test.ts b/packages/dev-server/test/gateway.test.ts
index 1bde0bbc..e366f059 100644
--- a/packages/dev-server/test/gateway.test.ts
+++ b/packages/dev-server/test/gateway.test.ts
@@ -5,6 +5,7 @@ import {
   forwardAuthFailure,
   forwardAuthHeaders,
   gatewayProxyHeaders,
+  gatewayWebSocketBackendHeaders,
   stripUntrustedInternalHeaders,
   gatewayRestartDelay,
   internalError,
@@ -37,6 +38,35 @@ test("gateway disables compression for its internal proxy hop", () => {
   expect(headers.get("x-forwarded-for")).toBe("127.0.0.1");
 });
 
+test("gateway WebSocket bridge forwards validated application identity", () => {
+  const request = new Request("http://web.localhost:3000/__wrnexus/hmr", {
+    headers: {
+      host: "web.localhost:3000",
+      origin: "http://web.localhost:3000",
+      cookie: "session=abc",
+      connection: "Upgrade",
+      upgrade: "websocket",
+      "sec-websocket-key": "test-key",
+    },
+  });
+  const headers = gatewayWebSocketBackendHeaders(
+    request,
+    new URL(request.url),
+    "127.0.0.1",
+    true,
+    "http://127.0.0.1:3001",
+  );
+
+  expect(headers.origin).toBe("http://127.0.0.1:3001");
+  expect(headers.cookie).toBe("session=abc");
+  expect(headers["x-forwarded-host"]).toBe("web.localhost:3000");
+  expect(headers["x-forwarded-proto"]).toBe("http");
+  expect(headers.host).toBeUndefined();
+  expect(headers.connection).toBeUndefined();
+  expect(headers.upgrade).toBeUndefined();
+  expect(headers["sec-websocket-key"]).toBeUndefined();
+});
+
 test("gateway proxy headers do not preserve the RPC internal marker", () => {
   const request = new Request("http://localhost:3000/path", {
     headers: { "x-wrnexus-internal": "1" },