ZJIT: Check for VM stack overflow

Previously, the included test crashed or turned into an infinite
loop due to the missing check.

Author: XrXr

test/ruby/test_zjit.rb

@@ -2584,6 +2584,26 @@ def test(x)
     }, insns: [:opt_case_dispatch]
   end
 
+  def test_stack_overflow
+    assert_compiles 'nil', %q{
+      def recurse(n)
+        return if n == 0
+        recurse(n-1)
+        nil # no tail call
+      end
+
+      recurse(2)
+      recurse(2)
+      begin
+        recurse(20_000)
+      rescue SystemStackError
+        # Not asserting an exception is raised here since main
+        # thread stack size is environment-sensitive. Only
+        # that we don't crash or infinite loop.
+      end
+    }, call_threshold: 2
+  end
+
   def test_invokeblock
     assert_compiles '42', %q{
       def test

zjit/src/codegen.rs

@@ -20,6 +20,7 @@ use crate::hir::{iseq_to_hir, Block, BlockId, BranchEdge, Invariant, RangeType,
 use crate::hir::{Const, FrameState, Function, Insn, InsnId};
 use crate::hir_type::{types, Type};
 use crate::options::get_option;
+use crate::cast::IntoUsize;
 
 /// Ephemeral code generation state
 struct JITState {
@@ -1042,6 +1043,19 @@ fn gen_send_without_block_direct(
     args: Vec<Opnd>,
     state: &FrameState,
 ) -> lir::Opnd {
+    let local_size = unsafe { get_iseq_body_local_table_size(iseq) }.as_usize();
+    // Stack overflow check: fails if CFP<=SP at any point in the callee.
+    asm_comment!(asm, "stack overflow check");
+    let stack_growth = state.stack_size() + local_size + unsafe { get_iseq_body_stack_max(iseq) }.as_usize();
+    // vm_push_frame() checks it against a decremented cfp, and CHECK_VM_STACK_OVERFLOW0
+    // adds to the margin another control frame with `&bounds[1]`.
+    const { assert!(RUBY_SIZEOF_CONTROL_FRAME % SIZEOF_VALUE == 0, "sizeof(rb_control_frame_t) is a multiple of sizeof(VALUE)"); }
+    let cfp_growth = 2 * (RUBY_SIZEOF_CONTROL_FRAME / SIZEOF_VALUE);
+    let peak_offset = SIZEOF_VALUE * (stack_growth + cfp_growth);
+    let stack_limit = asm.add(SP, peak_offset.into());
+    asm.cmp(CFP, stack_limit);
+    asm.jbe(side_exit(jit, state, StackOverflow));
+
     // Save cfp->pc and cfp->sp for the caller frame
     gen_prepare_call_with_gc(asm, state);
     gen_save_sp(asm, state.stack().len() - args.len() - 1); // -1 for receiver
@@ -1059,8 +1073,7 @@ fn gen_send_without_block_direct(
     });
 
     asm_comment!(asm, "switch to new SP register");
-    let local_size = unsafe { get_iseq_body_local_table_size(iseq) } as usize;
-    let sp_offset = (state.stack().len() + local_size - args.len() + VM_ENV_DATA_SIZE as usize) * SIZEOF_VALUE;
+    let sp_offset = (state.stack().len() + local_size - args.len() + VM_ENV_DATA_SIZE.as_usize()) * SIZEOF_VALUE;
     let new_sp = asm.add(SP, sp_offset.into());
     asm.mov(SP, new_sp);
 

zjit/src/hir.rs

@@ -464,6 +464,7 @@ pub enum SideExitReason {
     Interrupt,
     BlockParamProxyModified,
     BlockParamProxyNotIseqOrIfunc,
+    StackOverflow,
 }
 
 impl std::fmt::Display for SideExitReason {

zjit/src/stats.rs

@@ -102,6 +102,7 @@ make_counters! {
         exit_callee_side_exit,
         exit_obj_to_string_fallback,
         exit_interrupt,
+        exit_stackoverflow,
         exit_optional_arguments,
         exit_block_param_proxy_modified,
         exit_block_param_proxy_not_iseq_or_ifunc,
@@ -232,6 +233,7 @@ pub fn exit_counter_ptr(reason: crate::hir::SideExitReason) -> *mut u64 {
         CalleeSideExit                => exit_callee_side_exit,
         ObjToStringFallback           => exit_obj_to_string_fallback,
         Interrupt                     => exit_interrupt,
+        StackOverflow                 => exit_stackoverflow,
         BlockParamProxyModified       => exit_block_param_proxy_modified,
         BlockParamProxyNotIseqOrIfunc => exit_block_param_proxy_not_iseq_or_ifunc,
     };