1 files changed, 197 insertions, 57 deletions
diff --git a/src/bun.js/ws.exports.js b/src/bun.js/ws.exports.js
index c141c39f6..e1f042220 100644
--- a/src/bun.js/ws.exports.js
+++ b/src/bun.js/ws.exports.js
@@ -9,85 +9,225 @@ const kBunInternals = Symbol.for("::bunternal::");
 const readyStates = ["CONNECTING", "OPEN", "CLOSING", "CLOSED"];
 const encoder = new TextEncoder();
 
-class BunWebSocket extends globalThis.WebSocket {
-  constructor(url, ...args) {
-    super(url, ...args);
-    this.#wrappedHandlers = new WeakMap();
+const emittedWarnings = new Set();
+function emitWarning(type, message) {
+  if (emittedWarnings.has(type)) return;
+  emittedWarnings.add(type);
+  // process.emitWarning(message); // our printing is bad
+  console.warn("[bun] Warning:", message);
+}
+
+/*
+ * deviations: we do not implement these events
+ * - "unexpected-response"
+ * - "upgrade"
+ * - "ping"
+ * - "pong"
+ * - "redirect"
+ */
+class BunWebSocket extends EventEmitter {
+  static CONNECTING = 0;
+  static OPEN = 1;
+  static CLOSING = 2;
+  static CLOSED = 3;
+
+  #ws;
+  #paused = false;
+  #fragments = false;
+  #binaryType = "nodebuffer";
+  readyState = BunWebSocket.CONNECTING;
+
+  constructor(url, protocols, options) {
+    // deviation: we don't support anything in `options`
+    super();
+    let ws = (this.#ws = new WebSocket(url, protocols));
+    ws.binaryType = "nodebuffer"; // bun's WebSocket supports "nodebuffer"
+    ws.addEventListener("open", () => {
+      this.readyState = BunWebSocket.OPEN;
+      this.emit("open");
+    });
+    ws.addEventListener("error", err => {
+      this.readyState = BunWebSocket.CLOSED;
+      this.emit("error", err);
+    });
+    ws.addEventListener("close", ev => {
+      this.readyState = BunWebSocket.CLOSED;
+      this.emit("close", ev.code, ev.reason);
+    });
+    ws.addEventListener("message", ev => {
+      const isBinary = typeof ev.data !== "string";
+      if (isBinary) {
+        this.emit("message", this.#fragments ? [ev.data] : ev.data, isBinary);
+      } else {
+        var encoded = encoder.encode(ev.data);
+        if (this.#binaryType !== "arraybuffer") {
+          encoded = Buffer.from(encoded.buffer, encoded.byteOffset, encoded.byteLength);
+        }
+        this.emit("message", this.#fragments ? [encoded] : encoded, isBinary);
+      }
+    });
+  }
+
+  on(event, listener) {
+    if (
+      event === "unexpected-response" ||
+      event === "upgrade" ||
+      event === "ping" ||
+      event === "pong" ||
+      event === "redirect"
+    ) {
+      emitWarning(event, "ws.WebSocket '" + event + "' event is not implemented in bun");
+    }
+    return super.on(event, listener);
+  }
+
+  send(data, opts, cb) {
+    this.#ws.send(data, opts?.compress);
+    // deviation: this should be called once the data is written, not immediately
+    typeof cb === "function" && cb();
+  }
+
+  close(code, reason) {
+    this.#ws.close(code, reason);
   }
-  #binaryType;
-  #wrappedHandlers = new WeakMap();
 
   get binaryType() {
     return this.#binaryType;
   }
-  set binaryType(type) {
-    if (type !== "nodebuffer" && type !== "blob" && type !== "arraybuffer") {
-      throw new TypeError("binaryType must be either 'blob', 'arraybuffer' or 'nodebuffer'");
-    }
-    if (type !== "blob") {
-      super.binaryType = type;
+
+  set binaryType(value) {
+    if (value) this.#ws.binaryType = value;
+  }
+
+  set binaryType(value) {
+    if (value === "nodebuffer" || value === "arraybuffer") {
+      this.#ws.binaryType = this.#binaryType = value;
+      this.#fragments = false;
+    } else if (value === "fragments") {
+      this.#ws.binaryType = "nodebuffer";
+      this.#binaryType = "fragments";
+      this.#fragments = true;
     }
-    this.#binaryType = type;
   }
 
-  send(data, opts, cb) {
-    super.send(data, opts?.compress);
-    typeof cb === "function" && cb();
+  get protocol() {
+    return this.#ws.protocol;
   }
 
-  on(event, callback) {
-    if (event === "message") {
-      var handler = ({ data }) => {
-        try {
-          if (this.#binaryType == "blob") {
-            data = new Blob([data]);
-          }
-          callback(data);
-        } catch (e) {
-          globalThis.reportError(e);
-        }
-      };
+  get extensions() {
+    return this.#ws.extensions;
+  }
 
-      this.#wrappedHandlers.set(callback, handler);
-      this.addEventListener(event, handler);
-    } else {
-      this.addEventListener(event, callback);
+  // deviation: this does not support `message` with `binaryType = "fragments"`
+  addEventListener(type, listener, options) {
+    this.#ws.addEventListener(type, listener, options);
+  }
+
+  removeEventListener(type, listener) {
+    this.#ws.removeEventListener(type, listener);
+  }
+
+  get onopen() {
+    return this.#ws.onopen;
+  }
+
+  set onopen(value) {
+    this.#ws.onopen = value;
+  }
+
+  get onerror() {
+    return this.#ws.onerror;
+  }
+
+  set onerror(value) {
+    this.#ws.onerror = value;
+  }
+
+  get onclose() {
+    return this.#ws.onclose;
+  }
+
+  set onclose(value) {
+    this.#ws.onclose = value;
+  }
+
+  get onmessage() {
+    return this.#ws.onmessage;
+  }
+
+  // deviation: this does not support `binaryType = "fragments"`
+  set onmessage(value) {
+    this.#ws.onmessage = value;
+  }
+
+  get bufferedAmount() {
+    return this.#ws.bufferedAmount;
+  }
+
+  get isPaused() {
+    return this.#paused;
+  }
+
+  ping(data, mask, cb) {
+    if (this.readyState === BunWebSocket.CONNECTING) {
+      throw new Error("WebSocket is not open: readyState 0 (CONNECTING)");
+    }
+
+    if (typeof data === "function") {
+      cb = data;
+      data = mask = undefined;
+    } else if (typeof mask === "function") {
+      cb = mask;
+      mask = undefined;
     }
+
+    if (typeof data === "number") data = data.toString();
+
+    // deviation: we don't support ping
+    emitWarning("ping()", "ws.WebSocket.ping() is not implemented in bun");
+    typeof cb === "function" && cb();
   }
 
-  once(event, callback) {
-    if (event === "message") {
-      var handler = ({ data }) => {
-        try {
-          callback(data);
-        } catch (e) {
-          globalThis.reportError(e);
-        }
-      };
+  pong(data, mask, cb) {
+    if (this.readyState === BunWebSocket.CONNECTING) {
+      throw new Error("WebSocket is not open: readyState 0 (CONNECTING)");
+    }
 
-      this.#wrappedHandlers.set(callback, handler);
-      this.addEventListener(event, handler, { once: true });
-    } else {
-      this.addEventListener(event, callback, { once: true });
+    if (typeof data === "function") {
+      cb = data;
+      data = mask = undefined;
+    } else if (typeof mask === "function") {
+      cb = mask;
+      mask = undefined;
     }
+
+    if (typeof data === "number") data = data.toString();
+
+    // deviation: we don't support pong
+    emitWarning("pong()", "ws.WebSocket.pong() is not implemented in bun");
+    typeof cb === "function" && cb();
   }
 
-  emit(event, data) {
-    if (event === "message") {
-      this.dispatchEvent(new MessageEvent("message", { data }));
-    } else {
-      this.dispatchEvent(new CustomEvent(event, { detail: data }));
+  pause() {
+    if (this.readyState === WebSocket.CONNECTING || this.readyState === WebSocket.CLOSED) {
+      return;
     }
+
+    this.#paused = true;
+
+    // deviation: we dont support pause()
+    emitWarning("pause()", "ws.WebSocket.pause() is not implemented in bun");
   }
 
-  off(event, callback) {
-    var wrapped = this.#wrappedHandlers.get(callback);
-    if (wrapped) {
-      this.removeEventListener(event, wrapped);
-      this.#wrappedHandlers.delete(callback);
-    } else {
-      this.removeEventListener(event, callback);
+  resume() {
+    if (this.readyState === WebSocket.CONNECTING || this.readyState === WebSocket.CLOSED) {
+      return;
     }
+
+    this.#paused = false;
+
+    // deviation: we dont support resume()
+    emitWarning("resume()", "ws.WebSocket.resume() is not implemented in bun");
   }
 }