[WebSocket] Implement `"nodebuffer"` binaryType

5 files changed, 125 insertions, 5 deletions

diff --git a/.vscode/c_cpp_properties.json b/.vscode/c_cpp_properties.json
index 7aec502d3..7600b5d20 100644
--- a/.vscode/c_cpp_properties.json
+++ b/.vscode/c_cpp_properties.json
@@ -33,6 +33,7 @@
           "${workspaceFolder}/src/bun.js/bindings/*",
           "${workspaceFolder}/src/bun.js/bindings/sqlite/",
           "${workspaceFolder}/src/bun.js/bindings/webcrypto/",
+          "${workspaceFolder}/src/bun.js/bindings/webcore/",
           "${workspaceFolder}/src/bun.js/builtins/*",
           "${workspaceFolder}/src/bun.js/builtins/cpp/*",
           "${workspaceFolder}/src/bun.js/modules/*",
diff --git a/packages/bun-types/globals.d.ts b/packages/bun-types/globals.d.ts
index 5a00d5ebb..b8388b49c 100644
--- a/packages/bun-types/globals.d.ts
+++ b/packages/bun-types/globals.d.ts
@@ -1,4 +1,7 @@
-type BinaryType = "arraybuffer" | "blob";
+/**
+ * "blob" is not supported yet
+ */
+type BinaryType = "arraybuffer" | "nodebuffer" | "blob";
 type Transferable = ArrayBuffer;
 type MessageEventSource = undefined;
 type Encoding = "utf-8" | "windows-1252" | "utf-16";
@@ -1803,7 +1806,7 @@ declare var CustomEvent: {
 interface WebSocketEventMap {
   close: CloseEvent;
   error: Event;
-  message: MessageEvent;
+  message: MessageEvent<Buffer | ArrayBuffer | string>;
   open: Event;
 }
 
@@ -1812,7 +1815,9 @@ interface WebSocket extends EventTarget {
   /**
    * Returns a string that indicates how binary data from the WebSocket object is exposed to scripts:
    *
-   * Can be set, to change how binary data is returned. The default is "blob".
+   * Can be set, to change how binary data is returned. The default is `"arraybuffer"`.
+   *
+   * Unlike in browsers, you can also set `binaryType` to `"nodebuffer"` to receive a {@link Buffer} object.
    */
   binaryType: BinaryType;
   /**
@@ -1825,7 +1830,9 @@ interface WebSocket extends EventTarget {
   readonly extensions: string;
   onclose: ((this: WebSocket, ev: CloseEvent) => any) | null;
   onerror: ((this: WebSocket, ev: Event) => any) | null;
-  onmessage: ((this: WebSocket, ev: MessageEvent) => any) | null;
+  onmessage:
+    | ((this: WebSocket, ev: WebSocketEventMap["message"]) => any)
+    | null;
   onopen: ((this: WebSocket, ev: Event) => any) | null;
   /** Returns the subprotocol selected by the server, if any. It can be used in conjunction with the array form of the constructor's second argument to perform subprotocol negotiation. */
   readonly protocol: string;
diff --git a/src/bun.js/bindings/webcore/WebSocket.cpp b/src/bun.js/bindings/webcore/WebSocket.cpp
index 404336050..4af556a7b 100644
--- a/src/bun.js/bindings/webcore/WebSocket.cpp
+++ b/src/bun.js/bindings/webcore/WebSocket.cpp
@@ -69,6 +69,8 @@
 #include <wtf/text/CString.h>
 #include <wtf/text/StringBuilder.h>
 
+#include "JSBuffer.h"
+
 // #if USE(WEB_THREAD)
 // #include "WebCoreThreadRun.h"
 // #endif
@@ -680,6 +682,8 @@ String WebSocket::binaryType() const
     //     return "blob"_s;
     case BinaryType::ArrayBuffer:
         return "arraybuffer"_s;
+    case BinaryType::NodeBuffer:
+        return "nodebuffer"_s;
     }
     ASSERT_NOT_REACHED();
     return String();
@@ -694,6 +698,9 @@ ExceptionOr<void> WebSocket::setBinaryType(const String& binaryType)
     if (binaryType == "arraybuffer"_s) {
         m_binaryType = BinaryType::ArrayBuffer;
         return {};
+    } else if (binaryType == "nodebuffer"_s) {
+        m_binaryType = BinaryType::NodeBuffer;
+        return {};
     }
     // scriptExecutionContext()->addConsoleMessage(MessageSource::JS, MessageLevel::Error, "'" + binaryType + "' is not a valid value for binaryType; binaryType remains unchanged.");
     return Exception { SyntaxError, makeString("'"_s, binaryType, "' is not a valid value for binaryType; binaryType remains unchanged."_s) };
@@ -860,6 +867,49 @@ void WebSocket::didReceiveBinaryData(Vector<uint8_t>&& binaryData)
 
         break;
     }
