Backport WASI setjmp handler memory leak fixes to Ruby 3.3

6 files changed, 57 insertions, 0 deletions

diff --git a/cont.c b/cont.c
index 55040d3d38..d22bcc5137 100644
--- a/cont.c
+++ b/cont.c
@@ -1518,6 +1518,51 @@ cont_restore_thread(rb_context_t *cont)
             rb_raise(rb_eRuntimeError, "can't call across trace_func");
         }
 
+#if defined(__wasm__) && !defined(__EMSCRIPTEN__)
+        if (th->ec->tag != sec->tag) {
+            /* find the lowest common ancestor tag of the current EC and the saved EC */
+
+            struct rb_vm_tag *lowest_common_ancestor = NULL;
+            size_t num_tags = 0;
+            size_t num_saved_tags = 0;
+            for (struct rb_vm_tag *tag = th->ec->tag; tag != NULL; tag = tag->prev) {
+                ++num_tags;
+            }
+            for (struct rb_vm_tag *tag = sec->tag; tag != NULL; tag = tag->prev) {
+                ++num_saved_tags;
+            }
+
+            size_t min_tags = num_tags <= num_saved_tags ? num_tags : num_saved_tags;
+
+            struct rb_vm_tag *tag = th->ec->tag;
+            while (num_tags > min_tags) {
+                tag = tag->prev;
+                --num_tags;
+            }
+
+            struct rb_vm_tag *saved_tag = sec->tag;
+            while (num_saved_tags > min_tags) {
+                saved_tag = saved_tag->prev;
+                --num_saved_tags;
+            }
+
+            while (min_tags > 0) {
+                if (tag == saved_tag) {
+                    lowest_common_ancestor = tag;
+                    break;
+                }
+                tag = tag->prev;
+                saved_tag = saved_tag->prev;
+                --min_tags;
+            }
+
+            /* free all the jump buffers between the current EC's tag and the lowest common ancestor tag */
+            for (struct rb_vm_tag *tag = th->ec->tag; tag != lowest_common_ancestor; tag = tag->prev) {
+                rb_vm_tag_jmpbuf_deinit(&tag->buf);
+            }
+        }
+#endif
+
         /* copy vm stack */
 #ifdef CAPTURE_JUST_VALID_VM_STACK
         MEMCPY(th->ec->vm_stack,
diff --git a/signal.c b/signal.c
index 589ee05cbd..ff116c86d4 100644
--- a/signal.c
+++ b/signal.c
@@ -850,6 +850,7 @@ check_stack_overflow(int sig, const uintptr_t addr, const ucontext_t *ctx)
              * otherwise it can cause stack overflow again at the same
              * place. */
             if ((crit = (!ec->tag->prev || !--uplevel)) != FALSE) break;
+            rb_vm_tag_jmpbuf_deinit(&ec->tag->buf);
             ec->tag = ec->tag->prev;
         }
         reset_sigmask(sig);
diff --git a/vm.c b/vm.c
index 9585135dff..259d16b85a 100644
--- a/vm.c
+++ b/vm.c
@@ -2729,6 +2729,7 @@ vm_exec_handle_exception(rb_execution_context_t *ec, enum ruby_tag_type state, V
             if (VM_FRAME_FINISHED_P(ec->cfp)) {
                 rb_vm_pop_frame(ec);
                 ec->errinfo = (VALUE)err;
+                rb_vm_tag_jmpbuf_deinit(&ec->tag->buf);
                 ec->tag = ec->tag->prev;
                 EC_JUMP_TAG(ec, state);
             }
diff --git a/vm_trace.c b/vm_trace.c
index 7050d1efc2..123708f803 100644
--- a/vm_trace.c
+++ b/vm_trace.c
@@ -455,6 +455,7 @@ rb_exec_event_hooks(rb_trace_arg_t *trace_arg, rb_hook_list_t *hooks, int pop_p)
             if (state) {
                 if (pop_p) {
                     if (VM_FRAME_FINISHED_P(ec->cfp)) {
+                        rb_vm_tag_jmpbuf_deinit(&ec->tag->buf);
                         ec->tag = ec->tag->prev;
                     }
                     rb_vm_pop_frame(ec);
diff --git a/wasm/setjmp.c b/wasm/setjmp.c
index ebbf8949c1..32ede68c09 100644
--- a/wasm/setjmp.c
+++ b/wasm/setjmp.c
@@ -143,9 +143,11 @@ rb_wasm_try_catch_init(struct rb_wasm_try_catch *try_catch,
     try_catch->try_f = try_f;
     try_catch->catch_f = catch_f;
     try_catch->context = context;
+    try_catch->stack_pointer = NULL;
 }
 
 // NOTE: This function is not processed by Asyncify due to a call of asyncify_stop_rewind
+__attribute__((noinline))
 void
 rb_wasm_try_catch_loop_run(struct rb_wasm_try_catch *try_catch, rb_wasm_jmp_buf *target)
 {
@@ -154,6 +156,10 @@ rb_wasm_try_catch_loop_run(struct rb_wasm_try_catch *try_catch, rb_wasm_jmp_buf
 
     target->state = JMP_BUF_STATE_CAPTURED;
 
+    if (try_catch->stack_pointer == NULL) {
+        try_catch->stack_pointer = rb_wasm_get_stack_pointer();
+    }
+
     switch ((enum try_catch_phase)try_catch->state) {
     case TRY_CATCH_PHASE_MAIN:
         // may unwind
@@ -175,6 +181,8 @@ rb_wasm_try_catch_loop_run(struct rb_wasm_try_catch *try_catch, rb_wasm_jmp_buf
             // stop unwinding
             // (but call stop_rewind to update the asyncify state to "normal" from "unwind")
             asyncify_stop_rewind();
+            // reset the stack pointer to what it was before the most recent call to try_f or catch_f
+            rb_wasm_set_stack_pointer(try_catch->stack_pointer);
             // clear the active jmpbuf because it's already stopped
             _rb_wasm_active_jmpbuf = NULL;
             // reset jmpbuf state to be able to unwind again
diff --git a/wasm/setjmp.h b/wasm/setjmp.h
index cc14df33be..e65bfc0ca0 100644
--- a/wasm/setjmp.h
+++ b/wasm/setjmp.h
@@ -65,6 +65,7 @@ struct rb_wasm_try_catch {
     rb_wasm_try_catch_func_t try_f;
     rb_wasm_try_catch_func_t catch_f;
     void *context;
+    void *stack_pointer;
     int state;
 };