ws: support upgrade and unexpected-response events

src/http_jsc/websocket_client/CppWebSocket.rs

@@ -22,6 +22,19 @@ bun_opaque::opaque_ffi! {
     pub struct CppWebSocket;
 }
 
+/// A single HTTP response header passed across the FFI to
+/// `WebSocket__didReceiveHandshakeResponse`. Name/value point into the
+/// PicoHTTP-parsed response head, which the caller keeps alive for the
+/// duration of the synchronous dispatch. Matches `WebCore::WebSocket::
+/// HandshakeRawHeader` (src/jsc/bindings/webcore/WebSocket.h).
+#[repr(C)]
+pub struct RawHeader {
+    pub name_ptr: *const u8,
+    pub name_len: usize,
+    pub value_ptr: *const u8,
+    pub value_len: usize,
+}
+
 // FFI surface for `WebCore::WebSocket` (src/jsc/bindings/webcore/WebSocket.cpp).
 // Kept private to this module — the safe wrappers below are the only callers.
 //
@@ -45,6 +58,16 @@ unsafe extern "C" {
         deflate_params: *const websocket_deflate::Params,
     );
     safe fn WebSocket__didAbruptClose(websocket_context: &CppWebSocket, reason: ErrorCode);
+    fn WebSocket__didReceiveHandshakeResponse(
+        websocket_context: &CppWebSocket,
+        status_code: u16,
+        status_message: *const u8,
+        status_message_len: usize,
+        headers: *const RawHeader,
+        headers_len: usize,
+        body: *const u8,
+        body_len: usize,
+    );
     fn WebSocket__didClose(websocket_context: &CppWebSocket, code: u16, reason: *const BunString);
     fn WebSocket__didReceiveText(
         websocket_context: &CppWebSocket,
@@ -78,6 +101,39 @@ impl CppWebSocket {
         event_loop.exit();
     }
 
+    /// Dispatch the native `'handshake'` event to the `ws` shim.
+    ///
+    /// `status_message`, `headers`, and `body` are borrowed for the call only
+    /// — the C++ side copies whatever it needs into JS values synchronously.
+    // Forwards raw pointers to C++ without dereferencing on the Rust side.
+    #[allow(clippy::not_unsafe_ptr_arg_deref)]
+    pub(crate) fn did_receive_handshake_response(
+        &self,
+        status_code: u16,
+        status_message: &[u8],
+        headers: &[RawHeader],
+        body: &[u8],
+    ) {
+        // SAFETY: VirtualMachine::get() returns the live current-thread VM;
+        // event_loop() yields its raw event-loop pointer (live for VM lifetime).
+        let event_loop = VirtualMachine::get().event_loop_mut();
+        event_loop.enter();
+        // SAFETY: self is a valid C++ WebCore::WebSocket; the slices outlive the call.
+        unsafe {
+            WebSocket__didReceiveHandshakeResponse(
+                self,
+                status_code,
+                status_message.as_ptr(),
+                status_message.len(),
+                headers.as_ptr(),
+                headers.len(),
+                body.as_ptr(),
+                body.len(),
+            )
+        };
+        event_loop.exit();
+    }
+
     pub(crate) fn did_close(&self, code: u16, reason: &mut BunString) {
         // SAFETY: VirtualMachine::get() returns the live current-thread VM;
         // event_loop() yields its raw event-loop pointer (live for VM lifetime).

src/js/builtins/BunBuiltinNames.h

@@ -164,6 +164,7 @@ using namespace JSC;
     macro(pullAlgorithm) \
     macro(pulling) \
     macro(queue) \
+    macro(rawHeaders) \
     macro(read) \
     macro(readIntoRequests) \
     macro(readRequests) \
@@ -193,6 +194,8 @@ using namespace JSC;
     macro(started) \
     macro(state) \
     macro(status) \
+    macro(statusCode) \
+    macro(statusMessage) \
     macro(statusText) \
     macro(storedError) \
     macro(strategy) \

src/js/thirdparty/ws.js

@@ -80,8 +80,74 @@
   error: 4,
   ping: 5,
   pong: 6,
+  upgrade: 7,
+  "unexpected-response": 8,
 };
 
+// Stable singleton used by `#armNativeBridge` to install the native-side
+// forwarder through `#onOrOnce` without pushing a user-visible listener. It
+// must be identity-stable so `super.off` can remove it later if needed.
+function noopBridgeListener() {}
+
+// Headers for which Node's http.IncomingMessage keeps only the first value and
+// discards later duplicates, matching lib/_http_incoming.js matchKnownFields.
+// Real `ws` hands consumers that Node object, so the synthetic IncomingMessage
+// passed to 'unexpected-response' must coalesce the same way.
+const kSingletonHeaders = new Set([
+  "age",
+  "authorization",
+  "content-length",
+  "content-type",
+  "etag",
+  "expires",
+  "from",
+  "host",
+  "if-modified-since",
+  "if-unmodified-since",
+  "last-modified",
+  "location",
+  "max-forwards",
+  "proxy-authorization",
+  "referer",
+  "retry-after",
+  "server",
+  "user-agent",
+]);
+
+let lazyReadable;
+function makeHandshakeResponse(statusCode, statusMessage, rawHeaders, body) {
+  lazyReadable ??= require("node:stream").Readable;
+  const res = new lazyReadable({ read() {} });
+  const headers = (res.headers = { __proto__: null });
+  res.rawHeaders = rawHeaders;
+  for (let i = 0; i < rawHeaders.length; i += 2) {
+    const lower = rawHeaders[i].toLowerCase();
+    const value = rawHeaders[i + 1];
+    const prev = headers[lower];
+    if (lower === "set-cookie") {
+      if (prev === undefined) headers[lower] = [value];
+      else prev.push(value);
+    } else if (prev === undefined) {
+      headers[lower] = value;
+    } else if (kSingletonHeaders.has(lower)) {
+      // Keep the first occurrence; discard the duplicate.
+    } else if (lower === "cookie") {
+      headers[lower] = prev + "; " + value;
+    } else {
+      headers[lower] = prev + ", " + value;
+    }
+  }
+  res.statusCode = statusCode;
+  res.statusMessage = statusMessage;
+  res.httpVersion = "1.1";
+  res.httpVersionMajor = 1;
+  res.httpVersionMinor = 1;
+  res.socket = res.connection = null;
+  if (body && body.length) res.push(body);
+  res.push(null);
+  return res;
+}
+
 const emittedWarnings = new Set();
 function emitWarning(type, message) {
   if (emittedWarnings.has(type)) return;
@@ -124,6 +190,7 @@
   #paused = false;
   #fragments = false;
   #binaryType = "nodebuffer";
+  #unexpectedResponseEmitted = false;
   // Bitset to track whether event handlers are set.
   #eventId = 0;
 
@@ -266,14 +333,105 @@
     }
     let ws = (this.#ws = new WebSocket(url, wsOptions));
     ws.binaryType = "nodebuffer";
+    // NOTE: the native 'handshake' listener is registered lazily from
+    // #ensureHandshakeListener() when the user subscribes to 'upgrade' or
+    // 'unexpected-response'. Keeping it off by default means callers that
+    // only listen to 'open'/'message'/'close' never exercise the Zig
+    // handshake-dispatch path (which has historically been fragile under
+    // ASAN, e.g. ws-proxy.test.ts).
 
     return ws;
   }
 
+  #handshakeListenerRegistered = false;
+  #ensureHandshakeListener() {
+    if (this.#handshakeListenerRegistered) return;
+    this.#handshakeListenerRegistered = true;
+    this.#ws.addEventListener("handshake", event => this.#onHandshake(event.data), onceObject);
+  }
+
+  // ws emits `'upgrade'` / `'unexpected-response'` with an `http.ClientRequest`
+  // as the first argument. We bypass node:http and talk to the native
+  // WebSocket directly, so there is no real ClientRequest to expose. Pass a
+  // minimal stub — `method`, `path`, and the no-op header helpers — so user
+  // code that inspects the request object (e.g. `req.getHeader(...)`) does
+  // not crash. Build it lazily on first emission to keep the common open/
+  // close path zero-cost.
+  #syntheticRequest;
+  #getSyntheticRequest() {
+    let req = this.#syntheticRequest;
+    if (req) return req;
+    const url = this.#ws?.url;
+    let path = "/";
+    try {
+      if (url) {
+        const parsed = new URL(url);
+        // Node's http.ClientRequest.path is pathname + search (no fragment)
+        // per RFC 7230 §5.3. Preserve the query string so handlers that log
+        // or reconstruct the failing URL see the full request target.
+        path = (parsed.pathname || "/") + (parsed.search || "");
+      }
+    } catch {}
+    req = this.#syntheticRequest = {
+      __proto__: Object.create(EventEmitter.prototype),
+      method: "GET",
+      path,
+      url,
+      headers: { __proto__: null },
+      rawHeaders: [],
+      getHeader() {},
+      getHeaders() {
+        return { __proto__: null };
+      },
+      setHeader() {},
+      removeHeader() {},
+      hasHeader() {
+        return false;
+      },
+      abort() {},
+      end() {},
+      write() {},
+      writeHead() {},
+      headersSent: true,
+      finished: true,
+      socket: null,
+      [Symbol.toStringTag]: "ClientRequest",
+    };
+    return req;
+  }
+
+  #onHandshake(data) {
+    const { statusCode, statusMessage, rawHeaders, body } = data;
+    // On 101, any bytes arriving after the header block are the first
+    // WebSocket frame, not HTTP response body — node's http layer delivers
+    // them as the `head` buffer on the 'upgrade' event and the native
+    // WebSocket client forwards them to the protocol reader on didConnect.
+    // Don't leak them into the `upgrade` event's IncomingMessage stream.
+    const res = makeHandshakeResponse(statusCode, statusMessage, rawHeaders, statusCode === 101 ? null : body);
+    if (statusCode === 101) {
+      // `upgrade` emits `(response)` per ws docs.
+      this.emit("upgrade", res);
+      return;
+    }
+    this.#unexpectedResponseEmitted = true;
+    if (this.listenerCount("unexpected-response") > 0) {
+      // `unexpected-response` emits `(request, response)` per ws docs.
+      this.emit("unexpected-response", this.#getSyntheticRequest(), res);
+    } else {
+      this.emit("error", new Error("Unexpected server response: " + statusCode));
+    }
+  }
+
   #onOrOnce(event, listener, once) {