+    case BinaryType::NodeBuffer: {
+
+        if (this->hasEventListeners("message"_s)) {
+            // the main reason for dispatching on a separate tick is to handle when you haven't yet attached an event listener
+            this->incPendingActivityCount();
+            JSUint8Array* buffer = jsCast<JSUint8Array*>(JSValue::decode(JSBuffer__bufferFromLength(scriptExecutionContext()->jsGlobalObject(), binaryData.size())));
+            if (binaryData.size() > 0)
+                memcpy(buffer->vector(), binaryData.data(), binaryData.size());
+            JSC::EnsureStillAliveScope ensureStillAlive(buffer);
+            MessageEvent::Init init;
+            init.data = buffer;
+            init.origin = this->m_url.string();
+
+            dispatchEvent(MessageEvent::create(eventNames().messageEvent, WTFMove(init), EventIsTrusted::Yes));
+            this->decPendingActivityCount();
+            return;
+        }
+
+        if (auto* context = scriptExecutionContext()) {
+            auto arrayBuffer = JSC::ArrayBuffer::tryCreate(binaryData.data(), binaryData.size());
+
+            this->incPendingActivityCount();
+
+            context->postTask([this, buffer = WTFMove(arrayBuffer), protectedThis = Ref { *this }](ScriptExecutionContext& context) {
+                ASSERT(scriptExecutionContext());
+                size_t length = buffer->byteLength();
+                JSUint8Array* uint8array = JSUint8Array::create(
+                    scriptExecutionContext()->jsGlobalObject(),
+                    reinterpret_cast<Zig::GlobalObject*>(scriptExecutionContext()->jsGlobalObject())->JSBufferSubclassStructure(),
+                    WTFMove(buffer.copyRef()),
+                    0,
+                    length);
+                JSC::EnsureStillAliveScope ensureStillAlive(uint8array);
+                MessageEvent::Init init;
+                init.data = uint8array;
+                init.origin = protectedThis->m_url.string();
+                protectedThis->dispatchEvent(MessageEvent::create(eventNames().messageEvent, WTFMove(init), EventIsTrusted::Yes));
+                protectedThis->decPendingActivityCount();
+            });
+        }
+
+        break;
+    }
     }
     // });
 }
diff --git a/src/bun.js/bindings/webcore/WebSocket.h b/src/bun.js/bindings/webcore/WebSocket.h
index 2c1bc3524..42261cfc4 100644
--- a/src/bun.js/bindings/webcore/WebSocket.h
+++ b/src/bun.js/bindings/webcore/WebSocket.h
@@ -164,7 +164,9 @@ private:
     void failAsynchronously();
 
     enum class BinaryType { Blob,
-        ArrayBuffer };
+        ArrayBuffer,
+        // non-standard:
+        NodeBuffer };
 
     State m_state { CONNECTING };
     URL m_url;
diff --git a/test/js/web/websocket/websocket.test.js b/test/js/web/websocket/websocket.test.js
index 99d60f292..867b86123 100644
--- a/test/js/web/websocket/websocket.test.js
+++ b/test/js/web/websocket/websocket.test.js
@@ -71,6 +71,66 @@ describe("WebSocket", () => {
     });
     const ws = new WebSocket(`http://${server.hostname}:${server.port}`, {});
   });
+  describe("nodebuffer", () => {
+    it("should support 'nodebuffer' binaryType", done => {
+      const server = Bun.serve({
+        port: 0,
+        fetch(req, server) {
+          if (server.upgrade(req)) {
+            return;
+          }
+
+          return new Response();
+        },
+        websocket: {
+          open(ws) {
+            ws.sendBinary(new Uint8Array([1, 2, 3]));
+          },
+        },
+      });
+      const ws = new WebSocket(`http://${server.hostname}:${server.port}`, {});
+      ws.binaryType = "nodebuffer";
+      expect(ws.binaryType).toBe("nodebuffer");
+      Bun.gc(true);
+      ws.onmessage = ({ data }) => {
+        expect(Buffer.isBuffer(data)).toBe(true);
+        expect(data).toEqual(new Uint8Array([1, 2, 3]));
+        server.stop(true);
+        Bun.gc(true);
+        done();
+      };
+    });
+
+    it("should support 'nodebuffer' binaryType when the handler is not immediately provided", done => {
+      var client;
+      const server = Bun.serve({
+        port: 0,
+        fetch(req, server) {
+          if (server.upgrade(req)) {
+            return;
+          }
+
+          return new Response();
+        },
+        websocket: {
+          open(ws) {
+            ws.sendBinary(new Uint8Array([1, 2, 3]));
+            setTimeout(() => {
+              client.onmessage = ({ data }) => {
+                expect(Buffer.isBuffer(data)).toBe(true);
+                expect(data).toEqual(new Uint8Array([1, 2, 3]));
+                server.stop(true);
+                done();
+              };
+            }, 0);
+          },
+        },
+      });
+      client = new WebSocket(`http://${server.hostname}:${server.port}`, {});
+      client.binaryType = "nodebuffer";
+      expect(client.binaryType).toBe("nodebuffer");
+    });
+  });
 
   it("should send and receive messages", async () => {
     const ws = new WebSocket(TEST_WEBSOCKET_HOST);