[Bun.serve] Support `"nodebuffer"` binaryType in `ServerWebSocket`

4 files changed, 122 insertions, 77 deletions

diff --git a/packages/bun-types/bun.d.ts b/packages/bun-types/bun.d.ts
index 00fe2963f..d06a7235e 100644
--- a/packages/bun-types/bun.d.ts
+++ b/packages/bun-types/bun.d.ts
@@ -1332,11 +1332,11 @@ declare module "bun" {
     cork: (callback: (ws: ServerWebSocket<T>) => any) => void | Promise<void>;
 
     /**
-     * Configure the {@link WebSocketHandler.message} callback to return a {@link ArrayBuffer} instead of a {@link Uint8Array}
+     * Configure the {@link WebSocketHandler.message} callback to return a {@link ArrayBuffer} or {@link Buffer} instead of a {@link Uint8Array}
      *
      * @default "uint8array"
      */
-    binaryType?: "arraybuffer" | "uint8array";
+    binaryType?: "arraybuffer" | "uint8array" | "nodebuffer";
   }
 
   type WebSocketCompressor =
@@ -1646,7 +1646,7 @@ declare module "bun" {
      * File path to a TLS key
      *
      * To enable TLS, this option is required.
-     * 
+     *
      * @deprecated since v0.6.3 - Use `key: Bun.file(path)` instead.
      */
     keyFile: string;
@@ -1654,7 +1654,7 @@ declare module "bun" {
      * File path to a TLS certificate
      *
      * To enable TLS, this option is required.
-     * 
+     *
      * @deprecated since v0.6.3 - Use `cert: Bun.file(path)` instead.
      */
     certFile: string;
@@ -1665,7 +1665,7 @@ declare module "bun" {
     passphrase?: string;
     /**
      *  File path to a .pem file for a custom root CA
-     * 
+     *
      * @deprecated since v0.6.3 - Use `ca: Bun.file(path)` instead.
      */
     caFile?: string;
@@ -1687,41 +1687,56 @@ declare module "bun" {
      */
     lowMemoryMode?: boolean;
 
-     /**
+    /**
      * Optionally override the trusted CA certificates. Default is to trust
      * the well-known CAs curated by Mozilla. Mozilla's CAs are completely
      * replaced when CAs are explicitly specified using this option.
      */
-     ca?: string | Buffer | BunFile | Array<string | Buffer | BunFile> | undefined;
-     /**
-      *  Cert chains in PEM format. One cert chain should be provided per
-      *  private key. Each cert chain should consist of the PEM formatted
-      *  certificate for a provided private key, followed by the PEM
-      *  formatted intermediate certificates (if any), in order, and not
-      *  including the root CA (the root CA must be pre-known to the peer,
-      *  see ca). When providing multiple cert chains, they do not have to
-      *  be in the same order as their private keys in key. If the
-      *  intermediate certificates are not provided, the peer will not be
-      *  able to validate the certificate, and the handshake will fail.
-      */
-     cert?: string | Buffer | BunFile | Array<string | Buffer | BunFile> | undefined;
-     /**
-      * Private keys in PEM format. PEM allows the option of private keys
-      * being encrypted. Encrypted keys will be decrypted with
-      * options.passphrase. Multiple keys using different algorithms can be
-      * provided either as an array of unencrypted key strings or buffers,
-      * or an array of objects in the form {pem: <string|buffer>[,
-      * passphrase: <string>]}. The object form can only occur in an array.
-      * object.passphrase is optional. Encrypted keys will be decrypted with
-      * object.passphrase if provided, or options.passphrase if it is not.
-      */
-     key?: string | Buffer | BunFile | Array<string | Buffer | BunFile> | undefined;
-     /**
-      * Optionally affect the OpenSSL protocol behavior, which is not
-      * usually necessary. This should be used carefully if at all! Value is
-      * a numeric bitmask of the SSL_OP_* options from OpenSSL Options
-      */
-     secureOptions?: number | undefined; // Value is a numeric bitmask of the `SSL_OP_*` options
+    ca?:
+      | string
+      | Buffer
+      | BunFile
+      | Array<string | Buffer | BunFile>
+      | undefined;
+    /**
+     *  Cert chains in PEM format. One cert chain should be provided per
+     *  private key. Each cert chain should consist of the PEM formatted
+     *  certificate for a provided private key, followed by the PEM
+     *  formatted intermediate certificates (if any), in order, and not
+     *  including the root CA (the root CA must be pre-known to the peer,
+     *  see ca). When providing multiple cert chains, they do not have to
+     *  be in the same order as their private keys in key. If the
+     *  intermediate certificates are not provided, the peer will not be
+     *  able to validate the certificate, and the handshake will fail.
+     */
+    cert?:
+      | string
+      | Buffer
+      | BunFile
+      | Array<string | Buffer | BunFile>
+      | undefined;
+    /**
+     * Private keys in PEM format. PEM allows the option of private keys
+     * being encrypted. Encrypted keys will be decrypted with
+     * options.passphrase. Multiple keys using different algorithms can be
+     * provided either as an array of unencrypted key strings or buffers,
+     * or an array of objects in the form {pem: <string|buffer>[,
+     * passphrase: <string>]}. The object form can only occur in an array.
+     * object.passphrase is optional. Encrypted keys will be decrypted with
+     * object.passphrase if provided, or options.passphrase if it is not.
+     */
+    key?:
+      | string
+      | Buffer
+      | BunFile
+      | Array<string | Buffer | BunFile>
+      | undefined;
+    /**
+     * Optionally affect the OpenSSL protocol behavior, which is not
+     * usually necessary. This should be used carefully if at all! Value is
+     * a numeric bitmask of the SSL_OP_* options from OpenSSL Options
+     */
+    secureOptions?: number | undefined; // Value is a numeric bitmask of the `SSL_OP_*` options
   }
 
   export interface TLSServeOptions extends ServeOptions, TLSOptions {
diff --git a/src/bun.js/api/server.zig b/src/bun.js/api/server.zig
index 91cfb0b54..d3a202364 100644
--- a/src/bun.js/api/server.zig
+++ b/src/bun.js/api/server.zig
@@ -3325,7 +3325,7 @@ pub const ServerWebSocket = struct {
     this_value: JSValue = .zero,
     websocket: uws.AnyWebSocket = undefined,
     closed: bool = false,
-    binary_type: JSC.JSValue.JSType = .Uint8Array,
+    binary_type: JSC.BinaryType = .Uint8Array,
     opened: bool = false,
 
     pub usingnamespace JSC.Codegen.JSServerWebSocket;
@@ -3430,6 +3430,12 @@ pub const ServerWebSocket = struct {
                         message,
                         .Uint8Array,
                     )
+                else if (this.binary_type == .Buffer)
+                    JSC.ArrayBuffer.create(
+                        globalObject,
+                        message,
+                        .Buffer,
+                    )
                 else
                     JSC.ArrayBuffer.create(
                         globalObject,
@@ -4208,18 +4214,13 @@ pub const ServerWebSocket = struct {
         log("getBinaryType()", .{});
 
         return switch (this.binary_type) {
-            .Uint8Array => ZigString.static("uint8array").toValue(globalThis),
-            else => ZigString.static("arraybuffer").toValue(globalThis),
+            .Uint8Array => ZigString.static("uint8array").toValueGC(globalThis),
+            .Buffer => ZigString.static("nodebuffer").toValueGC(globalThis),
+            .ArrayBuffer => ZigString.static("arraybuffer").toValueGC(globalThis),
+            else => @panic("Invalid binary type"),
         };
     }
 
-    pub const BinaryType = bun.ComptimeStringMap(JSC.JSValue.JSType, .{
-        &.{ "uint8array", .Uint8Array },
-        &.{ "Uint8Array", .Uint8Array },
-        &.{ "arraybuffer", .ArrayBuffer },
-        &.{ "ArrayBuffer", .ArrayBuffer },
-    });
-
     pub fn setBinaryType(
         this: *ServerWebSocket,
         globalThis: *JSC.JSGlobalObject,
@@ -4227,27 +4228,15 @@ pub const ServerWebSocket = struct {
     ) callconv(.C) bool {
         log("setBinaryType()", .{});
 
-        if (value.isEmptyOrUndefinedOrNull() or !value.isString()) {
-            globalThis.throw("binaryType must be either \"uint8array\" or \"arraybuffer\"", .{});
-            return false;
-        }
-
-        switch (BinaryType.getWithEql(
-            value.getZigString(globalThis),
-            ZigString.eqlComptime,
-        ) orelse // random value
-            .Uint8ClampedArray) {
-            .Uint8Array => {
-                this.binary_type = .Uint8Array;
-
-                return true;
-            },
-            .ArrayBuffer => {
-                this.binary_type = .ArrayBuffer;
+        switch (JSC.BinaryType.fromJSValue(globalThis, value) orelse
+            // some other value which we don't support
+            .Float64Array) {
+            .ArrayBuffer, .Buffer, .Uint8Array => |val| {
+                this.binary_type = val;
                 return true;
             },
             else => {
-                globalThis.throw("binaryType must be either \"uint8array\" or \"arraybuffer\"", .{});
+                globalThis.throw("binaryType must be either \"uint8array\" or \"arraybuffer\" or \"nodebuffer\"", .{});
                 return false;
             },
         }
diff --git a/src/bun.js/base.zig b/src/bun.js/base.zig
index 9cb2400b6..66c83dd5f 100644
--- a/src/bun.js/base.zig
+++ b/src/bun.js/base.zig
@@ -1761,10 +1761,11 @@ pub const ArrayBuffer = extern struct {
         return Stream{ .pos = 0, .buf = this.slice() };
     }
 
-    pub fn create(globalThis: *JSC.JSGlobalObject, bytes: []const u8, comptime kind: JSC.JSValue.JSType) JSValue {
+    pub fn create(globalThis: *JSC.JSGlobalObject, bytes: []const u8, comptime kind: BinaryType) JSValue {
         JSC.markBinding(@src());
         return switch (comptime kind) {
             .Uint8Array => Bun__createUint8ArrayForCopy(globalThis, bytes.ptr, bytes.len, false),
+            .Buffer => Bun__createUint8ArrayForCopy(globalThis, bytes.ptr, bytes.len, true),
             .ArrayBuffer => Bun__createArrayBufferForCopy(globalThis, bytes.ptr, bytes.len),
             else => @compileError("Not implemented yet"),
         };
@@ -4078,7 +4079,7 @@ pub const BinaryType = enum {
         return this.toJSType().toC();
     }
 
-    const Map = bun.ComptimeStringMap(
+    pub const Map = bun.ComptimeStringMap(
         BinaryType,
         .{
             .{ "ArrayBuffer", .ArrayBuffer },
diff --git a/test/js/bun/websocket/websocket-server.test.ts b/test/js/bun/websocket/websocket-server.test.ts
index a9b9b866b..014178b84 100644
--- a/test/js/bun/websocket/websocket-server.test.ts
+++ b/test/js/bun/websocket/websocket-server.test.ts
@@ -436,15 +436,54 @@ describe("websocket server", () => {
           server.stop();
         },
         message(ws, msg) {
-          if (ws.binaryType === "uint8array") {
-            expect(ws.binaryType).toBe("uint8array");
-            ws.binaryType = "arraybuffer";
-            expect(ws.binaryType).toBe("arraybuffer");
-            expect(msg instanceof Uint8Array).toBe(true);
-          } else {
-            expect(ws.binaryType).toBe("arraybuffer");
-            expect(msg instanceof ArrayBuffer).toBe(true);
-            done = true;
+          // The first message is supposed to be "uint8array"
+          // Then after uint8array, we switch it to "nodebuffer"
+          // Then after nodebuffer, we switch it to "arraybuffer"
+          // and then we're done
+          switch (ws.binaryType) {
+            case "uint8array": {
+              for (let badType of [
+                123,
+                NaN,
+                Symbol("uint8array"),
+                "uint16array",
+                "uint32array",
+                "float32array",
+                "float64array",
+                "garbage",
+              ]) {
+                expect(() => {
+                  /* @ts-ignore */
+                  ws.binaryType = badType;
+                }).toThrow();
+              }
+              expect(ws.binaryType).toBe("uint8array");
+              ws.binaryType = "nodebuffer";
+              expect(ws.binaryType).toBe("nodebuffer");
+              expect(msg instanceof Uint8Array).toBe(true);
+              expect(Buffer.isBuffer(msg)).toBe(false);
+              break;
+            }
+
+            case "nodebuffer": {
+              expect(ws.binaryType).toBe("nodebuffer");
+              ws.binaryType = "arraybuffer";
+              expect(ws.binaryType).toBe("arraybuffer");
+              expect(msg instanceof Uint8Array).toBe(true);
+              expect(Buffer.isBuffer(msg)).toBe(true);
+              break;
+            }
+
+            case "arraybuffer": {
+              expect(ws.binaryType).toBe("arraybuffer");
+              expect(msg instanceof ArrayBuffer).toBe(true);
+              done = true;
+              break;
+            }
+
+            default: {
+              throw new Error("unknown binaryType");
+            }
           }
 
           ws.send("hello world");
@@ -463,7 +502,7 @@ describe("websocket server", () => {
       },
     });
 
-    await new Promise<boolean>((resolve, reject) => {
+    const isDone = await new Promise<boolean>((resolve, reject) => {
       var counter = 0;
       const websocket = new WebSocket(`ws://${server.hostname}:${server.port}`);
       websocket.onopen = () => {
@@ -473,7 +512,7 @@ describe("websocket server", () => {
         try {
           expect(e.data).toBe("hello world");
 
-          if (counter++ > 0) {
+          if (counter++ > 2) {
             websocket.close();
             resolve(done);
           }
@@ -487,6 +526,7 @@ describe("websocket server", () => {
         reject(e);
       };
     });
+    expect(isDone).toBe(true);
   });
 
   it("does not upgrade for non-websocket connections", async () => {