Implement client WebSocket.bufferedAmount

src/http_jsc/websocket_client.rs

@@ -2092,6 +2092,24 @@ impl<const SSL: bool> WebSocket<SSL> {
         // This is under-estimated a little, as we don't include usockets context.
         cost
     }
+
+    /// Bytes queued by `send()` that have not yet been written to the socket.
+    /// Backs the client `WebSocket.bufferedAmount` getter. Includes the framing
+    /// bytes of buffered frames (the send buffer holds fully framed messages),
+    /// plus any encrypted bytes the proxy tunnel still holds.
+    //
+    // `extern "C"` entrypoint; `this` is non-null by C++ contract (see SAFETY comment below).
+    #[allow(clippy::not_unsafe_ptr_arg_deref)]
+    pub extern "C" fn get_buffered_amount(this: *const Self) -> usize {
+        // SAFETY: called from C++ with a valid pointer
+        let this = unsafe { &*this };
+        let mut buffered = this.send_buffer.readable_length();
+        if let Some(tunnel) = &this.proxy_tunnel {
+            // SAFETY: `tunnel` holds a live ref (RefPtr has no `Deref`).
+            buffered += unsafe { tunnel.as_ref() }.buffered_amount();
+        }
+        buffered
+    }
 }
 
 // ──────────────────────────────────────────────────────────────────────────
