Include headers from WebSocket handshake in Responses, closes #151

Headers included in `new Response`s to WebSocket upgrade requests
are now returned to the client when using the HTTP server.

Headers received in the upgrade response when using `fetch` as a
WebSocket client are now included in returned `Response`.

Author: mrbbot

packages/core/src/standards/http.ts

@@ -120,6 +120,31 @@ Headers.prototype.getAll = function (key: string): string[] {
   return value ? splitCookiesString(value) : [];
 };
 
+/** @internal */
+export function _headersFromIncomingRequest(
+  req: http.IncomingMessage
+): Headers {
+  const headers = new Headers();
+  for (const [name, values] of Object.entries(req.headers)) {
+    // These headers are unsupported in undici fetch requests, they're added
+    // automatically
+    if (
+      name === "transfer-encoding" ||
+      name === "connection" ||
+      name === "keep-alive" ||
+      name === "expect"
+    ) {
+      continue;
+    }
+    if (Array.isArray(values)) {
+      for (const value of values) headers.append(name, value);
+    } else if (values !== undefined) {
+      headers.append(name, values);
+    }
+  }
+  return headers;
+}
+
 // Instead of subclassing our customised Request and Response classes from
 // BaseRequest and BaseResponse, we instead compose them and implement the same
 // interface.

packages/core/src/standards/index.ts

