fix(nif): harden callback suspend/resume lifetime

Keep the worker resource alive while a callback is suspended (it could
be GC'd mid-suspension -> use-after-free on resume); release it in the
suspended-state destructor and NULL-check before the replay deref. Free
any prior result_data before a resume stores a new one (result_data, not
the toggling has_result flag, is the real pending-result indicator, so a
duplicate/raced resume no longer leaks). Clear the pending-callback TLS at
the worker request boundary so a reused worker thread doesn't trip the
stale-TLS invariant. Run the callback-response pipe writes on dirty IO
schedulers with non-blocking, deadline-bounded writes so a stalled reader
or large payload can't wedge a scheduler or desync the framed protocol.
Also Py_XDECREF callback_args in the GIL-held destructor branch.

Adds py_reentrant_SUITE:test_reentrant_resume_stress.

Author: benoitc

CHANGELOG.md

@@ -4,6 +4,14 @@
 
 ### Security
 
+- **Callback suspend/resume lifetime hardening** - The worker resource is now kept
+  alive for the lifetime of a suspended callback (it could previously be GC'd mid-
+  suspension, causing a use-after-free on resume). A resume frees any prior result
+  before storing a new one (no leak/double-replay on a duplicate resume), the
+  pending-callback thread-local is cleared at the worker request boundary, and the
+  callback-response pipe writes run on dirty schedulers with non-blocking, deadline-
+  bounded writes so a stalled reader or large payload can't wedge a scheduler or
+  desync the framed protocol.
 - **Zero-copy buffer pinning** - `py_buffer` no longer relocates (and frees) its
   storage while a Python `memoryview` points into it. A write that would grow the
   buffer while a view is held now returns an error instead of dangling the view into

c_src/py_callback.c

@@ -399,6 +399,14 @@ static suspended_state_t *create_suspended_state_ex(
     } else {
         state->worker = source->data.existing->worker;
     }
+    /* Keep the worker resource alive for as long as the suspended state exists.
+     * Without this the worker can be GC'd while a callback is suspended, and
+     * nif_resume_callback_dirty would dereference a freed worker (use-after-free
+     * with the GIL held). Mirrors the enif_keep_resource(ctx) on the context path;
+     * suspended_state_destructor releases it. */
+    if (state->worker != NULL) {
+        enif_keep_resource(state->worker);
+    }
 
     state->callback_id = PyLong_AsUnsignedLongLong(callback_id_obj);
 
@@ -4316,7 +4324,14 @@ static ERL_NIF_TERM nif_resume_callback(ErlNifEnv *env, int argc, const ERL_NIF_
     /* Store the result in the suspended state */
     pthread_mutex_lock(&state->mutex);
 
-    /* Copy result data */
+    /* Copy result data. Free any prior result first: a duplicate/raced resume
+     * would otherwise leak the previous buffer. (has_result is not a one-shot
+     * flag -- it toggles during nested replay -- so result_data is the real
+     * pending-result indicator.) */
+    if (state->result_data != NULL) {
+        enif_free(state->result_data);
+        state->result_data = NULL;
+    }
     state->result_data = enif_alloc(result_bin.size);
     if (state->result_data == NULL) {
         pthread_mutex_unlock(&state->mutex);
@@ -4364,6 +4379,12 @@ static ERL_NIF_TERM nif_resume_callback_dirty(ErlNifEnv *env, int argc, const ER
         return make_error(env, "no_result");
     }
 
+    /* The worker is kept alive for the lifetime of the suspended state, but
+     * guard rather than dereference NULL in the replay below. */
+    if (state->worker == NULL) {
+        return make_error(env, "no_worker");
+    }
+
     /* Set up thread-local state for replay */
     tl_current_worker = state->worker;
     tl_callback_env = env;

c_src/py_exec.c

@@ -304,11 +304,14 @@ static void process_request(py_request_t *req) {
                     suspended_state_t *suspended = create_suspended_state(env, exc_args, req);
                     Py_DECREF(exc_args);
                     if (suspended == NULL) {
-                        tl_pending_callback = false;
-                        Py_CLEAR(tl_pending_args);
+                        clear_pending_callback_tls();
                         req->result = make_error(env, "create_suspended_state_failed");
                     } else {
                         req->result = build_suspended_result(env, suspended);
+                        /* func_name/args are copied into the suspended state; clear
+                         * the pending-callback TLS so a later request on this reused
+                         * worker thread doesn't trip the stale-TLS entry invariant. */
+                        clear_pending_callback_tls();
                     }
                 }
             } else {
@@ -408,11 +411,11 @@ static void process_request(py_request_t *req) {
                         suspended_state_t *suspended = create_suspended_state(env, exc_args, req);
                         Py_DECREF(exc_args);
                         if (suspended == NULL) {
-                            tl_pending_callback = false;
-                            Py_CLEAR(tl_pending_args);
+                            clear_pending_callback_tls();
                             req->result = make_error(env, "create_suspended_state_failed");
                         } else {
                             req->result = build_suspended_result(env, suspended);
+                            clear_pending_callback_tls();
                         }
                     }
                 } else {

c_src/py_nif.c

@@ -513,11 +513,18 @@ static void suspended_state_destructor(ErlNifEnv *env, void *obj) {
     (void)env;
     suspended_state_t *state = (suspended_state_t *)obj;
 
+    /* Release the worker resource kept alive in create_suspended_state_ex. */
+    if (state->worker != NULL) {
+        enif_release_resource(state->worker);
+        state->worker = NULL;
+    }
+
     /* Clean up Python objects if Python is still initialized.
      * suspended_state_t is used with the worker-based API which runs in
      * the main interpreter, so we always use PyGILState_Ensure. */
     if (runtime_is_running() && state->callback_args != NULL) {
         if (PyGILState_GetThisThreadState() != NULL || PyGILState_Check()) {
+            Py_XDECREF(state->callback_args);
             state->callback_args = NULL;
         } else {
             PyGILState_STATE gstate = PyGILState_Ensure();
@@ -1700,6 +1707,11 @@ static ERL_NIF_TERM nif_set_callback_handler(ErlNifEnv *env, int argc, const ERL
     if (pipe(worker->callback_pipe) < 0) {
         return make_error(env, "pipe_failed");
     }
+    /* Non-blocking write end so write_all_with_deadline can bound the write. */
+    {
+        int wfl = fcntl(worker->callback_pipe[1], F_GETFL, 0);
+        if (wfl >= 0) (void)fcntl(worker->callback_pipe[1], F_SETFL, wfl | O_NONBLOCK);
+    }
 
     worker->has_callback_handler = true;
 
@@ -1708,6 +1720,10 @@ static ERL_NIF_TERM nif_set_callback_handler(ErlNifEnv *env, int argc, const ERL
         enif_make_int(env, worker->callback_pipe[1]));
 }
 
+/* Bound for callback-response pipe writes: a stalled reader must not block a
+ * dirty scheduler forever (the pipe write ends are set non-blocking). */
+#define CALLBACK_RESPONSE_IO_TIMEOUT_MS 30000
+
 static ERL_NIF_TERM nif_send_callback_response(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
     (void)argc;
     int fd;
@@ -1721,15 +1737,16 @@ static ERL_NIF_TERM nif_send_callback_response(ErlNifEnv *env, int argc, const E
         return make_error(env, "invalid_response");
     }
 
-    /* Write length then data */
+    /* Write length then data with a timed, non-blocking writer (the pipe write
+     * end is O_NONBLOCK) so a stalled reader or a large payload can't block a
+     * dirty scheduler forever or desync the length-framed protocol on EINTR. */
     uint32_t len = (uint32_t)response.size;
-    ssize_t n = write(fd, &len, sizeof(len));
-    if (n != sizeof(len)) {
+    if (write_all_with_deadline(fd, &len, sizeof(len),
+                                CALLBACK_RESPONSE_IO_TIMEOUT_MS) != WRITE_OK) {
         return make_error(env, "write_length_failed");
     }
-
-    n = write(fd, response.data, response.size);
-    if (n != (ssize_t)response.size) {
+    if (write_all_with_deadline(fd, response.data, response.size,
+                                CALLBACK_RESPONSE_IO_TIMEOUT_MS) != WRITE_OK) {
         return make_error(env, "write_data_failed");
     }
 
@@ -4702,6 +4719,11 @@ static ERL_NIF_TERM nif_context_create(ErlNifEnv *env, int argc, const ERL_NIF_T
         enif_release_resource(ctx);
         return make_error(env, "pipe_create_failed");
     }
+    /* Non-blocking write end so write_all_with_deadline can bound the write. */
+    {
+        int wfl = fcntl(ctx->callback_pipe[1], F_GETFL, 0);
+        if (wfl >= 0) (void)fcntl(ctx->callback_pipe[1], F_SETFL, wfl | O_NONBLOCK);
+    }
 
 #ifdef HAVE_SUBINTERPRETERS
     ctx->uses_own_gil = false;
@@ -6579,15 +6601,17 @@ static ERL_NIF_TERM nif_context_write_callback_response(ErlNifEnv *env, int argc
         return make_error(env, "pipe_not_initialized");
     }
 
-    /* Write length prefix (4 bytes, native endianness - must match read_length_prefixed_data) */
+    /* Write length prefix + data with a timed, non-blocking writer (the pipe
+     * write end is O_NONBLOCK) so a stalled reader or large payload can't block a
+     * dirty scheduler forever or desync the framed protocol. 4-byte native-endian
+     * length must match read_length_prefixed_data. */
     uint32_t len = (uint32_t)data.size;
-    ssize_t written = write(ctx->callback_pipe[1], &len, sizeof(len));
-    if (written != sizeof(len)) {
+    if (write_all_with_deadline(ctx->callback_pipe[1], &len, sizeof(len),
+                                CALLBACK_RESPONSE_IO_TIMEOUT_MS) != WRITE_OK) {
         return make_error(env, "write_failed");
     }
-
-    written = write(ctx->callback_pipe[1], data.data, data.size);
-    if (written != (ssize_t)data.size) {
+    if (write_all_with_deadline(ctx->callback_pipe[1], data.data, data.size,
+                                CALLBACK_RESPONSE_IO_TIMEOUT_MS) != WRITE_OK) {
         return make_error(env, "write_failed");
     }
 
@@ -6639,7 +6663,13 @@ static ERL_NIF_TERM nif_context_resume(ErlNifEnv *env, int argc, const ERL_NIF_T
         return make_error(env, "context_mismatch");
     }
 
-    /* Store the callback result */
+    /* Store the callback result. Free any prior result first to avoid leaking it
+     * on a duplicate/raced resume (result_data, not the toggling has_result flag,
+     * is the real pending-result indicator). */
+    if (state->result_data != NULL) {
+        enif_free(state->result_data);
+        state->result_data = NULL;
+    }
     state->result_data = enif_alloc(result_bin.size);
     if (state->result_data == NULL) {
         return make_error(env, "alloc_failed");
@@ -7812,7 +7842,7 @@ static ErlNifFunc nif_funcs[] = {
 
     /* Callback support */
     {"set_callback_handler", 2, nif_set_callback_handler, 0},
-    {"send_callback_response", 2, nif_send_callback_response, 0},
+    {"send_callback_response", 2, nif_send_callback_response, ERL_NIF_DIRTY_JOB_IO_BOUND},
     {"resume_callback", 2, nif_resume_callback, 0},
 
     /* Async worker management */
@@ -7961,7 +7991,7 @@ static ErlNifFunc nif_funcs[] = {
     {"context_interp_id", 1, nif_context_interp_id, 0},
     {"context_set_callback_handler", 2, nif_context_set_callback_handler, 0},
     {"context_get_callback_pipe", 1, nif_context_get_callback_pipe, 0},
-    {"context_write_callback_response", 2, nif_context_write_callback_response, 0},
+    {"context_write_callback_response", 2, nif_context_write_callback_response, ERL_NIF_DIRTY_JOB_IO_BOUND},
     {"context_resume", 3, nif_context_resume, ERL_NIF_DIRTY_JOB_CPU_BOUND},
     {"context_cancel_resume", 2, nif_context_cancel_resume, 0},
     {"context_get_event_loop", 1, nif_context_get_event_loop, 0},

test/py_reentrant_SUITE.erl

@@ -25,7 +25,8 @@
     test_callback_with_try_except/1,
     test_async_call/1,
     test_callback_name_registry/1,
-    test_etf_decode_safe/1
+    test_etf_decode_safe/1,
+    test_reentrant_resume_stress/1
 ]).
 
 all() ->
@@ -40,7 +41,8 @@ all() ->
         test_callback_with_try_except,
         test_async_call,
         test_callback_name_registry,
-        test_etf_decode_safe
+        test_etf_decode_safe,
+        test_reentrant_resume_stress
     ].
 
 init_per_suite(Config) ->
@@ -69,6 +71,7 @@ end_per_testcase(_TestCase, _Config) ->
     try py:unregister_function(test_registry_func) catch _:_ -> ok end,
     try py:unregister_function(etf_probe_ok) catch _:_ -> ok end,
     try py:unregister_function(etf_probe_novel) catch _:_ -> ok end,
+    try py:unregister_function(rs_double) catch _:_ -> ok end,
     ok.
 
 %%% ============================================================================
@@ -93,6 +96,25 @@ test_basic_reentrant(_Config) ->
 
     ok.
 
+%% @doc Many reentrant suspend/resume cycles must all succeed with no leak,
+%% stale-TLS invariant trip, or crash. Regression for the suspend/resume lifetime
+%% fixes (worker keep, TLS clear, result_data free, hardened callback-pipe writes);
+%% exercises the callback pipe and context resume under repetition.
+test_reentrant_resume_stress(_Config) ->
+    py:register_function(rs_double, fun([X]) ->
+        {ok, R} = py:eval(iolist_to_binary(io_lib:format("~p * 2", [X]))),
+        R
+    end),
+    Got = [begin
+        Code = iolist_to_binary(
+            io_lib:format("__import__('erlang').call('rs_double', ~p) + 1", [N])),
+        {ok, V} = py:eval(Code),
+        V
+    end || N <- lists:seq(1, 100)],
+    Got = [N * 2 + 1 || N <- lists:seq(1, 100)],
+    py:unregister_function(rs_double),
+    ok.
+
 %% @doc Regression for the binary_to_term SAFE-flag hardening (atom exhaustion).
 %% A `__etf__:` callback result that encodes a brand-new atom must be rejected (not
 %% decoded) so it cannot mint non-GC'd atoms, while a valid existing-atom payload