Upgrade ws type definitions, test for invalid client close code

Author: mrbbot

package-lock.json

@@ -18,7 +18,6 @@ import {
 import { randomHex } from "@miniflare/shared";
 import { coupleWebSocket } from "@miniflare/web-sockets";
 import { BodyInit, Headers } from "undici";
-// @ts-expect-error ws's type definitions haven't been updated yet
 import StandardWebSocket, { WebSocketServer } from "ws";
 import { getAccessibleHosts } from "./helpers";
 import { HTTPPlugin, RequestMeta } from "./plugin";
@@ -364,7 +363,7 @@ export async function createServer<Plugins extends HTTPPluginSignatures>(
     const { pathname } = new URL(request.url ?? "", "http://localhost");
     const server =
       pathname === "/cdn-cgi/mf/reload" ? liveReloadServer : webSocketServer;
-    server.handleUpgrade(request, socket, head, (ws: StandardWebSocket) =>
+    server.handleUpgrade(request, socket as any, head, (ws) =>
       server.emit("connection", ws, request)
     );
   });

packages/http-server/src/index.ts

@@ -32,7 +32,7 @@ import {
 } from "@miniflare/shared-test";
 import { MessageEvent, WebSocketPlugin } from "@miniflare/web-sockets";
 import test, { ExecutionContext, Macro } from "ava";
-import StandardWebSocket from "ws";
+import StandardWebSocket, { Data, Event as WebSocketEvent } from "ws";
 
 function listen(
   t: ExecutionContext,
@@ -502,7 +502,7 @@ test("createServer: handles web socket upgrades", async (t) => {
   const port = await listen(t, await createServer(mf));
 
   const ws = new StandardWebSocket(`ws://localhost:${port}`);
-  const [eventTrigger, eventPromise] = triggerPromise<string>();
+  const [eventTrigger, eventPromise] = triggerPromise<Data>();
   ws.addEventListener("message", (e) => {
     eventTrigger(e.data);
   });
@@ -546,7 +546,7 @@ test("createServer: notifies connected live reload clients on reload", async (t)
 
   const port = await listen(t, await createServer(mf));
   const ws = new StandardWebSocket(`ws://localhost:${port}/cdn-cgi/mf/reload`);
-  const [openTrigger, openPromise] = triggerPromise<void>();
+  const [openTrigger, openPromise] = triggerPromise<WebSocketEvent>();
   ws.addEventListener("open", openTrigger);
   let receivedClose = false;
   const [closeTrigger, closePromise] = triggerPromise<void>();

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

@@ -2,7 +2,6 @@ import http from "http";
 import { AddressInfo } from "net";
 import { URL } from "url";
 import { ExecutionContext } from "ava";
-// @ts-expect-error ws's type definitions haven't been updated yet
 import StandardWebSocket, { WebSocketServer } from "ws";
 
 export async function useServer(

packages/shared-test/src/http.ts

@@ -43,6 +43,6 @@
   },
   "devDependencies": {
     "@miniflare/shared-test": "2.0.0-next.3",
-    "@types/ws": "^7.4.7"
+    "@types/ws": "^8.2.0"
   }
 }

packages/web-sockets/package.json

@@ -1,4 +1,6 @@
 import assert from "assert";
+import http from "http";
+import { AddressInfo } from "net";
 import { viewToBuffer } from "@miniflare/shared";
 import {
   noop,
@@ -10,7 +12,10 @@ import {
 import { WebSocketPair, coupleWebSocket } from "@miniflare/web-sockets";
 import { CloseEvent, MessageEvent } from "@miniflare/web-sockets";
 import test from "ava";
-import StandardWebSocket from "ws";
+import StandardWebSocket, {
+  Event as WebSocketEvent,
+  WebSocketServer,
+} from "ws";
 
 test("coupleWebSocket: throws if already coupled", async (t) => {
   const server = await useServer(t, noop, (ws) => ws.send("test"));
@@ -97,7 +102,30 @@ test("coupleWebSocket: closes worker socket on client close", async (t) => {
   t.is(event.code, 1000);
   t.is(event.reason, "Test Closure");
 });
-// TODO: add test for invalid WebSocket close code
+test("coupleWebSocket: closes worker socket with invalid client close code", async (t) => {
+  const server = http.createServer();
+  const wss = new WebSocketServer({ server });
+  wss.on("connection", (ws) => {
+    // Close WebSocket without code, defaults to 1005 (No Status Received)
+    // which would be an invalid code if passed normally
+    ws.close();
+  });
+  const port = await new Promise<number>((resolve) => {
+    server.listen(0, () => {
+      resolve((server.address() as AddressInfo).port);
+    });
+  });
+  const ws = new StandardWebSocket(`ws://localhost:${port}`);
+  const [client, worker] = Object.values(new WebSocketPair());
+
+  const [eventTrigger, eventPromise] = triggerPromise<CloseEvent>();
+  worker.addEventListener("close", eventTrigger);
+  worker.accept();
+  await coupleWebSocket(ws, client);
+
+  const event = await eventPromise;
+  t.is(event.code, 1005);
+});
 
 test("coupleWebSocket: forwards messages from worker to client before coupling", async (t) => {
   const [eventTrigger, eventPromise] = triggerPromise<{ data: any }>();
@@ -177,7 +205,7 @@ test("coupleWebSocket: accepts worker socket immediately if already open", async
   // Send before coupling, simulates sending message in worker code before returning response
   worker.send("test");
   // Make sure socket is open before terminating
-  const [openTrigger, openPromise] = triggerPromise<void>();
+  const [openTrigger, openPromise] = triggerPromise<WebSocketEvent>();
   ws.addEventListener("open", openTrigger);
   await openPromise;
   await coupleWebSocket(ws, client);
@@ -192,7 +220,7 @@ test("coupleWebSocket: throws if web socket already closed", async (t) => {
 
   worker.accept();
   // Make sure socket is open before closing
-  const [openTrigger, openPromise] = triggerPromise<void>();
+  const [openTrigger, openPromise] = triggerPromise<WebSocketEvent>();
   ws.addEventListener("open", openTrigger);
   await openPromise;
   // Make sure socket is closed before terminating