node:tls: complete write(cb) and end(cb) over a Duplex when the stream completes the ciphertext

packages/bun-types/bun.d.ts

@@ -6824,6 +6824,8 @@ declare module "bun" {
      * Calls the `shutdown(2)` syscall internally.
      *
      * @param halfClose If `true`, only shuts down the write side (allows receiving). If `false` or omitted, shuts down both read and write. Defaults to `false`.
+     * @returns `true` when the transport has taken the shutdown. A TCP socket always returns `true`. A TLS socket returns
+     * `false` while a sealed record still waits for the kernel, or while a wrapped stream holds the close_notify. The `drain` handler follows once it completes.
      * @example
      * ```ts
      * // Stop sending data, but allow receiving
@@ -6833,7 +6835,7 @@ declare module "bun" {
      * socket.shutdown();
      * ```
      */
-    shutdown(halfClose?: boolean): void;
+    shutdown(halfClose?: boolean): boolean;
 
     /**
      * The ready state of the socket.

src/js/node/net.ts

@@ -334,6 +334,7 @@ const kOnreadPendingEnd = Symbol("kOnreadPendingEnd");
 const kOnreadReadRequested = Symbol("kOnreadReadRequested");
 const kOnreadEmptyTail = Buffer.alloc(0);
 const kwriteCallback = Symbol("writeCallback");
+const kshutdownCallback = Symbol("shutdownCallback");
 const kSocketClass = Symbol("kSocketClass");
 
 // A completed write whose status is a negative errno: Node hands it to the write
@@ -387,12 +388,23 @@ function writeErrnoException(negErrno) {
   }
   return er;
 }
-function endNT(socket, callback, err) {
+function endNT(socket, callback, self) {
   // Node's _final half-closes the writable side (sends FIN) and leaves the
   // readable side open; the Duplex's allowHalfOpen drives the eventual destroy.
   // https://github.com/nodejs/node/blob/614050b657e9757c1097aa85f92f2cb51149dc0d/lib/net.js#L500
-  socket.shutdown();
-  callback(err);
+  if (socket.shutdown()) {
+    callback();
+    return;
+  }
+  // A wrapped Duplex still holds the close_notify; its drain or close completes this.
+  self[kshutdownCallback] = callback;
+}
+function completeShutdown(self, socket) {
+  const callback = self[kshutdownCallback];
+  if (callback && (!socket || socket.shutdown())) {
+    self[kshutdownCallback] = null;
+    callback();
+  }
 }
 function emitCloseNT(self, hasError) {
   self.emit("close", hasError);
@@ -417,6 +429,8 @@ function destroyNT(self, err) {
 // Node's wrap 'close' -> destroy(): https://github.com/nodejs/node/blob/v26.3.0/lib/internal/tls/wrap.js#L739-L741
 function onUpgradedClose(self, connection) {
   if (self[kupgraded] !== connection) return;
+  // Before the destroy, so _final still finishes (Node: doClose -> ECANCELED, ignored).
+  completeShutdown(self, null);
   // The stream-level engine reads its transport with no backpressure, so the
   // transport can close after the peer's EOF with plaintext still unread.
   if ((self[kended] || self[kOnreadPendingEnd]) && !self.readableEnded) self.once("end", self[kOnUpgradedClose]);
@@ -696,6 +710,7 @@ const SocketHandlers = {
 
       self[kBytesWritten] = socket.bytesWritten;
     }
+    completeShutdown(self, socket);
   },
   end(socket) {
     const self = socket.data;
@@ -893,6 +908,7 @@ function SocketEmitEndNT(self, _err?) {
       const er = new ErrnoException(errErrno, "read") as Error & { code?: string };
       if (typeof er.code === "string" && /^E[A-Z0-9]+$/.test(er.code)) {
         self.destroy(er);
+        completeShutdown(self, null);
         return;
       }
     }
@@ -911,6 +927,7 @@ function SocketEmitEndNT(self, _err?) {
       // Any other coded error (ETIMEDOUT, EPIPE, ...) keeps its identity.
       self.destroy(_err);
     }
+    completeShutdown(self, null);
     return;
   }
   if (!self[kended]) {
@@ -931,6 +948,7 @@ function SocketEmitEndNT(self, _err?) {
     self[kwriteCallback] = null;
     pendingWrite(_err ?? $ERR_SOCKET_CLOSED());
   }
+  if (self[kclosed] || self.destroyed) completeShutdown(self, null);
 }
 
 // --- SNICallback dispatch helpers (hoisted: no per-handshake closures) ---
@@ -1537,6 +1555,7 @@ const SocketHandlers2 = {
         self._pendingData = null;
       }
     }
+    completeShutdown(self, socket);
   },
   end(socket) {
     $debug("Bun.Socket end");
@@ -1593,6 +1612,7 @@ const SocketHandlers2 = {
         // enum values are filtered out in NewSocket::on_close).
         self.destroy(err);
       }
+      completeShutdown(self, null);
       return;
     }
     if (!leftToTLSSocket && !deferEndForOnreadTail(self)) finishSocketEnd(self);
@@ -1604,6 +1624,7 @@ const SocketHandlers2 = {
       self[kwriteCallback] = null;
       pendingWrite($ERR_SOCKET_CLOSED());
     }
+    completeShutdown(self, null);
   },
   handshake(socket, success, verifyError) {
     $debug("Bun.Socket handshake");
@@ -1827,6 +1848,7 @@ function Socket(options?): void {
   // https://github.com/nodejs/node/blob/v26.3.0/lib/net.js#L401
   this[kTimeout] = null;
   this[kwriteCallback] = undefined;
+  this[kshutdownCallback] = undefined;
   this._pendingData = undefined;
   this._pendingEncoding = undefined; // for compatibility
   this._hadError = false;
@@ -2486,7 +2508,7 @@ Socket.prototype._final = function _final(callback) {
   if (!socket) return callback();
 
   // emit FIN allowHalfOpen only allow the readable side to close first
-  process.nextTick(endNT, socket, callback);
+  process.nextTick(endNT, socket, callback, this);
 };
 
 Object.defineProperty(Socket.prototype, "localAddress", {

src/runtime/socket/UpgradedDuplex.rs

@@ -44,6 +44,12 @@ pub(crate) struct UpgradedDuplex {
     pub on_end_callback: Cell<JSValue>,
     pub on_writable_callback: Cell<JSValue>,
     pub on_close_callback: Cell<JSValue>,
+    /// Passed as `cb` to every `duplex.write(chunk, cb)` / `duplex.end(null, cb)`.
+    pub on_write_done_callback: Cell<JSValue>,
+    /// Transport writes whose `cb` has not run yet. Zero means the Duplex took everything.
+    pub in_flight: Cell<u32>,
+    /// Nesting depth of `duplex.write()` / `duplex.end()` calls on the stack.
+    pub write_depth: Cell<u32>,
     pub event_loop_timer: JsCell<EventLoopTimer>,
     pub current_timeout: Cell<u32>,
     /// Transport bytes that arrived before the TLS engine existed.
@@ -274,26 +280,47 @@ impl UpgradedDuplex {
             }
         };
 
-        if let Some(data) = data {
-            let buffer = match bun_jsc::array_buffer::BinaryType::Buffer.to_js(data, &global) {
+        let payload = match data {
+            Some(data) => match bun_jsc::array_buffer::BinaryType::Buffer.to_js(data, &global) {
                 Ok(b) => b,
                 Err(err) => {
                     (self.handlers.on_error)(self.handlers.ctx, global.take_error(err));
                     return;
                 }
-            };
-            buffer.ensure_still_alive();
-
-            if let Err(err) = write_or_end.call(&global, duplex, &[buffer]) {
-                (self.handlers.on_error)(self.handlers.ctx, global.take_error(err));
-            }
-        } else {
-            if let Err(err) = write_or_end.call(&global, duplex, &[JSValue::NULL]) {
-                (self.handlers.on_error)(self.handlers.ctx, global.take_error(err));
-            }
+            },
+            None => JSValue::NULL,
+        };
+        payload.ensure_still_alive();
+
+        let done = self.write_done_handler(&global);
+        // Before the call: an erroring Duplex can run `done` inside it.
+        self.in_flight.set(self.in_flight.get() + 1);
+        self.write_depth.set(self.write_depth.get() + 1);
+        let result = write_or_end.call(&global, duplex, &[payload, done]);
+        self.write_depth.set(self.write_depth.get() - 1);
+        if let Err(err) = result {
+            // Fatal: the count stays high, so nothing completes before teardown resets it.
+            (self.handlers.on_error)(self.handlers.ctx, global.take_error(err));
         }
     }
 
+    fn write_done_handler(&self, global: &JSGlobalObject) -> JSValue {
+        lazy_js_handler(
+            &self.on_write_done_callback,
+            self.js_wrapper,
+            js_TLSSocket::duplex_on_write_done_set_cached,
+            global,
+            __jsc_host_on_write_done,
+            std::ptr::from_ref(self).cast_mut().cast::<c_void>(),
+        )
+    }
+
+    /// No `duplex.write()` callback is outstanding.
+    #[uws_callback(export = "UpgradedDuplex__transport_idle", no_catch)]
+    pub(crate) fn transport_idle(&self) -> bool {
+        self.in_flight.get() == 0
+    }
+
     fn internal_write(this: *mut Self, encoded_data: &[u8]) {
         // SAFETY: see handler note above.
         unsafe { &*this }.write_encrypted(encoded_data);
@@ -410,6 +437,9 @@ impl UpgradedDuplex {
             on_end_callback: Cell::new(JSValue::ZERO),
             on_writable_callback: Cell::new(JSValue::ZERO),
             on_close_callback: Cell::new(JSValue::ZERO),
+            on_write_done_callback: Cell::new(JSValue::ZERO),
+            in_flight: Cell::new(0),
+            write_depth: Cell::new(0),
             event_loop_timer: JsCell::new(EventLoopTimer::init_paused(
                 EventLoopTimerTag::UpgradedDuplex,
             )),
@@ -686,13 +716,15 @@ impl UpgradedDuplex {
             &self.on_end_callback,
             &self.on_writable_callback,
             &self.on_close_callback,
+            &self.on_write_done_callback,
         ] {
             let value = cb.get();
             if !value.is_empty() {
                 host_fn::set_function_data(value, None);
                 cb.set(JSValue::ZERO);
             }
         }
+        self.in_flight.set(0);
         self.ssl_error.set(CertError::default());
         self.pending_data.set(Vec::new());
         self.pending_close.set(false);
@@ -707,9 +739,9 @@ impl Drop for UpgradedDuplex {
     }
 }
 
-// SAFETY (all four host fns): the function data is the `*mut UpgradedDuplex`
-// installed by `get_js_handlers`; `teardown` clears it before the storage is
-// freed, so a non-null data pointer is live for the call.
+// SAFETY (all five host fns): the function data is the `*mut UpgradedDuplex`
+// installed by `get_js_handlers` and `write_done_handler`; `teardown` clears
+// it before the storage is freed, so a non-null data pointer is live for the call.
 
 #[bun_jsc::host_fn]
 fn on_received_data(global: &JSGlobalObject, frame: &CallFrame) -> JsResult<JSValue> {
@@ -781,6 +813,58 @@ fn on_writable(_global: &JSGlobalObject, frame: &CallFrame) -> JsResult<JSValue>
     Ok(JSValue::UNDEFINED)
 }
 
+/// `cb` of one `duplex.write()`; re-entered next tick with `err == js_wrapper` (error) or `err == origin` (drain).
+#[bun_jsc::host_fn]
+fn on_write_done(global: &JSGlobalObject, frame: &CallFrame) -> JsResult<JSValue> {
+    bun_output::scoped_log!(UpgradedDuplex, "onWriteDone");
+
+    let function = frame.callee();
+    let [err] = frame.arguments_as_array::<1>();
+
+    if let Some(self_ptr) = host_fn::get_function_data(function) {
+        // SAFETY: see host-fn note above.
+        let this = unsafe { &*self_ptr.cast::<UpgradedDuplex>() };
+        // A drain deferred from inside a write (see below).
+        if err == this.origin.get() && !err.is_empty() {
+            if this.in_flight.get() == 0 {
+                (this.handlers.on_writable)(this.handlers.ctx);
+            }
+            return Ok(JSValue::UNDEFINED);
+        }
+        if err == this.js_wrapper {
+            if !this.origin.get().is_empty() {
+                let epipe = bun_sys::Error::from_code_int(
+                    bun_sys::SystemErrno::EPIPE as core::ffi::c_int,
+                    bun_sys::Tag::write,
+                );
+                let err_value = <bun_sys::Error as bun_jsc::SysErrorJsc>::to_js(&epipe, global);
+                (this.handlers.on_error)(this.handlers.ctx, err_value);
+            }
+            return Ok(JSValue::UNDEFINED);
+        }
+        let Some(remaining) = this.in_flight.get().checked_sub(1) else {
+            debug_assert!(false, "a write callback without a counted write");
+            return Ok(JSValue::UNDEFINED);
+        };
+        this.in_flight.set(remaining);
+        if this.origin.get().is_empty() {
+            return Ok(JSValue::UNDEFINED);
+        }
+        if !err.is_empty_or_undefined_or_null() && !this.is_shutdown() {
+            JSValue::call_next_tick_1(function, global, this.js_wrapper)?;
+        } else if remaining == 0 {
+            if this.write_depth.get() > 0 {
+                // Inside the write that issued it: a drain now re-enters the flush too early.
+                JSValue::call_next_tick_1(function, global, this.origin.get())?;
+            } else {
+                (this.handlers.on_writable)(this.handlers.ctx);
+            }
+        }
+    }
+
+    Ok(JSValue::UNDEFINED)
+}
+
 #[bun_jsc::host_fn]
 fn on_close_js(_global: &JSGlobalObject, frame: &CallFrame) -> JsResult<JSValue> {
     bun_output::scoped_log!(UpgradedDuplex, "onCloseJS");

src/runtime/socket/socket_body.rs

@@ -950,13 +950,17 @@ impl<const SSL: bool> NewSocket<SSL> {
             return Ok(());
         }
         if this.native_callback.get().on_writable() {
-            return Ok(());
+            // The native drain ran JS that can close the socket; past a shutdown a parked end callback needs the JS drain.
+            if !this.has_handlers()
+                || this.socket.get().is_detached()
+                || !this.socket.get().is_shutdown()
+            {
+                return Ok(());
+            }
         }
         let handlers = this.get_handlers();
+        // The native flush (and the end-after-flush close) runs with or without a JS drain handler.
         let callback = handlers.on_writable();
-        if callback.is_empty() {
-            return Ok(());
-        }
 
         // Hold the socket alive for the rest of the dispatch: `internal_flush`
         // and the drain callback can both re-enter JS and close it.
@@ -1007,7 +1011,10 @@ impl<const SSL: bool> NewSocket<SSL> {
             this.buffered_data_for_node_net.get().len()
         );
         // is not writable if we have buffered data or if we are already detached
-        if this.buffered_data_for_node_net.get().len() > 0 || this.socket.get().is_detached() {
+        if this.buffered_data_for_node_net.get().len() > 0
+            || this.socket.get().is_detached()
+            || callback.is_empty()
+        {
             return Ok(());
         }
 
@@ -2702,7 +2709,9 @@ impl<const SSL: bool> NewSocket<SSL> {
                         // node:net fails the write like Node's onWriteComplete;
                         // -1 stays the legacy closed/shutdown sentinel.
                         JSValue::js_number(f64::from(wrote))
-                    } else if usize::try_from(wrote.max(0)).expect("int cast") == total {
+                    } else if usize::try_from(wrote.max(0)).expect("int cast") == total
+                        && this.flushed_to_transport()
+                    {
                         JSValue::TRUE
                     } else {
                         JSValue::FALSE
@@ -2712,6 +2721,12 @@ impl<const SSL: bool> NewSocket<SSL> {
         )
     }
 
+    /// `false` parks the node:net callback until the transport drains (a BYPASS_TLS twin gets none).
+    #[inline]
+    fn flushed_to_transport(&self) -> bool {
+        !SSL || self.flags.get().contains(Flags::BYPASS_TLS) || self.socket.get().transport_idle()
+    }
+
     #[bun_jsc::host_fn(method)]
     pub(crate) fn end_buffered(
         this: &Self,
@@ -2734,7 +2749,10 @@ impl<const SSL: bool> NewSocket<SSL> {
                     let _ = this.internal_flush();
                 }
 
-                JSValue::from(usize::try_from(wrote.max(0)).expect("int cast") == total)
+                JSValue::from(
+                    usize::try_from(wrote.max(0)).expect("int cast") == total
+                        && this.flushed_to_transport(),
+                )
             }
         };
         Ok(result)
@@ -3125,6 +3143,7 @@ impl<const SSL: bool> NewSocket<SSL> {
             && flags.contains(Flags::END_AFTER_FLUSH)
             && !flags.contains(Flags::EMPTY_PACKET_PENDING)
             && self.buffered_data_for_node_net.get().len() == 0
+            && self.flushed_to_transport()
     }
 
     /// Flushes the node:net buffered tail. Returns 0, or the positive errno of
@@ -3246,7 +3265,8 @@ impl<const SSL: bool> NewSocket<SSL> {
             this.socket.get().shutdown();
         }
 
-        Ok(JSValue::UNDEFINED)
+        // `false`: a drain completes it, like a `$write` that returns false.
+        Ok(JSValue::from(this.flushed_to_transport()))
     }
 
     #[bun_jsc::host_fn(method)]