undici: export ping() and publish undici:websocket:ping/pong diagnostics channels

src/js/thirdparty/undici.js

@@ -1,8 +1,10 @@
 const EventEmitter = require("node:events");
+const diagnosticsChannel = require("node:diagnostics_channel");
 const { _ReadableFromWeb: ReadableFromWeb } = require("internal/webstreams_adapters");
 
 const ObjectCreate = Object.create;
 const kEmptyObject = ObjectCreate(null);
+const NativeArrayBuffer = ArrayBuffer;
 
 var fetch = Bun.fetch;
 const bindings = $cpp("Undici.cpp", "createUndiciInternalBinding");
@@ -14,7 +16,7 @@ const File = bindings[4];
 const URL = bindings[5];
 const AbortSignal = bindings[6];
 const URLSearchParams = bindings[7];
-const WebSocket = bindings[8];
+const BunWebSocket = bindings[8];
 const CloseEvent = bindings[9];
 const ErrorEvent = bindings[10];
 const MessageEvent = bindings[11];
@@ -364,6 +366,52 @@ const util = {
   },
 };
 
+const pingChannel = diagnosticsChannel.channel("undici:websocket:ping");
+const pongChannel = diagnosticsChannel.channel("undici:websocket:pong");
+
+// undici publishes a Buffer regardless of the socket's binaryType
+function controlFramePayload(data) {
+  if (Buffer.isBuffer(data)) return data;
+  if (data instanceof NativeArrayBuffer) return Buffer.from(data);
+  if ($isTypedArrayView(data)) return Buffer.from(data.buffer, data.byteOffset, data.byteLength);
+  return data;
+}
+
+class WebSocket extends BunWebSocket {
+  constructor(...args) {
+    super(...args);
+    this.addEventListener("ping", ({ data }) => {
+      if (pingChannel.hasSubscribers) {
+        pingChannel.publish({ payload: controlFramePayload(data), websocket: this });
+      }
+    });
+    this.addEventListener("pong", ({ data }) => {
+      if (pongChannel.hasSubscribers) {
+        pongChannel.publish({ payload: controlFramePayload(data), websocket: this });
+      }
+    });
+  }
+}
+
+function ping(ws, payload) {
+  if (!(ws instanceof BunWebSocket)) {
+    throw new TypeError("Expected a WebSocket instance");
+  }
+
+  if (Buffer.isBuffer(payload)) {
+    if (payload.length > 125) {
+      throw new TypeError("A PING frame cannot have a body larger than 125 bytes.");
+    }
+  } else if (payload !== undefined) {
+    throw new TypeError("Expected buffer payload");
+  }
+
+  if (ws.readyState === BunWebSocket.OPEN) {
+    if (payload === undefined) ws.ping();
+    else ws.ping(payload);
+  }
+}
+
 class EventSource extends EventTarget {
   static CONNECTING = 0;
   static OPEN = 1;
@@ -474,6 +522,7 @@ const moduleExports = {
   mockErrors,
   MockPool,
   parseMIMEType,
+  ping,
   pipeline,
   Pool,
   ProxyAgent,

test/js/first_party/undici/undici.test.ts

@@ -1,6 +1,7 @@
 import { afterAll, beforeAll, describe, expect, it } from "bun:test";
+import diagnosticsChannel from "node:diagnostics_channel";
 import { Readable } from "node:stream";
-import { request, fetch as undiciFetch } from "undici";
+import { ping, request, fetch as undiciFetch, WebSocket as UndiciWebSocket } from "undici";
 
 import { createServer } from "../../../http-test-server";
 
@@ -267,3 +268,140 @@ describe("undici.request maxRedirections", () => {
     }
   });
 });
