Skip to content

Commit a325348

Browse files
benoitcbenoitc
committed
Fix memory safety issues in C NIF code
- Fix memory leaks in nif_async_worker_new error paths (msg_env not freed) - Fix memory leak in multi_executor_stop (shutdown_req not freed) - Fix memory leaks in create_suspended_state when binary allocations fail - Add NULL checks after enif_alloc_resource calls to prevent crashes - Add NULL checks after enif_alloc_env calls in async_future_callback and erlang_call_impl - Properly clean up Python objects (Py_DECREF) on allocation failures
1 parent 7a4cc6a commit a325348

1 file changed

Lines changed: 95 additions & 23 deletions

File tree

‎c_src/py_nif.c‎

Lines changed: 95 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1142,6 +1142,10 @@ static ERL_NIF_TERM nif_tracemalloc_stop(ErlNifEnv *env, int argc, const ERL_NIF
11421142
*/
11431143
static void async_future_callback(py_async_worker_t *worker, async_pending_t *pending) {
11441144
ErlNifEnv *msg_env = enif_alloc_env();
1145+
if (msg_env == NULL) {
1146+
/* Cannot send result - just log and return */
1147+
return;
1148+
}
11451149
PyObject *py_result = PyObject_CallMethod(pending->future, "result", NULL);
11461150

11471151
ERL_NIF_TERM result_term;
@@ -1311,6 +1315,7 @@ static ERL_NIF_TERM nif_async_worker_new(ErlNifEnv *env, int argc, const ERL_NIF
13111315

13121316
/* Create notification pipe */
13131317
if (pipe(worker->notify_pipe) < 0) {
1318+
enif_free_env(worker->msg_env);
13141319
enif_release_resource(worker);
13151320
return make_error(env, "pipe_failed");
13161321
}
@@ -1323,6 +1328,7 @@ static ERL_NIF_TERM nif_async_worker_new(ErlNifEnv *env, int argc, const ERL_NIF
13231328
close(worker->notify_pipe[0]);
13241329
close(worker->notify_pipe[1]);
13251330
pthread_mutex_destroy(&worker->queue_mutex);
1331+
enif_free_env(worker->msg_env);
13261332
enif_release_resource(worker);
13271333
return make_error(env, "thread_create_failed");
13281334
}
@@ -2507,26 +2513,56 @@ static suspended_state_t *create_suspended_state(ErlNifEnv *env, PyObject *exc_a
25072513
state->orig_env = enif_alloc_env();
25082514
if (state->orig_env == NULL) {
25092515
Py_DECREF(callback_args);
2516+
state->callback_args = NULL;
25102517
enif_free(state->callback_func_name);
2518+
state->callback_func_name = NULL;
25112519
enif_release_resource(state);
25122520
return NULL;
25132521
}
25142522

25152523
/* Copy request-specific data */
25162524
if (req->type == PY_REQ_CALL) {
25172525
/* Copy module and function binaries */
2518-
enif_alloc_binary(req->module_bin.size, &state->orig_module);
2526+
if (!enif_alloc_binary(req->module_bin.size, &state->orig_module)) {
2527+
Py_DECREF(callback_args);
2528+
state->callback_args = NULL;
2529+
enif_free(state->callback_func_name);
2530+
state->callback_func_name = NULL;
2531+
enif_free_env(state->orig_env);
2532+
state->orig_env = NULL;
2533+
enif_release_resource(state);
2534+
return NULL;
2535+
}
25192536
memcpy(state->orig_module.data, req->module_bin.data, req->module_bin.size);
25202537

2521-
enif_alloc_binary(req->func_bin.size, &state->orig_func);
2538+
if (!enif_alloc_binary(req->func_bin.size, &state->orig_func)) {
2539+
enif_release_binary(&state->orig_module);
2540+
Py_DECREF(callback_args);
2541+
state->callback_args = NULL;
2542+
enif_free(state->callback_func_name);
2543+
state->callback_func_name = NULL;
2544+
enif_free_env(state->orig_env);
2545+
state->orig_env = NULL;
2546+
enif_release_resource(state);
2547+
return NULL;
2548+
}
25222549
memcpy(state->orig_func.data, req->func_bin.data, req->func_bin.size);
25232550

25242551
/* Copy args and kwargs to our environment */
25252552
state->orig_args = enif_make_copy(state->orig_env, req->args_term);
25262553
state->orig_kwargs = enif_make_copy(state->orig_env, req->kwargs_term);
25272554
} else if (req->type == PY_REQ_EVAL) {
25282555
/* Copy code binary */
2529-
enif_alloc_binary(req->code_bin.size, &state->orig_code);
2556+
if (!enif_alloc_binary(req->code_bin.size, &state->orig_code)) {
2557+
Py_DECREF(callback_args);
2558+
state->callback_args = NULL;
2559+
enif_free(state->callback_func_name);
2560+
state->callback_func_name = NULL;
2561+
enif_free_env(state->orig_env);
2562+
state->orig_env = NULL;
2563+
enif_release_resource(state);
2564+
return NULL;
2565+
}
25302566
memcpy(state->orig_code.data, req->code_bin.data, req->code_bin.size);
25312567

25322568
/* Copy locals */
@@ -2679,6 +2715,11 @@ static PyObject *erlang_call_impl(PyObject *self, PyObject *args) {
26792715
if (!tl_allow_suspension) {
26802716
/* Fall back to blocking behavior - send message and wait on pipe */
26812717
ErlNifEnv *msg_env = enif_alloc_env();
2718+
if (msg_env == NULL) {
2719+
Py_DECREF(call_args);
2720+
PyErr_SetString(PyExc_MemoryError, "Failed to allocate message environment");
2721+
return NULL;
2722+
}
26822723
ERL_NIF_TERM func_term;
26832724
{
26842725
unsigned char *buf = enif_make_new_binary(msg_env, func_name_len, &func_term);
@@ -3367,11 +3408,16 @@ static void process_request(py_request_t *req) {
33673408
}
33683409
} else if (PyGen_Check(py_result) || PyIter_Check(py_result)) {
33693410
py_object_t *wrapper = enif_alloc_resource(PYOBJ_RESOURCE_TYPE, sizeof(py_object_t));
3370-
wrapper->obj = py_result;
3371-
ERL_NIF_TERM gen_ref = enif_make_resource(env, wrapper);
3372-
enif_release_resource(wrapper);
3373-
req->result = enif_make_tuple2(env, ATOM_OK,
3374-
enif_make_tuple2(env, ATOM_GENERATOR, gen_ref));
3411+
if (wrapper == NULL) {
3412+
Py_DECREF(py_result);
3413+
req->result = make_error(env, "alloc_failed");
3414+
} else {
3415+
wrapper->obj = py_result;
3416+
ERL_NIF_TERM gen_ref = enif_make_resource(env, wrapper);
3417+
enif_release_resource(wrapper);
3418+
req->result = enif_make_tuple2(env, ATOM_OK,
3419+
enif_make_tuple2(env, ATOM_GENERATOR, gen_ref));
3420+
}
33753421
} else {
33763422
ERL_NIF_TERM term_result = py_to_term(env, py_result);
33773423
Py_DECREF(py_result);
@@ -3469,11 +3515,16 @@ static void process_request(py_request_t *req) {
34693515
}
34703516
} else if (PyGen_Check(py_result) || PyIter_Check(py_result)) {
34713517
py_object_t *wrapper = enif_alloc_resource(PYOBJ_RESOURCE_TYPE, sizeof(py_object_t));
3472-
wrapper->obj = py_result;
3473-
ERL_NIF_TERM gen_ref = enif_make_resource(env, wrapper);
3474-
enif_release_resource(wrapper);
3475-
req->result = enif_make_tuple2(env, ATOM_OK,
3476-
enif_make_tuple2(env, ATOM_GENERATOR, gen_ref));
3518+
if (wrapper == NULL) {
3519+
Py_DECREF(py_result);
3520+
req->result = make_error(env, "alloc_failed");
3521+
} else {
3522+
wrapper->obj = py_result;
3523+
ERL_NIF_TERM gen_ref = enif_make_resource(env, wrapper);
3524+
enif_release_resource(wrapper);
3525+
req->result = enif_make_tuple2(env, ATOM_OK,
3526+
enif_make_tuple2(env, ATOM_GENERATOR, gen_ref));
3527+
}
34773528
} else {
34783529
ERL_NIF_TERM term_result = py_to_term(env, py_result);
34793530
Py_DECREF(py_result);
@@ -3537,11 +3588,16 @@ static void process_request(py_request_t *req) {
35373588
}
35383589
} else if (PyGen_Check(item) || PyIter_Check(item)) {
35393590
py_object_t *wrapper = enif_alloc_resource(PYOBJ_RESOURCE_TYPE, sizeof(py_object_t));
3540-
wrapper->obj = item;
3541-
ERL_NIF_TERM gen_ref = enif_make_resource(env, wrapper);
3542-
enif_release_resource(wrapper);
3543-
req->result = enif_make_tuple2(env, ATOM_OK,
3544-
enif_make_tuple2(env, ATOM_GENERATOR, gen_ref));
3591+
if (wrapper == NULL) {
3592+
Py_DECREF(item);
3593+
req->result = make_error(env, "alloc_failed");
3594+
} else {
3595+
wrapper->obj = item;
3596+
ERL_NIF_TERM gen_ref = enif_make_resource(env, wrapper);
3597+
enif_release_resource(wrapper);
3598+
req->result = enif_make_tuple2(env, ATOM_OK,
3599+
enif_make_tuple2(env, ATOM_GENERATOR, gen_ref));
3600+
}
35453601
} else {
35463602
ERL_NIF_TERM term_result = py_to_term(env, item);
35473603
Py_DECREF(item);
@@ -3564,10 +3620,15 @@ static void process_request(py_request_t *req) {
35643620
req->result = make_py_error(env);
35653621
} else {
35663622
py_object_t *wrapper = enif_alloc_resource(PYOBJ_RESOURCE_TYPE, sizeof(py_object_t));
3567-
wrapper->obj = module;
3568-
ERL_NIF_TERM mod_ref = enif_make_resource(env, wrapper);
3569-
enif_release_resource(wrapper);
3570-
req->result = enif_make_tuple2(env, ATOM_OK, mod_ref);
3623+
if (wrapper == NULL) {
3624+
Py_DECREF(module);
3625+
req->result = make_error(env, "alloc_failed");
3626+
} else {
3627+
wrapper->obj = module;
3628+
ERL_NIF_TERM mod_ref = enif_make_resource(env, wrapper);
3629+
enif_release_resource(wrapper);
3630+
req->result = enif_make_tuple2(env, ATOM_OK, mod_ref);
3631+
}
35713632
}
35723633
break;
35733634
}
@@ -3617,6 +3678,7 @@ static void process_request(py_request_t *req) {
36173678
enif_make_map_put(env, result_map,
36183679
enif_make_atom(env, "gc_stats"), gc_stats_list, &result_map);
36193680
}
3681+
/* If gen_stats alloc failed, we skip gc_stats but continue with other stats */
36203682
}
36213683
Py_DECREF(stats);
36223684
}
@@ -4053,6 +4115,9 @@ static void multi_executor_stop(void) {
40534115
return;
40544116
}
40554117

4118+
/* Allocate shutdown requests for all executors */
4119+
py_request_t *shutdown_reqs[MAX_EXECUTORS] = {0};
4120+
40564121
/* Signal shutdown and send shutdown requests to all executors */
40574122
for (int i = 0; i < g_num_executors; i++) {
40584123
executor_t *exec = &g_executors[i];
@@ -4062,18 +4127,25 @@ static void multi_executor_stop(void) {
40624127
if (shutdown_req != NULL) {
40634128
request_init(shutdown_req);
40644129
shutdown_req->type = PY_REQ_SHUTDOWN;
4130+
shutdown_reqs[i] = shutdown_req;
40654131
multi_executor_enqueue(i, shutdown_req);
40664132
}
40674133
/* If alloc fails, the shutdown flag is already set, so executor
40684134
* will exit when it checks the flag */
40694135
}
40704136

4071-
/* Wait for all executors to finish */
4137+
/* Wait for all executors to finish and clean up shutdown requests */
40724138
for (int i = 0; i < g_num_executors; i++) {
40734139
executor_t *exec = &g_executors[i];
40744140
pthread_join(exec->thread, NULL);
40754141
pthread_mutex_destroy(&exec->mutex);
40764142
pthread_cond_destroy(&exec->cond);
4143+
4144+
/* Clean up the shutdown request */
4145+
if (shutdown_reqs[i] != NULL) {
4146+
request_cleanup(shutdown_reqs[i]);
4147+
enif_free(shutdown_reqs[i]);
4148+
}
40774149
}
40784150

40794151
g_multi_executor_initialized = false;

0 commit comments

Comments
 (0)