uws: make stream-buffer conversion a one-shot ownership transfer

.github/workflows/miri.yml

@@ -22,6 +22,7 @@ on:
       - "src/ptr/**"
       - "src/resolve_builtins/**"
       - "src/shell_parser/**"
+      - "src/uws_sys/**"
       - "src/wyhash/**"
       - "scripts/rust-miri.ts"
       - "Cargo.toml"

scripts/rust-miri.ts

@@ -44,6 +44,7 @@ const MIRI_CRATES = [
   "bun_ptr",
   "bun_resolve_builtins",
   "bun_shell_parser",
+  "bun_uws_sys",
   "bun_wyhash",
 ];
 

src/runtime/socket/uws_jsc.rs

@@ -129,12 +129,12 @@ pub(crate) unsafe extern "C" fn us_socket_buffered_js_write(
     // `&mut *socket` / `&mut *buffer` under Stacked Borrows, so raw pointers with
     // no uniqueness assertion are used throughout.
 
-    // Convert `data`/`encoding` BEFORE materializing the stream buffer into an owning
-    // `Vec<u8>`: the conversion can run arbitrary JS (toString/Symbol.toPrimitive,
-    // Request/Response body coercion) which can re-enter this function on the same
-    // socket. Taking the buffer first would leave two owning `Vec`s over the same
-    // `list_ptr`; the inner call's realloc would free the allocation out from under
-    // the outer frame (use-after-free).
+    // Convert `data`/`encoding` BEFORE taking the stream buffer: the conversion can
+    // run arbitrary JS (toString/Symbol.toPrimitive, Request/Response body coercion)
+    // which can re-enter this function on the same socket. A re-entrant call while
+    // this frame holds the taken buffer would take an empty buffer, stash its own
+    // data there, and the final `update` below would drop that data and reorder the
+    // stream.
     let node_buffer: BlobOrStringOrBuffer = if data.is_undefined() {
         BlobOrStringOrBuffer::StringOrBuffer(StringOrBuffer::EMPTY)
     } else {
@@ -169,9 +169,9 @@ pub(crate) unsafe extern "C" fn us_socket_buffered_js_write(
     }
 
     // SAFETY: caller (JSNodeHTTPServerSocket.cpp) guarantees `buffer` is valid for the call.
-    // No JS executes between here and the `update()` below, so this owning `Vec` is the
-    // sole owner of `list_ptr` for the remainder of the function.
-    let mut stream_buffer = unsafe { &mut *buffer }.to_stream_buffer();
+    // The take nulls the raw parts, so this owning `Vec` is the allocation's sole owner;
+    // no JS executes between here and the `update()` below (see the ordering note above).
+    let mut stream_buffer = unsafe { &mut *buffer }.take_stream_buffer();
     let mut total_written: usize = 0;
 
     // Labeled block + post-block cleanup so the `buffer.update` / `buffer.wrote`

src/uws_sys/us_socket_t.rs

@@ -517,8 +517,16 @@ pub struct StreamBuffer {
     pub cursor: usize,
 }
 
+// Ownership invariant for the raw parts: a non-null `list_ptr` means the
+// struct owns the raw parts of exactly one `Vec<u8>` decomposed by `update`.
+// `take_stream_buffer` nulls the parts when it transfers ownership out, and
+// `update` drops any buffer still owned before overwriting them, so the
+// invariant holds on every path (C++ only ever zero-initializes the struct).
 impl us_socket_stream_buffer_t {
     pub fn update(&mut self, stream_buffer: StreamBuffer) {
+        // Drop whatever is currently owned so overwriting the raw parts below
+        // can't leak a previous buffer.
+        drop(self.take_stream_buffer());
         // Decompose the Vec<u8> backing `stream_buffer.list` into raw parts so
         // the C side can read ptr/len/cap directly.
         let mut list = core::mem::ManuallyDrop::new(stream_buffer.list);
@@ -536,36 +544,39 @@ impl us_socket_stream_buffer_t {
         self.total_bytes_written = self.total_bytes_written.saturating_add(written);
     }
 
-    pub fn to_stream_buffer(&self) -> StreamBuffer {
-        StreamBuffer {
-            list: if !self.list_ptr.is_null() {
-                unsafe {
-                    // SAFETY: list_ptr/list_len/list_cap were produced by decomposing a
-                    // Vec<u8> in `update`; global allocator (mimalloc) matches.
-                    Vec::from_raw_parts(self.list_ptr, self.list_len, self.list_cap)
-                }
-            } else {
-                Vec::new()
-            },
-            cursor: self.cursor,
-        }
+    /// One-shot ownership transfer: rebuilds the owned `Vec<u8>` from the raw
+    /// parts and nulls them, so the returned `StreamBuffer` is the allocation's
+    /// sole owner and a second take (or a later `destroy`) sees an empty
+    /// buffer. `total_bytes_written` is cumulative socket state, not buffer
+    /// contents, and survives the take.
+    pub fn take_stream_buffer(&mut self) -> StreamBuffer {
+        let list = if !self.list_ptr.is_null() {
+            // SAFETY: per the ownership invariant above, the raw parts came
+            // from a Vec<u8> decomposed in `update` (global allocator
+            // matches), and nulling them below ends this struct's ownership.
+            unsafe { Vec::from_raw_parts(self.list_ptr, self.list_len, self.list_cap) }
+        } else {
+            Vec::new()
+        };
+        let cursor = self.cursor;
+        self.list_ptr = ptr::null_mut();
+        self.list_len = 0;
+        self.list_cap = 0;
+        self.cursor = 0;
+        StreamBuffer { list, cursor }
     }
 
     /// Explicit teardown — this struct is `#[repr(C)]` and freed via the
-    /// exported `us_socket_free_stream_buffer`, so no `Drop` impl.
+    /// exported `us_socket_free_stream_buffer`, so no `Drop` impl. Idempotent:
+    /// the take nulls the raw parts, so a second call is a no-op.
     ///
     /// SAFETY: `this` must point to a live `us_socket_stream_buffer_t` whose
     /// `list_ptr`/`list_cap` were produced by `update` (decomposed `Vec<u8>` on
-    /// the global mimalloc allocator). Not called more than once.
+    /// the global mimalloc allocator).
     pub unsafe fn destroy(this: *mut Self) {
         // SAFETY: caller contract — `this` is non-null and exclusively borrowed
         let this = unsafe { &mut *this };
-        if !this.list_ptr.is_null() {
-            unsafe {
-                // SAFETY: list_ptr/list_cap came from a decomposed Vec<u8> (global mimalloc).
-                drop(Vec::from_raw_parts(this.list_ptr, 0, this.list_cap));
-            }
-        }
+        drop(this.take_stream_buffer());
     }
 }
 
@@ -575,3 +586,97 @@ pub(crate) extern "C" fn us_socket_free_stream_buffer(buffer: *mut us_socket_str
     unsafe { us_socket_stream_buffer_t::destroy(buffer) };
 }
 // us_socket_buffered_js_write moved to src/runtime/socket/uws_jsc.rs
+
+// Pure Rust (no FFI at test runtime), so these run under `cargo miri test`,
+// which catches double-frees and leaks of the raw-part round-trips.
+#[cfg(test)]
+mod stream_buffer_tests {
+    use super::{StreamBuffer, us_socket_stream_buffer_t};
+
+    // https://github.com/oven-sh/bun/issues/31971
+    #[test]
+    fn take_transfers_ownership_once() {
+        let mut raw = us_socket_stream_buffer_t::default();
+        raw.update(StreamBuffer {
+            list: vec![1, 2, 3],
+            cursor: 1,
+        });
+
+        let first = raw.take_stream_buffer();
+        assert_eq!(first.list, [1, 2, 3]);
+        assert_eq!(first.cursor, 1);
+
+        // The take nulled the raw parts: a second take yields an empty buffer
+        // instead of a second owner of the same allocation.
+        let second = raw.take_stream_buffer();
+        assert!(second.list.is_empty());
+        assert_eq!(second.list.capacity(), 0);
+        assert_eq!(second.cursor, 0);
+        assert!(raw.list_ptr.is_null());
+    }
+
+    #[test]
+    fn destroy_after_take_is_a_noop() {
+        let mut raw = us_socket_stream_buffer_t::default();
+        raw.update(StreamBuffer {
+            list: vec![4, 5, 6],
+            cursor: 0,
+        });
+        let taken = raw.take_stream_buffer();
+        // SAFETY: `raw` is a live stack value whose parts came from `update`.
+        unsafe { us_socket_stream_buffer_t::destroy(&mut raw) };
+        drop(taken);
+    }
+
+    #[test]
+    fn destroy_is_idempotent() {
+        let mut raw = us_socket_stream_buffer_t::default();
+        raw.update(StreamBuffer {
+            list: vec![7; 32],
+            cursor: 0,
+        });
+        // SAFETY: `raw` is a live stack value whose parts came from `update`.
+        unsafe { us_socket_stream_buffer_t::destroy(&mut raw) };
+        unsafe { us_socket_stream_buffer_t::destroy(&mut raw) };
+        assert!(raw.list_ptr.is_null());
+    }
+
+    #[test]
+    fn update_drops_the_previously_owned_buffer() {
+        let mut raw = us_socket_stream_buffer_t::default();
+        raw.update(StreamBuffer {
+            list: vec![1; 16],
+            cursor: 2,
+        });
+        raw.update(StreamBuffer {
+            list: vec![9, 9],
+            cursor: 0,
+        });
+        let taken = raw.take_stream_buffer();
+        assert_eq!(taken.list, [9, 9]);
+        assert_eq!(taken.cursor, 0);
+    }
+
+    #[test]
+    fn total_bytes_written_survives_the_take() {
+        let mut raw = us_socket_stream_buffer_t::default();
+        raw.update(StreamBuffer {
+            list: vec![1, 2],
+            cursor: 0,
+        });
+        raw.wrote(5);
+        drop(raw.take_stream_buffer());
+        assert_eq!(raw.total_bytes_written, 5);
+    }
+
+    #[test]
+    fn empty_capacity_round_trips_as_null() {
+        let mut raw = us_socket_stream_buffer_t::default();
+        raw.update(StreamBuffer {
+            list: Vec::new(),
+            cursor: 0,
+        });
+        assert!(raw.list_ptr.is_null());
+        assert!(raw.take_stream_buffer().list.is_empty());
+    }
+}

test/internal/uws-stream-buffer-ownership.test.ts

@@ -0,0 +1,52 @@
+// Source-text consistency check for oven-sh/bun#31971.
+//
+// `us_socket_stream_buffer_t` hands the raw parts of a `Vec<u8>` back and
+// forth across the C++ boundary. Before #31971 the conversion back to an
+// owning `StreamBuffer` was `pub fn to_stream_buffer(&self)`: it rebuilt a
+// `Vec` via `Vec::from_raw_parts` without clearing `list_ptr`/`list_len`/
+// `list_cap`, so entirely safe Rust could mint two owners of the same
+// allocation (double free / use-after-free). The fix makes the conversion a
+// one-shot ownership transfer, `pub fn take_stream_buffer(&mut self)`, which
+// nulls the raw parts it transfers out.
+//
+// The primary guard is behavioral: `stream_buffer_tests` in
+// src/uws_sys/us_socket_t.rs runs under `cargo miri test -p bun_uws_sys`
+// (MIRI_CRATES in scripts/rust-miri.ts, CI lane .github/workflows/miri.yml),
+// where Miri deterministically reports the double free / leak if the take
+// ever stops clearing the raw parts. This file is a belt-and-suspenders lint
+// that pins the same invariant at the source-text layer so a reviewer can see
+// it without running cargo — same pattern + test/internal/ placement as
+// ban-words.test.ts and dead-code-escapes.test.ts (a coding-convention lint,
+// not a behavioral test). Not under test/regression/issue/ because the
+// conversion was unsound from the day it landed; there is no prior release
+// where it worked.
+
+import { expect, test } from "bun:test";
+import { readFileSync } from "node:fs";
+import { join } from "node:path";
+
+// test/internal/ is one level under repo root.
+const repoRoot = join(import.meta.dir, "..", "..");
+
+// Whitespace-tolerant so benign reformatting doesn't break the lint; only
+// signatures and call-site counts are asserted.
+const source = readFileSync(join(repoRoot, "src/uws_sys/us_socket_t.rs"), "utf8").replace(/\s+/g, " ");
+
+test("stream-buffer conversion is a one-shot take (&mut self), not a safe &self copy (#31971)", () => {
+  // The conversion must require exclusive access so it can null the raw
+  // parts whose ownership it transfers out.
+  expect(source).toMatch(/\bpub\s+fn\s+take_stream_buffer\s*\(\s*&\s*mut\s+self\s*\)\s*->\s*StreamBuffer\b/);
+  // The old `&self` conversion rebuilt an owning Vec without clearing the
+  // parts, so two calls (or one call plus `destroy`) double-freed.
+  expect(source).not.toMatch(/\bfn\s+to_stream_buffer\b/);
+  expect(source).not.toMatch(/\bfn\s+take_stream_buffer\s*\(\s*&\s*self\b/);
+});
+
+test("take_stream_buffer is the only place that rebuilds the Vec from raw parts (#31971)", () => {
+  // A single rebuild site keeps the ownership transfer auditable: `destroy`
+  // and `update` route through the take instead of re-deriving a Vec from
+  // `list_ptr` themselves. (`stream_buffer_tests` at the bottom of the file
+  // exercises the behavior; this only pins the structure.)
+  const rebuilds = source.match(/Vec::from_raw_parts/g) ?? [];
+  expect(rebuilds).toHaveLength(1);
+});