@@ -2110,6 +2128,7 @@ macro_rules! export_websocket_client {
         cancel = $cancel:ident,
         close = $close:ident,
         finalize = $finalize:ident,
+        get_buffered_amount = $get_buffered_amount:ident,
         init = $init:ident,
         init_with_tunnel = $init_with_tunnel:ident,
         memory_cost = $memory_cost:ident,
@@ -2130,6 +2149,10 @@ macro_rules! export_websocket_client {
             WebSocket::<$ssl>::finalize(this)
         }
         #[unsafe(no_mangle)]
+        pub extern "C" fn $get_buffered_amount(this: *const WebSocket<$ssl>) -> usize {
+            WebSocket::<$ssl>::get_buffered_amount(this)
+        }
+        #[unsafe(no_mangle)]
         pub extern "C" fn $init(
             outgoing: *mut CppWebSocket,
             input_socket: *mut c_void,
@@ -2200,6 +2223,7 @@ export_websocket_client!(
     cancel = Bun__WebSocketClient__cancel,
     close = Bun__WebSocketClient__close,
     finalize = Bun__WebSocketClient__finalize,
+    get_buffered_amount = Bun__WebSocketClient__getBufferedAmount,
     init = Bun__WebSocketClient__init,
     init_with_tunnel = Bun__WebSocketClient__initWithTunnel,
     memory_cost = Bun__WebSocketClient__memoryCost,
@@ -2212,6 +2236,7 @@ export_websocket_client!(
     cancel = Bun__WebSocketClientTLS__cancel,
     close = Bun__WebSocketClientTLS__close,
     finalize = Bun__WebSocketClientTLS__finalize,
+    get_buffered_amount = Bun__WebSocketClientTLS__getBufferedAmount,
     init = Bun__WebSocketClientTLS__init,
     init_with_tunnel = Bun__WebSocketClientTLS__initWithTunnel,
     memory_cost = Bun__WebSocketClientTLS__memoryCost,

src/http_jsc/websocket_client/WebSocketProxyTunnel.rs

@@ -600,6 +600,11 @@ impl WebSocketProxyTunnel {
     pub(crate) fn has_backpressure(&self) -> bool {
         self.write_buffer.is_not_empty()
     }
+
+    /// Encrypted bytes still buffered in the tunnel awaiting a writable socket.
+    pub(crate) fn buffered_amount(&self) -> usize {
+        self.write_buffer.size()
+    }
 }
 
 impl Drop for WebSocketProxyTunnel {

src/jsc/bindings/webcore/WebSocket.cpp

@@ -174,6 +174,13 @@
     return a + b;
 }
 
+static unsigned clampToUnsigned(size_t value)
+{
+    return value > std::numeric_limits<unsigned>::max()
+        ? std::numeric_limits<unsigned>::max()
+        : static_cast<unsigned>(value);
+}
+
 ASCIILiteral WebSocket::subprotocolSeparator()
 {
     return ", "_s;
@@ -857,8 +864,6 @@
     switch (m_connectedWebSocketKind) {
     case ConnectedWebSocketKind::Client: {
         Bun__WebSocketClient__writeBinaryData(this->m_connectedWebSocket.client, reinterpret_cast<const unsigned char*>(baseAddress), length, static_cast<uint8_t>(op));
-        // this->m_connectedWebSocket.client->send({ baseAddress, length }, opCode);
-        // this->m_bufferedAmount = this->m_connectedWebSocket.client->getBufferedAmount();
         break;
     }
     case ConnectedWebSocketKind::ClientSSL: {
@@ -887,8 +892,6 @@
     case ConnectedWebSocketKind::Client: {
         auto zigStr = Zig::toZigString(message);
         Bun__WebSocketClient__writeString(this->m_connectedWebSocket.client, &zigStr, static_cast<uint8_t>(op));
-        // this->m_connectedWebSocket.client->send({ baseAddress, length }, opCode);
-        // this->m_bufferedAmount = this->m_connectedWebSocket.client->getBufferedAmount();
         break;
     }
     case ConnectedWebSocketKind::ClientSSL: {
@@ -991,17 +994,19 @@
     m_state = CLOSING;
     switch (m_connectedWebSocketKind) {
     case ConnectedWebSocketKind::Client: {
+        // Snapshot the backlog before the connection (and its send buffer) is
+        // torn down: per spec bufferedAmount must not reset to 0 on close.
+        m_bufferedAmount = clampToUnsigned(Bun__WebSocketClient__getBufferedAmount(this->m_connectedWebSocket.client));
         ZigString reasonZigStr = Zig::toZigString(reason);
         Bun__WebSocketClient__close(this->m_connectedWebSocket.client, code, &reasonZigStr);
         updateHasPendingActivity();
-        // this->m_bufferedAmount = this->m_connectedWebSocket.client->getBufferedAmount();
         break;
     }
     case ConnectedWebSocketKind::ClientSSL: {
+        m_bufferedAmount = clampToUnsigned(Bun__WebSocketClientTLS__getBufferedAmount(this->m_connectedWebSocket.clientSSL));
         ZigString reasonZigStr = Zig::toZigString(reason);
         Bun__WebSocketClientTLS__close(this->m_connectedWebSocket.clientSSL, code, &reasonZigStr);
         updateHasPendingActivity();
-        // this->m_bufferedAmount = this->m_connectedWebSocket.clientSSL->getBufferedAmount();
         break;
     }
     // case ConnectedWebSocketKind::Server: {
@@ -1036,11 +1041,15 @@
     m_state = CLOSING;
     switch (m_connectedWebSocketKind) {
     case ConnectedWebSocketKind::Client: {
+        // Snapshot the backlog before cancel() frees the send buffer, so
+        // bufferedAmount does not reset to 0 (see bufferedAmount()).
+        m_bufferedAmount = clampToUnsigned(Bun__WebSocketClient__getBufferedAmount(this->m_connectedWebSocket.client));
         Bun__WebSocketClient__cancel(this->m_connectedWebSocket.client);
         updateHasPendingActivity();
         break;
     }
     case ConnectedWebSocketKind::ClientSSL: {
+        m_bufferedAmount = clampToUnsigned(Bun__WebSocketClientTLS__getBufferedAmount(this->m_connectedWebSocket.clientSSL));
         Bun__WebSocketClientTLS__cancel(this->m_connectedWebSocket.clientSSL);
         updateHasPendingActivity();
         break;
@@ -1224,7 +1233,24 @@
 
 unsigned WebSocket::bufferedAmount() const
 {
-    return saturateAdd(m_bufferedAmount, m_bufferedAmountAfterClose);
+    // While OPEN, query the live send-buffer size from the connection so
+    // backpressure is observable. Once closed the connection is gone, but the
+    // spec requires bufferedAmount not to reset to 0 — close()/terminate()
+    // snapshot the final backlog into m_bufferedAmount, and send() after close
+    // adds to m_bufferedAmountAfterClose, so the total only ever increases.
+    unsigned buffered = m_bufferedAmount;
+    switch (m_connectedWebSocketKind) {
+    case ConnectedWebSocketKind::Client:
+        buffered = clampToUnsigned(Bun__WebSocketClient__getBufferedAmount(this->m_connectedWebSocket.client));
+        break;
+    case ConnectedWebSocketKind::ClientSSL:
+        buffered = clampToUnsigned(Bun__WebSocketClientTLS__getBufferedAmount(this->m_connectedWebSocket.clientSSL));
+        break;
+    case ConnectedWebSocketKind::None:
+        break;
+    }
+
+    return saturateAdd(buffered, m_bufferedAmountAfterClose);
 }
 
 String WebSocket::protocol() const
@@ -1588,7 +1614,11 @@
 
     bool wasClean = m_state == CLOSING && !unhandledBufferedAmount && code != 0; // WebSocketChannel::CloseEventCodeAbnormalClosure;
     m_state = CLOSED;
-    m_bufferedAmount = unhandledBufferedAmount;
+    // Don't reset the backlog: close()/terminate() already snapshotted the
+    // unsent bytes into m_bufferedAmount, and the spec requires bufferedAmount
+    // not to drop to 0 once closed. Keep whichever is larger.
+    if (unhandledBufferedAmount > m_bufferedAmount)
+        m_bufferedAmount = unhandledBufferedAmount;
     ASSERT(scriptExecutionContext());
     this->m_connectedWebSocketKind = ConnectedWebSocketKind::None;
     this->m_upgradeClient = nullptr;

test/js/web/websocket/websocket-buffered-amount.test.ts

@@ -0,0 +1,104 @@
+import { describe, expect, test } from "bun:test";
+import crypto from "node:crypto";
+import net from "node:net";
+
+const WS_MAGIC = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
+
+// Raw TCP server that completes the WebSocket handshake and then stops reading
+// from the socket (`pause()`), so the client's outbound frames cannot drain to
+// the peer and pile up in the in-process send buffer.
+function nonDrainingServer(): Promise<{ port: number; close: () => void }> {
+  return new Promise((resolve, reject) => {
+    const server = net.createServer(sock => {
+      let buf = "";
+      let upgraded = false;
+      sock.on("data", d => {
+        if (upgraded) return;
+        buf += d.toString("latin1");
+        if (!buf.includes("\r\n\r\n")) return;
+        const key = /sec-websocket-key:\s*(.+)\r\n/i.exec(buf)?.[1]?.trim() ?? "";
+        const accept = crypto
+          .createHash("sha1")
+          .update(key + WS_MAGIC)
+          .digest("base64");
+        sock.write(
+          "HTTP/1.1 101 Switching Protocols\r\n" +
+            "Upgrade: websocket\r\n" +
+            "Connection: Upgrade\r\n" +
+            `Sec-WebSocket-Accept: ${accept}\r\n\r\n`,
+        );
+        upgraded = true;
+        sock.pause(); // never read the client's frames
+      });
+      sock.on("error", () => {});
+    });
+    server.on("error", reject);
+    server.listen(0, "127.0.0.1", () => {
+      const address = server.address() as net.AddressInfo;
+      resolve({ port: address.port, close: () => server.close() });
+    });
+  });
+}
+
+describe("WebSocket.bufferedAmount (client)", () => {
+  test("reflects the backlog queued to a peer that stopped reading", async () => {
+    const { port, close } = await nonDrainingServer();
+    try {
+      const ws = new WebSocket(`ws://127.0.0.1:${port}/`);
+      const { promise, resolve, reject } = Promise.withResolvers<{ atOpen: number; max: number }>();
+      ws.onerror = () => reject(new Error("unexpected error event"));
+      ws.onopen = () => {
+        // Nothing queued yet: the baseline must be 0, not a constant.
+        const atOpen = ws.bufferedAmount;
+        const chunk = Buffer.alloc(64 * 1024, 0x79).toString();
+        let max = atOpen;
+        // 4000 * 64 KiB = ~250 MiB — far more than any socket buffer can accept,
+        // so the excess must queue in-process.
+        for (let i = 0; i < 4000; i++) {
+          ws.send(chunk);
+          if (ws.bufferedAmount > max) max = ws.bufferedAmount;
+        }
+        resolve({ atOpen, max });
+      };
+      const { atOpen, max } = await promise;
+      ws.close();
+
+      // Baseline with nothing queued.
+      expect(atOpen).toBe(0);
+      // Before the fix, bufferedAmount was hard-wired to 0 for the client
+      // WebSocket. It must now track the unsent backlog — which is far larger
+      // than a single 64 KiB frame once the peer stops reading.
+      expect(max).toBeGreaterThan(64 * 1024);
+    } finally {
+      close();
+    }
+  });
+
+  // Per the WHATWG spec, bufferedAmount "does not reset to zero once the
+  // connection closes" — after close() it only increases with further send().
+  test("does not reset to 0 after close() while a backlog is queued", async () => {
+    const { port, close } = await nonDrainingServer();
+    try {
+      const ws = new WebSocket(`ws://127.0.0.1:${port}/`);
+      const { promise, resolve, reject } = Promise.withResolvers<{ beforeClose: number; afterClose: number }>();
+      ws.onerror = () => reject(new Error("unexpected error event"));
+      ws.onopen = () => {
+        const chunk = Buffer.alloc(64 * 1024, 0x7a).toString();
+        for (let i = 0; i < 4000; i++) ws.send(chunk);
+        const beforeClose = ws.bufferedAmount;
+        ws.close();
+        // Reading immediately after close() must retain the queued backlog,
+        // not snap back to 0.
+        const afterClose = ws.bufferedAmount;
+        resolve({ beforeClose, afterClose });
+      };
+      const { beforeClose, afterClose } = await promise;
+
+      expect(beforeClose).toBeGreaterThan(64 * 1024);
+      // The backlog must survive the close() transition.
+      expect(afterClose).toBe(beforeClose);
+    } finally {
+      close();
+    }
+  });
+});