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

scripts/rust-miri.ts

@@ -52,6 +52,7 @@ const MIRI_CRATES = [
   "bun_shell_parser",
   "bun_threading",
   "bun_url",
+  "bun_uws_sys",
   "bun_wyhash",
 ];
 

src/runtime/socket/uws_jsc.rs

@@ -141,12 +141,7 @@ 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).
+    // Coerce `data` first: it can run JS that re-enters this function on the same socket.
     let node_buffer: BlobOrStringOrBuffer = if data.is_undefined() {
         BlobOrStringOrBuffer::StringOrBuffer(StringOrBuffer::EMPTY)
     } else {
@@ -180,10 +175,9 @@ 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();
+    // SAFETY: the C++ caller keeps `buffer` valid for the call. No JS runs
+    // between here and the `update()` below.
+    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

@@ -617,8 +617,10 @@ pub struct StreamBuffer {
     pub cursor: usize,
 }
 
+// Invariant: a non-null `list_ptr` means this struct owns one `Vec<u8>` decomposed by `update`.
 impl us_socket_stream_buffer_t {
     pub fn update(&mut self, stream_buffer: StreamBuffer) {
+        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);
@@ -636,35 +638,28 @@ 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,
-        }
+    /// Moves the owned buffer out and nulls the raw parts. `total_bytes_written` survives.
+    pub fn take_stream_buffer(&mut self) -> StreamBuffer {
+        let list = if !self.list_ptr.is_null() {
+            // SAFETY: the raw parts came from a Vec<u8> decomposed in `update`
+            // and are nulled below, so this is the only owner.
+            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.
+    /// Teardown called from C++ via `us_socket_free_stream_buffer`. Idempotent.
     ///
-    /// SAFETY: `list_ptr`/`list_cap` were produced by `update` (decomposed
-    /// `Vec<u8>` on the global mimalloc allocator). Not called more than once.
+    /// SAFETY: `list_ptr`/`list_cap` were produced by `update`.
     pub(crate) unsafe fn destroy(&mut self) {
-        if !self.list_ptr.is_null() {
-            unsafe {
-                // SAFETY: list_ptr/list_cap came from a decomposed Vec<u8> (global mimalloc).
-                drop(Vec::from_raw_parts(self.list_ptr, 0, self.list_cap));
-            }
-            self.list_ptr = core::ptr::null_mut();
-            self.list_cap = 0;
-        }
+        drop(self.take_stream_buffer());
     }
 }
 
@@ -674,3 +669,103 @@ extern "C" fn us_socket_free_stream_buffer(buffer: *mut us_socket_stream_buffer_
     unsafe { (*buffer).destroy() };
 }
 // us_socket_buffered_js_write moved to src/runtime/socket/uws_jsc.rs
+
+// FFI-free, so `cargo miri test` runs these (see scripts/rust-miri.ts).
+#[cfg(test)]
+mod stream_buffer_tests {
+    use super::{StreamBuffer, us_socket_stream_buffer_t};
+
+    // The C++ side zero-initializes the struct.
+    fn empty() -> us_socket_stream_buffer_t {
+        us_socket_stream_buffer_t {
+            list_ptr: core::ptr::null_mut(),
+            list_cap: 0,
+            list_len: 0,
+            total_bytes_written: 0,
+            cursor: 0,
+        }
+    }
+
+    // https://github.com/oven-sh/bun/issues/31971
+    #[test]
+    fn take_transfers_ownership_once() {
+        let mut raw = empty();
+        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);
+
+        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 = empty();
+        raw.update(StreamBuffer {
+            list: vec![4, 5, 6],
+            cursor: 0,
+        });
+        let taken = raw.take_stream_buffer();
+        unsafe { raw.destroy() };
+        drop(taken);
+    }
+
+    #[test]
+    fn destroy_is_idempotent() {
+        let mut raw = empty();
+        raw.update(StreamBuffer {
+            list: vec![7; 32],
+            cursor: 0,
+        });
+        unsafe { raw.destroy() };
+        unsafe { raw.destroy() };
+        assert!(raw.list_ptr.is_null());
+    }
+
+    #[test]
+    fn update_drops_the_previously_owned_buffer() {
+        let mut raw = empty();
+        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 = empty();
+        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 = empty();
+        raw.update(StreamBuffer {
+            list: Vec::new(),
+            cursor: 0,
+        });
+        assert!(raw.list_ptr.is_null());
+        assert!(raw.take_stream_buffer().list.is_empty());
+    }
+}

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

@@ -0,0 +1,37 @@
+// Source-text lint for oven-sh/bun#31971.
+//
+// `us_socket_stream_buffer_t` hands the raw parts of a `Vec<u8>` across the
+// C++ boundary. The conversion back to an owning `StreamBuffer` used to be
+// `pub fn to_stream_buffer(&self)`, which rebuilt a `Vec` via
+// `Vec::from_raw_parts` without clearing the parts, so safe Rust could mint
+// two owners of one allocation. The fix is a one-shot
+// `pub fn take_stream_buffer(&mut self)` that nulls the parts it moves out.
+//
+// The behavioral guard is `stream_buffer_tests` in src/uws_sys/us_socket_t.rs,
+// which `cargo miri test -p bun_uws_sys` runs (MIRI_CRATES in
+// scripts/rust-miri.ts, the miri job in .github/workflows/rust-lints.yml).
+// This file pins the same invariant at the source-text layer so a reviewer can
+// see it without running cargo.
+
+import { expect, test } from "bun:test";
+import { readFileSync } from "node:fs";
+import path from "node:path";
+
+const root = path.resolve(import.meta.dir, "..", "..", "..");
+
+// Whitespace-tolerant so reformatting doesn't break the lint.
+const source = readFileSync(path.join(root, "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/);
+  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)", () => {
+  // `destroy` and `update` route through the take instead of re-deriving a
+  // Vec from `list_ptr` themselves.
+  const rebuilds = source.match(/Vec::from_raw_parts/g) ?? [];
+  expect(rebuilds).toHaveLength(1);
+});