http_jsc: park the WebSocket upgrade client refs as OwnedRef fields (and Cell handshake state)

Closed
RO
3mo ago
+298-309

Stacked on #37665 (the base of this PR is that branch). Two commits: the original Cell fields change, and, following review ("can we use RAII types instead of ref_ and unref()"), a second commit that turns the refs held on the client into OwnedRef values.

What

1. Handshake state in Cells (first commit)

HTTPClient<SSL> in src/http_jsc/websocket_client/WebSocketUpgradeClient.rs (the WebSocket upgrade client) stored its handshake state as plain fields, and every socket handler, which only has a ThisPtr<Self> (&Self), updated them through a separate unsafe { (*this.as_ptr()).state = State::Failed } with its own SAFETY comment. tcp, to_send_len and state are now Cells, the shape the connected WebSocket<SSL> in websocket_client.rs already uses for the same fields, and every access goes through get / set; the Bun__WebSocketHTTP(S)Client__connect shims take the WebCore::WebSocket handle as &CppWebSocket, the mapping the same signature already uses for JSGlobalObject*; dispatch_handshake and the two did_connect* calls go through CppWebSocket::opaque_ref like the other dispatch sites in the file; the two handle_writable decrements become saturating_sub. That commit removes 21 unsafe blocks and adds 2.

2. The refs held on the client become values (second commit)

Two parties hold a ref on an HTTPClient while it exists, and both were implemented as hand-paired ref_() / Self::deref() calls:

  • the count connect() allocated with stood for the uSockets userdata pointer, and was released by hand at the four failure exits of connect(), in handle_close, handle_connect_error, cancel (only when the ext slot still held the pointer) and the non-tunnel hand-off in process_response;
  • the "+1 for cpp_websocket" taken at the two success exits of connect() stood for WebCore::WebSocket's m_upgradeClient, and was released wherever outgoing_websocket.take() returned Some (cancel, dispatch_abrupt_close, both hand-offs).

Now:

struct CppOwner<const SSL: bool> {
    websocket: NonNull<CppWebSocket>,
    _client_ref: OwnedRef<HTTPClient<SSL>>,
}

pub struct HTTPClient<const SSL: bool> {
    ref_count: Cell<u32>,
    socket_ref: Cell<Option<OwnedRef<Self>>>,   // the socket userdata's ref, parked here
    cpp_owner: Cell<Option<CppOwner<SSL>>>,     // replaces outgoing_websocket
    ...
}
  • connect() allocates with OwnedRef::new, so the local is the initial ref: each of the four failure exits is now a plain return None and the local's drop destroys the client. On success, park_refs() moves the local into socket_ref and parks a clone of it, paired with the WebSocket pointer, in cpp_owner; C++ receives and stores the same raw pointer as before. (The field is None until then; nothing can observe that window, because the socket functions install the client as userdata only after the connect call returns, and a failure before that destroys it.)
  • Every release site takes the matching field and lets the value drop at the point where the deref used to be, after the last call into C++: handle_close and handle_connect_error end with drop(this.socket_ref.take()); dispatch_abrupt_close takes cpp_owner, calls did_abrupt_close, and drops it (it is a safe fn now); the tunnel hand-off takes cpp_owner and drops it after did_connect_with_tunnel while socket_ref stays parked, since the client keeps forwarding socket data until handle_close; the non-tunnel hand-off drops cpp_owner and then socket_ref, the order of the two derefs it replaces. cancel() keeps its ref_guard bracket and still retires the socket only when the ext slot held the pointer (a close callback further up the stack otherwise releases socket_ref itself), so its behaviour is unchanged too. Reads of the C++ handle go through a small outgoing_websocket() accessor.
  • deinit asserts both fields are empty: a parked ref would have kept the count above zero, so destruction can only happen once both holders have let go. clear_data() does not touch either field; the release points are exactly the sites above. No ref_() or deref() on HTTPClient remains.
  • The same treatment for the proxy tunnel: WebSocketProxyTunnel::init returns OwnedRef<WebSocketProxyTunnel> (built with OwnedRef::new), WebSocketProxy.tunnel is Option<OwnedRef<..>>, and the proxy's Drop calls shutdown() and lets the field release the ref, in the order it used by hand. start_proxy_tls_handshake drops the local on its failure exits. One of those exits (the proxy was taken while start() ran) used to drop a bare NonNull and leak the tunnel, SSL state included; it now destroys it. That is the one behaviour change, and it is a leak fix: nothing else references the tunnel at that point and the client it points back to is still live while it is destroyed.

The second commit removes 18 unsafe blocks and adds 1 in the upgrade client (153 on main, 117 now across both commits), and dispatch_abrupt_close stops being an unsafe fn. The struct grows by two words (the two parked refs), which is what memoryCost reports for it.

Why

Who holds a ref on the client, and where each one is released, is now stated by two fields and readable from their types: a retirement path cannot forget to release (the field is taken), cannot release twice (it is None the second time) and cannot release the wrong one, and deinit checks the invariant. The Cells say which state the handlers update through a shared self-pointer. It is zero-cost: an OwnedRef is one pointer, a parked ref is the same increment the ref_() call was, and each take() + drop is the same decrement the removed deref was; Cell accessors on these Copy fields compile to the same loads and stores as the raw projections did.

Related to #30777: connect() is still an unsafe fn taking raw header pointers, so that issue stays open for a follow-up.

Part of a series of small type-system hardening changes.

Verification

cargo check and cargo clippy are clean for bun_http_jsc. Debug build succeeds. bun bd test on the six files that exercise this code (websocket-client.test.ts, websocket-proxy.test.ts, websocket-close-connecting.test.ts, websocket-proxy-tunnel-upgrade-leak.test.ts, websocket-proxy-tunnel-client-leak.test.ts, websocket-proxy-close-reentrancy.test.ts): 78 pass, 4 skip (Docker-only Squid proxy cases), 0 fail; the two tunnel leak tests are the ones that count new / destroy of these objects. websocket.test.js also passes apart from should connect many times over https, which times out identically on main in this container (300 TLS connections on a debug ASAN build on a loaded machine); run alone it passes here too, with the WebSocket instance count returning to its baseline.

cargo check --target x86_64-pc-windows-msvc also passes for bun_http_jsc on this branch, so no platform-gated caller of the removed hand-written ref/deref entry points remains.