+
+describe.concurrent("undici WebSocket ping", () => {
+  function serveWebSocket(handlers: Partial<Bun.WebSocketHandler> = {}) {
+    return Bun.serve({
+      port: 0,
+      fetch(req, server) {
+        if (server.upgrade(req)) return;
+        return new Response(null, { status: 400 });
+      },
+      websocket: {
+        message() {},
+        ...handlers,
+      },
+    });
+  }
+
+  it("exports ping as a function", () => {
+    expect(typeof ping).toBe("function");
+  });
+
+  it("ping() sends a ping frame and undici:websocket:pong fires on the reply", async () => {
+    const serverReceived = Promise.withResolvers<Buffer>();
+    await using server = serveWebSocket({
+      ping(_ws, data) {
+        serverReceived.resolve(data);
+      },
+    });
+
+    const ws = new UndiciWebSocket(`ws://localhost:${server.port}/`);
+    const pongMessage = Promise.withResolvers<{ payload: Buffer; websocket: unknown }>();
+    const pongChannel = diagnosticsChannel.channel("undici:websocket:pong");
+    // the channel is process-global, so only accept messages for this socket
+    const onPong = (message: any) => {
+      if (message.websocket === ws) pongMessage.resolve(message);
+    };
+    pongChannel.subscribe(onPong);
+    try {
+      const opened = Promise.withResolvers<void>();
+      ws.addEventListener("open", () => opened.resolve());
+      ws.addEventListener("error", (e: any) => opened.reject(e.error ?? new Error(e.message)));
+      await opened.promise;
+
+      ping(ws, Buffer.from("hello slack"));
+
+      expect((await serverReceived.promise).toString()).toBe("hello slack");
+
+      // the server answers the ping with a pong echoing the payload
+      const message = await pongMessage.promise;
+      expect(Buffer.isBuffer(message.payload)).toBe(true);
+      expect(message.payload.toString()).toBe("hello slack");
+      expect(message.websocket).toBe(ws);
+    } finally {
+      pongChannel.unsubscribe(onPong);
+      ws.close();
+    }
+  });
+
+  it("undici:websocket:ping fires when the server sends a ping", async () => {
+    await using server = serveWebSocket({
+      open(ws) {
+        ws.ping(Buffer.from("from server"));
+      },
+    });
+
+    const ws = new UndiciWebSocket(`ws://localhost:${server.port}/`);
+    const pingMessage = Promise.withResolvers<{ payload: Buffer; websocket: unknown }>();
+    const pingChannel = diagnosticsChannel.channel("undici:websocket:ping");
+    // the channel is process-global, so only accept messages for this socket
+    const onPing = (message: any) => {
+      if (message.websocket === ws) pingMessage.resolve(message);
+    };
+    pingChannel.subscribe(onPing);
+    try {
+      ws.addEventListener("error", (e: any) => pingMessage.reject(e.error ?? new Error(e.message)));
+      ws.addEventListener("close", (e: any) =>
+        pingMessage.reject(new Error(`socket closed before ping: ${e.code} ${e.reason}`)),
+      );
+      const message = await pingMessage.promise;
+      expect(Buffer.isBuffer(message.payload)).toBe(true);
+      expect(message.payload.toString()).toBe("from server");
+      expect(message.websocket).toBe(ws);
+    } finally {
+      pingChannel.unsubscribe(onPing);
+      ws.close();
+    }
+  });
+
+  it("publishes a Buffer payload even when binaryType is 'arraybuffer'", async () => {
+    await using server = serveWebSocket({
+      open(ws) {
+        ws.ping(Buffer.from("typed"));
+      },
+    });
+
+    const ws = new UndiciWebSocket(`ws://localhost:${server.port}/`);
+    ws.binaryType = "arraybuffer";
+    const pingMessage = Promise.withResolvers<{ payload: Buffer; websocket: unknown }>();
+    const pingChannel = diagnosticsChannel.channel("undici:websocket:ping");
+    const onPing = (message: any) => {
+      if (message.websocket === ws) pingMessage.resolve(message);
+    };
+    pingChannel.subscribe(onPing);
+    try {
+      ws.addEventListener("error", (e: any) => pingMessage.reject(e.error ?? new Error(e.message)));
+      ws.addEventListener("close", (e: any) =>
+        pingMessage.reject(new Error(`socket closed before ping: ${e.code} ${e.reason}`)),
+      );
+      const message = await pingMessage.promise;
+      expect(Buffer.isBuffer(message.payload)).toBe(true);
+      expect(message.payload.toString()).toBe("typed");
+    } finally {
+      pingChannel.unsubscribe(onPing);
+      ws.close();
+    }
+  });
+
+  it("ping() validates its arguments like undici", async () => {
+    await using server = serveWebSocket();
+    const ws = new UndiciWebSocket(`ws://localhost:${server.port}/`);
+    try {
+      expect(() => ping(undefined as any)).toThrow(TypeError);
+      expect(() => ping(null as any)).toThrow(TypeError);
+      expect(() => ping({} as any)).toThrow(TypeError);
+      expect(() => ping(ws, "not a buffer" as any)).toThrow("Expected buffer payload");
+      expect(() => ping(ws, null as any)).toThrow("Expected buffer payload");
+      expect(() => ping(ws, new Uint8Array(4) as any)).toThrow("Expected buffer payload");
+      expect(() => ping(ws, Buffer.alloc(126))).toThrow("A PING frame cannot have a body larger than 125 bytes.");
+      // a ping on a socket that is not open is a no-op, not an error
+      // (matches undici 7.x: ping() only sends while OPEN and never throws on state)
+      expect(ws.readyState).toBe(UndiciWebSocket.CONNECTING);
+      expect(() => ping(ws, Buffer.alloc(125))).not.toThrow();
+      expect(() => ping(ws)).not.toThrow();
+    } finally {
+      ws.close();
+    }
+  });
+});