Allow invalid code on WebSocket close from clients, closes #86

cloudflare · Jan 7, 2022 · 7196e8d · 7196e8d
1 parent a8205b9
commit 7196e8d
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 9 deletions.
diff --git a/packages/web-sockets/src/couple.ts b/packages/web-sockets/src/couple.ts
@@ -4,6 +4,7 @@ import {
   ErrorEvent,
   WebSocket,
   kAccepted,
+  kClose,
   kClosed,
   kCoupled,
 } from "./websocket";
@@ -32,7 +33,7 @@ export async function coupleWebSocket(
     }
   });
   ws.on("close", (code: number, reason: Buffer) => {
-    if (!pair[kClosed]) pair.close(code, reason.toString());
+    if (!pair[kClosed]) pair[kClose](code, reason.toString());
   });
   ws.on("error", (error) => {
     pair.dispatchEvent(new ErrorEvent(error));

diff --git a/packages/web-sockets/src/websocket.ts b/packages/web-sockets/src/websocket.ts
@@ -34,6 +34,8 @@ const kPair = Symbol("kPair");
 export const kAccepted = Symbol("kAccepted");
 export const kCoupled = Symbol("kCoupled");
 export const kClosed = Symbol("kClosed");
+// Internal close method exposed to bypass close code checking
+export const kClose = Symbol("kClose");
 
 export type WebSocketEventMap = {
   message: MessageEvent;
@@ -91,13 +93,6 @@ export class WebSocket extends InputGatedEventTarget<WebSocketEventMap> {
   }
 
   close(code?: number, reason?: string): void {
-    const pair = this[kPair];
-    if (!this[kAccepted]) {
-      throw new TypeError(
-        "You must call accept() on this WebSocket before sending messages."
-      );
-    }
-    if (this[kClosed]) throw new TypeError("WebSocket already closed");
     if (code) {
       // https://developer.mozilla.org/en-US/docs/Web/API/CloseEvent/code
       const validCode =
@@ -114,9 +109,20 @@ export class WebSocket extends InputGatedEventTarget<WebSocketEventMap> {
         "If you specify a WebSocket close reason, you must also specify a code."
       );
     }
+    this[kClose](code, reason);
+  }
 
+  [kClose](code?: number, reason?: string): void {
+    // Split from close() so we don't check the close code when forwarding close
+    // events from the client
+    if (!this[kAccepted]) {
+      throw new TypeError(
+        "You must call accept() on this WebSocket before sending messages."
+      );
+    }
+    if (this[kClosed]) throw new TypeError("WebSocket already closed");
     this[kClosed] = true;
-    pair[kClosed] = true;
+    this[kPair][kClosed] = true;
     void this.#dispatchCloseEvent(code, reason);
   }
 

diff --git a/packages/web-sockets/test/couple.spec.ts b/packages/web-sockets/test/couple.spec.ts
@@ -97,6 +97,7 @@ 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: forwards messages from worker to client before coupling", async (t) => {
   const [eventTrigger, eventPromise] = triggerPromise<{ data: any }>();