@@ -4,6 +4,7 @@ export * from "./domexception";
 export * from "./encoding";
 export * from "./event";
 export {
+  _headersFromIncomingRequest,
   _kInner,
   _isByteStream,
   Body,

packages/http-server/src/index.ts

@@ -1,7 +1,7 @@
 // noinspection ES6ConvertVarToLetConst
 
 import assert from "assert";
-import http, { OutgoingHttpHeaders } from "http";
+import http from "http";
 import https from "https";
 import { Transform, Writable } from "stream";
 import { ReadableStream } from "stream/web";
@@ -13,6 +13,7 @@ import {
   Request,
   Response,
   _getBodyLength,
+  _headersFromIncomingRequest,
   logResponse,
 } from "@miniflare/core";
 import { prefixError, randomHex } from "@miniflare/shared";
@@ -105,24 +106,7 @@ export async function convertNodeRequest(
   req.headers["host"] = url.host;
 
   // Build Headers object from request
-  const headers = new Headers();
-  for (const [name, values] of Object.entries(req.headers)) {
-    // These headers are unsupported in undici fetch requests, they're added
-    // automatically
-    if (
-      name === "transfer-encoding" ||
-      name === "connection" ||
-      name === "keep-alive" ||
-      name === "expect"
-    ) {
-      continue;
-    }
-    if (Array.isArray(values)) {
-      for (const value of values) headers.append(name, value);
-    } else if (values !== undefined) {
-      headers.append(name, values);
-    }
-  }
+  const headers = _headersFromIncomingRequest(req);
 
   // Create Request with additional Cloudflare specific properties:
   // https://developers.cloudflare.com/workers/runtime-apis/request#incomingrequestcfproperties
@@ -183,7 +167,7 @@ export function createRequestListener<Plugins extends HTTPPluginSignatures>(
         response = await mf.dispatchFetch(request);
         waitUntil = response.waitUntil();
         status = response.status;
-        const headers: OutgoingHttpHeaders = {};
+        const headers: http.OutgoingHttpHeaders = {};
         // eslint-disable-next-line prefer-const
         for (let [key, value] of response.headers) {
           key = key.toLowerCase();
@@ -334,6 +318,12 @@ export function createRequestListener<Plugins extends HTTPPluginSignatures>(
   };
 }
 
+const restrictedWebSocketUpgradeHeaders = [
+  "upgrade",
+  "connection",
+  "sec-websocket-accept",
+];
+
 export async function createServer<Plugins extends HTTPPluginSignatures>(
   mf: MiniflareCore<Plugins>,
   options?: http.ServerOptions & https.ServerOptions
@@ -356,6 +346,21 @@ export async function createServer<Plugins extends HTTPPluginSignatures>(
   // Setup WebSocket servers
   const webSocketServer = new WebSocketServer({ noServer: true });
   const liveReloadServer = new WebSocketServer({ noServer: true });
+
+  // Add custom headers included in response to WebSocket upgrade requests
+  const extraHeaders = new WeakMap<http.IncomingMessage, Headers>();
+  webSocketServer.on("headers", (headers, req) => {
+    const extra = extraHeaders.get(req);
+    extraHeaders.delete(req);
+    if (extra) {
+      for (const [key, value] of extra) {
+        if (!restrictedWebSocketUpgradeHeaders.includes(key)) {
+          headers.push(`${key}: ${value}`);
+        }
+      }
+    }
+  });
+
   server.on("upgrade", async (request, socket, head) => {
     // Only interested in pathname so base URL doesn't matter
     const { pathname } = new URL(request.url ?? "", "http://localhost");
@@ -382,6 +387,7 @@ export async function createServer<Plugins extends HTTPPluginSignatures>(
       }
 
       // Accept and couple the Web Socket
+      extraHeaders.set(request, response.headers);
       webSocketServer.handleUpgrade(request, socket as any, head, (ws) => {
         void coupleWebSocket(ws, webSocket);
         webSocketServer.emit("connection", ws, request);

packages/http-server/test/index.spec.ts

@@ -598,6 +598,29 @@ test("createServer: handles web socket upgrades", async (t) => {
   });
   t.is(await eventPromise, "worker:hello");
 });
+test("createServer: includes headers from web socket upgrade response", async (t) => {
+  // https://github.com/cloudflare/miniflare/issues/151
+  const mf = useMiniflareWithHandler(
+    { HTTPPlugin, WebSocketPlugin },
+    {},
+    async (globals) => {
+      const [client, worker] = Object.values(new globals.WebSocketPair());
+      worker.accept();
+      return new globals.Response(null, {
+        status: 101,
+        webSocket: client,
+        headers: { "Set-Cookie": "key=value" },
+      });
+    }
+  );
+  const port = await listen(t, await createServer(mf));
+
+  const ws = new StandardWebSocket(`ws://localhost:${port}`);
+  const [trigger, promise] = triggerPromise<http.IncomingMessage>();
+  ws.addListener("upgrade", (req) => trigger(req));
+  const req = await promise;
+  t.deepEqual(req.headers["set-cookie"], ["key=value"]);
+});
 test("createServer: expects status 101 and web socket response for upgrades", async (t) => {
   const log = new TestLog();
   log.error = (message) => log.logWithLevel(LogLevel.ERROR, message.toString());

packages/web-sockets/src/fetch.ts

@@ -4,9 +4,11 @@ import {
   RequestInfo,
   RequestInit,
   Response,
+  _headersFromIncomingRequest,
   fetch,
 } from "@miniflare/core";
 import { getRequestContext } from "@miniflare/shared";
+import { Headers } from "undici";
 import StandardWebSocket from "ws";
 import { coupleWebSocket } from "./couple";
 import { WebSocketPair } from "./websocket";
@@ -55,12 +57,22 @@ export async function upgradingFetch(
       headers,
     });
 
+    // Get response headers from upgrade
+    let headersResolve: (headers: Headers) => void;
+    const headersPromise = new Promise<Headers>((resolve) => {
+      headersResolve = resolve;
+    });
+    ws.once("upgrade", (req) => {
+      headersResolve(_headersFromIncomingRequest(req));
+    });
+
     // Couple web socket with pair and resolve
     const [worker, client] = Object.values(new WebSocketPair());
     await coupleWebSocket(ws, client);
     return new Response(null, {
       status: 101,
       webSocket: worker,
+      headers: await headersPromise,
     });
   }
 

packages/web-sockets/test/fetch.spec.ts

@@ -1,5 +1,7 @@
 import assert from "assert";
 import { Blob } from "buffer";
+import http from "http";
+import { AddressInfo } from "net";
 import { TransformStream } from "stream/web";
 import { URLSearchParams } from "url";
 import { CachePlugin } from "@miniflare/cache";
@@ -25,7 +27,10 @@ import {
 } from "@miniflare/web-sockets";
 import test from "ava";
 import { FormData } from "undici";
-import StandardWebSocket, { MessageEvent as WebSocketMessageEvent } from "ws";
+import StandardWebSocket, {
+  MessageEvent as WebSocketMessageEvent,
+  WebSocketServer,
+} from "ws";
 
 test("upgradingFetch: performs regular http request", async (t) => {
   const upstream = (await useServer(t, (req, res) => res.end("upstream"))).http;
@@ -98,6 +103,28 @@ test("upgradingFetch: performs web socket upgrade with Sec-WebSocket-Protocol he
   const event = await eventPromise;
   t.is(event.data, "protocol1,proto2,p3");
 });
+test("upgradingFetch: includes headers from web socket upgrade response", async (t) => {
+  const server = http.createServer();
+  const wss = new WebSocketServer({ server });
+  wss.on("connection", (ws) => {
+    ws.send("hello");
+    ws.close();
+  });
+  wss.on("headers", (headers) => {
+    headers.push("Set-Cookie: key=value");
+  });
+  const port = await new Promise<number>((resolve) => {
+    server.listen(0, () => {
+      t.teardown(() => server.close());
+      resolve((server.address() as AddressInfo).port);
+    });
+  });
+  const res = await upgradingFetch(`http://localhost:${port}`, {
+    headers: { upgrade: "websocket" },
+  });
+  t.not(res.webSocket, undefined);
+  t.is(res.headers.get("set-cookie"), "key=value");
+});
 test("upgradingFetch: throws on ws(s) protocols", async (t) => {
   await t.throwsAsync(
     upgradingFetch("ws://localhost/", {