YJIT: Plug native stack overflow

Previously, TestStack#test_machine_stack_size failed pretty consistently
on ARM64 macOS, with Rust code and part of the interpreter used for
per-instruction fallback (rb_vm_invokeblock() and friends) touching the
stack guard page and crashing with SEGV. I've also seen the same test
fail on x64 Linux, though with a different symptom.

5 files changed, 21 insertions, 3 deletions

diff --git a/vm_insnhelper.c b/vm_insnhelper.c
index a325d070a2..f98e6c6e75 100644
--- a/vm_insnhelper.c
+++ b/vm_insnhelper.c
@@ -93,6 +93,7 @@ rb_ec_stack_overflow(rb_execution_context_t *ec, int crit)
 #endif
 }
 
+static inline void stack_check(rb_execution_context_t *ec);
 
 #if VM_CHECK_MODE > 0
 static int
@@ -5565,6 +5566,7 @@ vm_sendish(
 VALUE
 rb_vm_send(rb_execution_context_t *ec, rb_control_frame_t *reg_cfp, CALL_DATA cd, ISEQ blockiseq)
 {
+    stack_check(ec);
     VALUE bh = vm_caller_setup_arg_block(ec, GET_CFP(), cd->ci, blockiseq, false);
     VALUE val = vm_sendish(ec, GET_CFP(), cd, bh, mexp_search_method);
     VM_EXEC(ec, val);
@@ -5574,6 +5576,7 @@ rb_vm_send(rb_execution_context_t *ec, rb_control_frame_t *reg_cfp, CALL_DATA cd
 VALUE
 rb_vm_opt_send_without_block(rb_execution_context_t *ec, rb_control_frame_t *reg_cfp, CALL_DATA cd)
 {
+    stack_check(ec);
     VALUE bh = VM_BLOCK_HANDLER_NONE;
     VALUE val = vm_sendish(ec, GET_CFP(), cd, bh, mexp_search_method);
     VM_EXEC(ec, val);
@@ -5583,6 +5586,7 @@ rb_vm_opt_send_without_block(rb_execution_context_t *ec, rb_control_frame_t *reg
 VALUE
 rb_vm_invokesuper(rb_execution_context_t *ec, rb_control_frame_t *reg_cfp, CALL_DATA cd, ISEQ blockiseq)
 {
+    stack_check(ec);
     VALUE bh = vm_caller_setup_arg_block(ec, GET_CFP(), cd->ci, blockiseq, true);
     VALUE val = vm_sendish(ec, GET_CFP(), cd, bh, mexp_search_super);
     VM_EXEC(ec, val);
@@ -5592,6 +5596,7 @@ rb_vm_invokesuper(rb_execution_context_t *ec, rb_control_frame_t *reg_cfp, CALL_
 VALUE
 rb_vm_invokeblock(rb_execution_context_t *ec, rb_control_frame_t *reg_cfp, CALL_DATA cd)
 {
+    stack_check(ec);
     VALUE bh = VM_BLOCK_HANDLER_NONE;
     VALUE val = vm_sendish(ec, GET_CFP(), cd, bh, mexp_search_invokeblock);
     VM_EXEC(ec, val);
diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index baa71fc6cf..db93672a3d 100644
--- a/yjit/bindgen/src/main.rs
+++ b/yjit/bindgen/src/main.rs
@@ -449,6 +449,7 @@ fn main() {
         .allowlist_function("rb_obj_class")
         .allowlist_function("rb_obj_is_proc")
         .allowlist_function("rb_vm_base_ptr")
+        .allowlist_function("rb_ec_stack_check")
 
         // We define VALUE manually, don't import it
         .blocklist_type("VALUE")
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index 57b742724b..8261051315 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -2482,9 +2482,6 @@ fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -
         (target.get_blockid(), target.get_ctx())
     };
 
-    let cb = CodegenGlobals::get_inline_cb();
-    let ocb = CodegenGlobals::get_outlined_cb();
-
     let (cfp, original_interp_sp) = unsafe {
         let cfp = get_ec_cfp(ec);
         let original_interp_sp = get_cfp_sp(cfp);
@@ -2506,9 +2503,18 @@ fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -
         // So we do it here instead.
         rb_set_cfp_sp(cfp, reconned_sp);
 
+        // Bail if we're about to run out of native stack space.
+        // We've just reconstructed interpreter state.
+        if rb_ec_stack_check(ec as _) != 0 {
+            return CodegenGlobals::get_stub_exit_code().raw_ptr();
+        }
+
         (cfp, original_interp_sp)
     };
 
+    let cb = CodegenGlobals::get_inline_cb();
+    let ocb = CodegenGlobals::get_outlined_cb();
+
     // Try to find an existing compiled version of this block
     let mut block = find_block_version(target_blockid, &target_ctx);
     let mut branch_modified = false;
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index ccf5b530d6..3e2961aa72 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -1173,6 +1173,7 @@ extern "C" {
     ) -> *const rb_callable_method_entry_t;
     pub fn rb_obj_info(obj: VALUE) -> *const ::std::os::raw::c_char;
     pub fn rb_class_allocate_instance(klass: VALUE) -> VALUE;
+    pub fn rb_ec_stack_check(ec: *mut rb_execution_context_struct) -> ::std::os::raw::c_int;
     pub fn rb_shape_id_offset() -> i32;
     pub fn rb_shape_get_shape_by_id(shape_id: shape_id_t) -> *mut rb_shape_t;
     pub fn rb_shape_get_shape_id(obj: VALUE) -> shape_id_t;
diff --git a/yjit/src/yjit.rs b/yjit/src/yjit.rs
index 97799923d2..2aed3c6a4b 100644
--- a/yjit/src/yjit.rs
+++ b/yjit/src/yjit.rs
@@ -114,6 +114,11 @@ fn rb_bug_panic_hook() {
 /// See [jit_compile_exception] for details.
 #[no_mangle]
 pub extern "C" fn rb_yjit_iseq_gen_entry_point(iseq: IseqPtr, ec: EcPtr, jit_exception: bool) -> *const u8 {
+    // Don't compile when there is insufficient native stack space
+    if unsafe { rb_ec_stack_check(ec as _) } != 0 {
+        return std::ptr::null();
+    }
+
     // Reject ISEQs with very large temp stacks,
     // this will allow us to use u8/i8 values to track stack_size and sp_offset
     let stack_max = unsafe { rb_get_iseq_body_stack_max(iseq) };