usockets (Windows): keep a closed socket allocated until the dispatch that closed it returns

packages/bun-usockets/src/eventing/libuv.c

@@ -52,8 +52,7 @@ static int us_internal_poll_cb_socket_is_probeable(struct us_poll_t *wp) {
   return !s->flags.is_closed && s->flags.is_paused;
 }
 
-/* uv_poll_t->data always (except for most times after calling us_poll_stop)
- * points to the us_poll_t */
+/* uv_poll_t->data always points to the us_poll_t */
 static void poll_cb(uv_poll_t *p, int status, int events) {
   /* UV_DISCONNECT (Windows AFD): the peer closed its write side. A FIN
    * arriving after this side already half-closed and stopped reading never
@@ -176,18 +175,13 @@ static void check_cb(uv_check_t *p) {
   us_internal_loop_post(loop);
 }
 
-/* Not used for polls, since polls need two frees */
+/* Not used for polls: their uv handle is a separate allocation */
 static void close_cb_free(uv_handle_t *h) { us_free(h->data); }
 
-/* This one is different for polls, since we need two frees here */
-static void close_cb_free_poll(uv_handle_t *h) {
-  /* It is only in case we called us_poll_stop then quickly us_poll_free that we
-   * enter this. Most of the time, actual freeing is done by us_poll_free. */
-  if (h->data) {
-    us_free(h->data);
-    us_free(h);
-  }
-}
+/* A poll's uv_poll_t is a separate allocation that libuv keeps linked into
+ * the loop until its close completes; us_poll_free hands it over with
+ * uv_close and this frees it. The us_poll_t itself is freed by us_poll_free. */
+static void close_cb_free_uv_poll(uv_handle_t *h) { us_free(h); }
 
 static void timer_cb(uv_timer_t *t) {
   struct us_internal_callback_t *cb = t->data;
@@ -208,23 +202,19 @@ void us_poll_init(struct us_poll_t *p, LIBUS_SOCKET_DESCRIPTOR fd,
 }
 
 void us_poll_free(struct us_poll_t *p, struct us_loop_t *loop) {
-  // poll was resized and dont own uv_poll_t anymore
-  if(!p->uv_p) {
-    us_free(p);
-    return;
-  }
-  /* The idea here is like so; in us_poll_stop we call uv_close after setting
-   * data of uv-poll to 0. This means that in close_cb_free we call free on 0
-   * with does nothing, since us_poll_stop should not really free the poll.
-   * HOWEVER, if we then call us_poll_free while still closing the uv-poll, we
-   * simply change back the data to point to our structure so that we actually
-   * do free it like we should. */
-  if (uv_is_closing((uv_handle_t *)p->uv_p)) {
-    p->uv_p->data = p;
-  } else {
-    us_free(p->uv_p);
-    us_free(p);
+  /* NULL when a resize moved the handle to the new block. */
+  if (p->uv_p) {
+    if (p->uv_p->type == UV_POLL) {
+      /* us_poll_start_rc initialised it (uv__handle_init linked it into the
+       * loop): only libuv can unlink it, and it is freed once that completes.
+       * This is the one place a poll handle is uv_close'd - see us_poll_stop. */
+      uv_close((uv_handle_t *)p->uv_p, close_cb_free_uv_poll);
+    } else {
+      /* Never initialised: libuv has not seen this block. */
+      us_free(p->uv_p);
+    }
   }
+  us_free(p);
 }
 
 int us_poll_start_rc(struct us_poll_t *p, struct us_loop_t *loop, int events) {
@@ -236,9 +226,9 @@ int us_poll_start_rc(struct us_poll_t *p, struct us_loop_t *loop, int events) {
   /* uv_poll_init_socket (win/poll.c) can fail either before uv__handle_init
    * (ioctlsocket FIONBIO) or after it (getsockopt SO_PROTOCOL_INFOW). The
    * latter leaves the handle linked into loop->handle_queue with
-   * submitted_events_* still unset. Zero first so, on failure, ->type
-   * distinguishes the two states and the fields uv__poll_close reads are 0
-   * rather than garbage. */
+   * submitted_events_* still unset. Zero first so, on failure, ->type tells
+   * us_poll_free which of the two states it is in and the fields
+   * uv__poll_close reads are 0 rather than garbage. */
   memset(p->uv_p, 0, sizeof(uv_poll_t));
   p->uv_p->data = p;
 
@@ -252,20 +242,8 @@ int us_poll_start_rc(struct us_poll_t *p, struct us_loop_t *loop, int events) {
 #endif
   rc = uv_poll_init_socket(loop->uv_loop, p->uv_p, p->fd);
   if (rc < 0) {
+    /* The caller's us_poll_free disposes of uv_p either way (see there). */
     int saved = LIBUS_ERR;
-    if (p->uv_p->type == UV_POLL) {
-      /* uv__handle_init ran: the handle is in loop->handle_queue. Close it
-       * through libuv so it is unlinked; the caller's us_poll_free sees
-       * uv_is_closing and hands ownership to close_cb_free_poll. */
-      p->uv_p->data = 0;
-      uv_close((uv_handle_t *)p->uv_p, close_cb_free_poll);
-    } else {
-      /* Never reached uv__handle_init: uv_p is still our raw block. Free it
-       * here and null the pointer so the caller's us_poll_free takes the
-       * !uv_p fast path (its uv_is_closing check would read garbage). */
-      us_free(p->uv_p);
-      p->uv_p = NULL;
-    }
     errno = saved ? saved : -rc;
     return rc;
   }
@@ -301,14 +279,17 @@ int us_poll_change(struct us_poll_t *p, struct us_loop_t *loop, int events) {
 
 void us_poll_stop(struct us_poll_t *p, struct us_loop_t *loop) {
   if(!p->uv_p) return;
+  /* Stop only; the uv_close waits for us_poll_free. A socket is routinely
+   * closed from inside its own poll_cb, and JS run from there may drive the
+   * loop again (waitForPromise) before poll_cb returns. Starting the close
+   * here would let that nested uv_run complete it and free the uv_poll_t
+   * while libuv's uv__fast_poll_process_poll_req for this very handle is still
+   * on the stack: once poll_cb returns it reads the freed handle and queues
+   * its endgame a second time (double close callback, corrupted handle
+   * queue). us_poll_free runs from the outermost us_internal_loop_post, which
+   * is never inside a poll_cb. uv_poll_stop alone already guarantees poll_cb
+   * is not invoked for this handle again. */
   uv_poll_stop(p->uv_p);
-
-  /* We normally only want to close the poll here, not free it. But if we stop
-   * it, then quickly "free" it with us_poll_free, we postpone the actual
-   * freeing to close_cb_free_poll whenever it triggers. That's why we set data
-   * to null here, so that us_poll_free can reset it if needed */
-  p->uv_p->data = 0;
-  uv_close((uv_handle_t *)p->uv_p, close_cb_free_poll);
 }
 
 int us_poll_events(struct us_poll_t *p) {
@@ -335,7 +316,9 @@ void us_loop_pump(struct us_loop_t *loop) {
    * bun:test) supply their own keep-going predicate, so force exactly one
    * non-blocking iteration; UV_RUN_NOWAIT keeps the poll timeout at 0. */
   loop->uv_loop->active_handles++;
+  loop->data.tick_depth++;
   uv_run(loop->uv_loop, UV_RUN_NOWAIT);
+  loop->data.tick_depth--;
   loop->uv_loop->active_handles--;
 }
 
@@ -414,14 +397,22 @@ void us_loop_run(struct us_loop_t *loop) {
     Bun__JSC_onBeforeWait(loop->data.jsc_vm, (uint64_t) uv_now(loop->uv_loop) * 1000000ULL);
   }
 
+  /* Same bracket as us_loop_run_bun_tick on epoll/kqueue: a socket callback
+   * dispatched inside this uv_run may drive the loop again (waitForPromise ->
+   * us_loop_pump), and us_internal_loop_post only frees closed sockets from
+   * the outermost tick so an outer dispatch never resumes on a freed socket. */
+  loop->data.tick_depth++;
   uv_run(loop->uv_loop, UV_RUN_ONCE);
+  loop->data.tick_depth--;
 }
 
 struct us_poll_t *us_create_poll(struct us_loop_t *loop, int fallthrough,
                                  unsigned int ext_size) {
   struct us_poll_t *p =
       (struct us_poll_t *)us_malloc(sizeof(struct us_poll_t) + ext_size);
   p->uv_p = us_malloc(sizeof(uv_poll_t));
+  /* Not a libuv handle until us_poll_start_rc; us_poll_free checks this. */
+  p->uv_p->type = UV_UNKNOWN_HANDLE;
   p->uv_p->data = p;
   return p;
 }

packages/bun-usockets/src/internal/loop_data.h

@@ -95,10 +95,11 @@ struct us_internal_loop_data_t {
     /* We do not care if this flips or not, it doesn't matter */
     size_t iteration_nr;
     void* jsc_vm;
-    /* Reentrancy depth of us_loop_run_bun_tick. When >1, we are inside a
-     * nested tick (e.g. waitForPromise from a poll callback). Freeing closed
-     * sockets must be deferred to the outermost tick so the outer dispatch
-     * doesn't read a freed poll. */
+    /* Reentrancy depth of the loop tick (us_loop_run / us_loop_run_bun_tick /
+     * us_loop_pump). When >1, we are inside a nested tick (e.g.
+     * waitForPromise from a poll callback). Freeing closed sockets must be
+     * deferred to the outermost tick so the outer dispatch doesn't read a
+     * freed poll. */
     int tick_depth;
 };
 

test/js/bun/net/socket.test.ts

@@ -4354,3 +4354,25 @@ describe.concurrent("a socket closed by data() while its peer's reset is being d
     expect(exitCode).toBe(0);
   });
 });
+
+describe.concurrent("a socket closed by data() which then re-enters the event loop before returning", () => {
+  // The fixture runs under `bun test` so that expect(promise).resolves can drive
+  // nested event-loop ticks from inside the data callback. The closed socket must
+  // stay allocated until the dispatch that invoked data() has returned; the loop
+  // used to free it from a nested tick on Windows, and the outer dispatch then
+  // read (and the allocator reused) freed memory.
+  it("is not freed until the dispatch that called data() has returned", async () => {
+    await using proc = Bun.spawn({
+      cmd: [bunExe(), "test", fileURLToPath(new URL("./close-inside-data-reentrant-fixture.ts", import.meta.url))],
+      env: bunEnv,
+      stdout: "pipe",
+      stderr: "pipe",
+    });
+    const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]);
+    // stdout carries only the runner's version banner; results go to stderr.
+    expect(stdout).toMatch(/^bun test v\S+ \(\S+\)\n$/);
+    expect(stderr).toContain(" 1 pass");
+    expect(proc.signalCode).toBeNull();
+    expect(exitCode).toBe(0);
+  });
+});