Add WebSocket support to dispatchFetch/fetch and custom service bindings (#453)

* Refactor test helpers into `test-shared` directory

* Add custom `Request`/`Response` classes

* Implement `WebSocket` class and coupler

* Implement WebSocket upgrading `fetch`

* Add WebSocket coupling to loopback requests

* Allow `request.cf` to be customised in `dispatchFetch` and loopbacks

* Add WebSocket kitchen sink test

* Fix type bundling errors

* Apply PR suggestions

Author: mrbbot

package-lock.json

@@ -27,29 +27,33 @@
     "prepublishOnly": "npm run lint && npm run clean && npm run build && npm run types:bundle && npm run test",
     "test": "npm run build && ava && rimraf ./.tmp",
     "types:build": "tsc",
-    "types:bundle": "npm run types:build && node scripts/types.mjs"
+    "___$comment:types:bundle": "api-extractor doesn't know to load index.ts instead of index.d.ts when resolving imported types",
+    "types:bundle": "cp node_modules/@cloudflare/workers-types/experimental/index.ts node_modules/@cloudflare/workers-types/experimental/index.d.ts && npm run types:build && node scripts/types.mjs"
   },
   "devDependencies": {
+    "@ava/get-port": "^2.0.0",
     "@ava/typescript": "^3.0.1",
-    "@microsoft/api-extractor": "^7.19.4",
+    "@microsoft/api-extractor": "^7.33.6",
     "@types/node": "^18.11.9",
     "@types/rimraf": "^3.0.2",
     "@types/which": "^2.0.1",
     "@typescript-eslint/eslint-plugin": "^5.9.1",
     "@typescript-eslint/parser": "^5.9.1",
     "ava": "^5.0.1",
     "capnpc-ts": "^0.7.0",
+    "capnp-ts": "^0.7.0",
     "esbuild": "^0.12.20",
     "eslint": "^8.6.0",
     "eslint-config-prettier": "^8.3.0",
     "eslint-plugin-es": "^4.1.0",
     "eslint-plugin-import": "^2.24.2",
     "eslint-plugin-prettier": "^3.4.1",
     "esm": "^3.2.25",
+    "expect-type": "^0.15.0",
     "patch-package": "^6.4.7",
     "prettier": "^2.3.2",
     "rimraf": "^3.0.2",
-    "typescript": "^4.5.4",
+    "typescript": "^4.8.4",
     "which": "^2.0.2"
   },
   "engines": {

package.json

@@ -0,0 +1,75 @@
+import http from "http";
+import { Headers, RequestInfo, fetch as baseFetch } from "undici";
+import NodeWebSocket from "ws";
+import { DeferredPromise } from "../shared";
+import { Request, RequestInit } from "./request";
+import { Response } from "./response";
+import { WebSocketPair, coupleWebSocket } from "./websocket";
+
+const ignored = ["transfer-encoding", "connection", "keep-alive", "expect"];
+function headersFromIncomingRequest(req: http.IncomingMessage): Headers {
+  const entries = Object.entries(req.headers).filter(
+    (pair): pair is [string, string | string[]] => {
+      const [name, value] = pair;
+      return !ignored.includes(name) && value !== undefined;
+    }
+  );
+  return new Headers(Object.fromEntries(entries));
+}
+
+export async function fetch(
+  input: RequestInfo,
+  init?: RequestInit | Request
+): Promise<Response> {
+  const request = new Request(input, init as RequestInit);
+
+  // Handle WebSocket upgrades
+  if (
+    request.method === "GET" &&
+    request.headers.get("upgrade") === "websocket"
+  ) {
+    const url = new URL(request.url);
+    if (url.protocol !== "http:" && url.protocol !== "https:") {
+      throw new TypeError(
+        `Fetch API cannot load: ${url.toString()}.\nMake sure you're using http(s):// URLs for WebSocket requests via fetch.`
+      );
+    }
+    url.protocol = url.protocol.replace("http", "ws");
+
+    // Normalise request headers to a format ws understands, extracting the
+    // Sec-WebSocket-Protocol header as ws treats this differently
+    const headers: Record<string, string> = {};
+    let protocols: string[] | undefined;
+    for (const [key, value] of request.headers.entries()) {
+      if (key.toLowerCase() === "sec-websocket-protocol") {
+        protocols = value.split(",").map((protocol) => protocol.trim());
+      } else {
+        headers[key] = value;
+      }
+    }
+
+    // Establish web socket connection
+    const ws = new NodeWebSocket(url, protocols, {
+      followRedirects: request.redirect === "follow",
+      headers,
+    });
+
+    // Get response headers from upgrade
+    const headersPromise = new DeferredPromise<Headers>();
+    ws.once("upgrade", (req) => {
+      headersPromise.resolve(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,
+    });
+  }
+
+  const response = await baseFetch(request);
+  return new Response(response.body, response);
+}

packages/tre/src/http/fetch.ts

@@ -0,0 +1,20 @@
+export * from "./fetch";
+export * from "./request";
+export * from "./response";
+export * from "./websocket";
+
+export { File, FormData, Headers } from "undici";
+export type {
+  BodyInit,
+  HeadersInit,
+  ReferrerPolicy,
+  RequestCache,
+  RequestCredentials,
+  RequestDestination,
+  RequestDuplex,
+  RequestInfo,
+  RequestMode,
+  RequestRedirect,
+  ResponseRedirectStatus,
+  ResponseType,
+} from "undici";

packages/tre/src/http/index.ts

@@ -0,0 +1,49 @@
+import type {
+  IncomingRequestCfProperties,
+  RequestInitCfProperties,
+} from "@cloudflare/workers-types/experimental";
+import {
+  Request as BaseRequest,
+  RequestInit as BaseRequestInit,
+  RequestInfo,
+} from "undici";
+
+export type RequestInitCfType =
+  | Partial<IncomingRequestCfProperties>
+  | RequestInitCfProperties;
+
+export interface RequestInit<
+  CfType extends RequestInitCfType = RequestInitCfType
+> extends BaseRequestInit {
+  cf?: CfType;
+}
+
+const kCf = Symbol("kCf");
+export class Request<
+  CfType extends RequestInitCfType = RequestInitCfType
+> extends BaseRequest {
+  // We should be able to use a private `#cf` property here instead of a symbol
+  // here, but we need to set this on a clone, which would otherwise lead to a
+  // "Cannot write private member to an object whose class did not declare it"
+  // error.
+  [kCf]?: CfType;
+
+  constructor(input: RequestInfo, init?: RequestInit<CfType>) {
+    super(input, init);
+    this[kCf] = init?.cf;
+  }
+
+  get cf() {
+    return this[kCf];
+  }
+
+  // JSDoc comment so retained when bundling types with api-extractor
+  /** @ts-expect-error `clone` is actually defined as a method internally */
+  clone(): Request<CfType> {
+    const request = super.clone() as Request<CfType>;
+    // Update prototype so cloning a clone clones `cf`
+    Object.setPrototypeOf(request, Request.prototype);
+    request[kCf] = this[kCf];
+    return request;
+  }
+}

packages/tre/src/http/request.ts

@@ -0,0 +1,81 @@
+import {
+  Response as BaseResponse,
+  ResponseInit as BaseResponseInit,
+  BodyInit,
+  ResponseRedirectStatus,
+} from "undici";
+import { WebSocket } from "./websocket";
+
+export interface ResponseInit extends BaseResponseInit {
+  webSocket?: WebSocket | null;
+}
+
+const kWebSocket = Symbol("kWebSocket");
+export class Response extends BaseResponse {
+  // We should be able to use a private `#webSocket` property here instead of a
+  // symbol here, but `undici` calls `this.status` in its constructor, which
+  // causes a "Cannot read private member from an object whose class did not
+  // declare it" error.
+  readonly [kWebSocket]: WebSocket | null;
+
+  // Override BaseResponse's static methods for building Responses to return
+  // our type instead. Ideally, we don't want to use `Object.setPrototypeOf`.
+  // Unfortunately, `error()` and `redirect()` set the internal header guard
+  // to "immutable".
+  static error(): Response {
+    const response = BaseResponse.error() as Response;
+    Object.setPrototypeOf(response, Response.prototype);
+    return response;
+  }
+  static redirect(url: string | URL, status: ResponseRedirectStatus): Response {
+    const response = BaseResponse.redirect(url, status) as Response;
+    Object.setPrototypeOf(response, Response.prototype);
+    return response;
+  }
+  static json(data: any, init?: ResponseInit): Response {
+    // https://fetch.spec.whatwg.org/#dom-response-json
+    const body = JSON.stringify(data);
+    const response = new Response(body, init);
+    response.headers.set("Content-Type", "application/json");
+    return response;
+  }
+
+  constructor(body?: BodyInit, init?: ResponseInit) {
+    // Status 101 Switching Protocols would normally throw a RangeError, but we
+    // need to allow it for WebSockets
+    if (init?.webSocket) {
+      if (init.status !== 101) {
+        throw new RangeError(
+          "Responses with a WebSocket must have status code 101."
+        );
+      }
+      init = { ...init, status: 200 };
+    }
+
+    super(body, init);
+    this[kWebSocket] = init?.webSocket ?? null;
+  }
+
+  // JSDoc comment so retained when bundling types with api-extractor
+  /** @ts-expect-error `status` is actually defined as a getter internally */
+  get status() {
+    // When passing a WebSocket, we validate that the passed status was actually
+    // 101, but we can't store this because `undici` rightfully complains.
+    return this[kWebSocket] ? 101 : super.status;
+  }
+
+  get webSocket() {
+    return this[kWebSocket];
+  }
+
+  // JSDoc comment so retained when bundling types with api-extractor
+  /** @ts-expect-error `clone` is actually defined as a method internally */
+  clone(): Response {
+    if (this[kWebSocket]) {
+      throw new TypeError("Cannot clone a response to a WebSocket handshake.");
+    }
+    const response = super.clone() as Response;
+    Object.setPrototypeOf(response, Response.prototype);
+    return response;
+  }
+}