-    if (event === "unexpected-response" || event === "upgrade" || event === "redirect") {
+    if (event === "redirect") {
       emitWarning(event, "ws.WebSocket '" + event + "' event is not implemented in bun");
     }
+    if (event === "upgrade" || event === "unexpected-response") {
+      // Lazy-register the native handshake listener so callers that never
+      // subscribe to these events don't exercise the Zig handshake dispatch.
+      this.#ensureHandshakeListener();
+      return once ? super.once(event, listener) : super.on(event, listener);
+    }
     const mask = 1 << eventIds[event];
     const hasPersistentListener = mask && (this.#eventId & mask) === mask;
     // Add a native listener if:
@@ -322,6 +480,7 @@
         this.#ws.addEventListener(
           "error",
           err => {
+            if (this.#unexpectedResponseEmitted) return;
             this.emit("error", err);
           },
           once,
@@ -355,6 +514,57 @@
     return this.#onOrOnce(event, listener, onceObject);
   }
 
+  // `on` is the conventional spelling, but ws / EventEmitter consumers also
+  // reach for `addListener` / `prependListener` / `prependOnceListener`. Each
+  // needs to arm the native bridge (the `#onOrOnce` path for standard events
+  // like 'open' / 'message', or `#ensureHandshakeListener()` for 'upgrade' /
+  // 'unexpected-response') — otherwise the handler sits on the EventEmitter
+  // list but the native event that would trigger `this.emit(...)` is never
+  // wired up, and the callback silently never fires.
+  addListener(event, listener) {
+    return this.#onOrOnce(event, listener, undefined);
+  }
+
+  prependListener(event, listener) {
+    this.#armNativeBridge(event);
+    return super.prependListener(event, listener);
+  }
+
+  prependOnceListener(event, listener) {
+    this.#armNativeBridge(event);
+    return super.prependOnceListener(event, listener);
+  }
+
+  // Install the native-side listener that eventually calls `this.emit(event,
+  // …)`, without pushing a new EventEmitter listener onto the list — that
+  // part is the caller's job (e.g. `super.prependListener`). Mirrors the
+  // bridge-installation branch of `#onOrOnce` for the cases where we need
+  // the side effect but not the `super.on`/`super.once` call.
+  #armNativeBridge(event) {
+    if (event === "upgrade" || event === "unexpected-response") {
+      this.#ensureHandshakeListener();
+      return;
+    }
+    // `1 << undefined` evaluates to `1` (ToInt32(undefined) = 0), so a
+    // naive `if (!mask) return` would silently fall through for unknown
+    // event names and install a spurious noopBridgeListener on bit 0.
+    // Guard explicitly against unknown events.
+    if (eventIds[event] === undefined) return;
+    const mask = 1 << eventIds[event];
+    const hasPersistentListener = (this.#eventId & mask) === mask;
+    if (hasPersistentListener) return;
+    // Install the native forwarder once. `#onOrOnce` flips the `#eventId`
+    // bit and registers the native listener as a side effect, but it also
+    // appends the listener we pass to the EventEmitter list. We only want
+    // the side effect here (the caller, e.g. `super.prependListener`, owns
+    // the real listener), so register `noopBridgeListener` and immediately
+    // remove it — leaving the native forwarder installed and the
+    // EventEmitter list untouched. Subsequent calls see
+    // `hasPersistentListener` and skip this branch.
+    this.#onOrOnce(event, noopBridgeListener, undefined);
+    super.off(event, noopBridgeListener);
+  }
+
   send(data, opts, cb) {
     if ($isCallable(opts)) {
       cb = opts;
@@ -422,10 +632,37 @@
 
   // deviation: this does not support `message` with `binaryType = "fragments"`
   addEventListener(type, listener, options) {
+    // 'upgrade' and 'unexpected-response' are Node-style events — they're
+    // emitted on the JS-side EventEmitter by `#onHandshake`, never by the
+    // native WebSocket. Register on `this` and arm the handshake listener.
+    // We deliberately don't wrap the handler in a DOM-style adapter: a
+    // wrapped closure would make `removeEventListener` fail to match the
+    // original listener, leaking the handler. Consumers that reach for
+    // `addEventListener` on a ws shim already accept a mixed API; for
+    // these two Node-only events they receive Node-style
+    // `(response)` / `(request, response)` args.
+    if (type === "upgrade" || type === "unexpected-response") {
+      this.#ensureHandshakeListener();
+      // `addEventListener` returns undefined per the DOM spec / real ws —
+      // use statement form so we don't leak EventEmitter's `this` return.
+      if (options && options.once) {
+        super.once(type, listener);
+      } else {
+        super.on(type, listener);
+      }
+      return;
+    }
     this.#ws.addEventListener(type, listener, options);
   }
 
   removeEventListener(type, listener) {
+    // Symmetric with addEventListener: upgrade/unexpected-response
+    // listeners live on `this` (the EventEmitter), not on the native
+    // WebSocket, so removing them means calling `super.off`.
+    if (type === "upgrade" || type === "unexpected-response") {
+      super.off(type, listener);
+      return;
+    }
     this.#ws.removeEventListener(type, listener);
   }