From d5f5a56bf291d2456366bfb824d4413d02465f87 Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashi.kokubun@shopify.com>
Date: Wed, 2 Jul 2025 13:01:24 -0700
Subject: [PATCH] ZJIT: Reject ISEQs with too-large stack_max (#13770)

---
 .github/workflows/zjit-macos.yml | 2 +-
 zjit/src/asm/arm64/mod.rs        | 2 +-
 zjit/src/codegen.rs              | 8 ++++++++
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/zjit-macos.yml b/.github/workflows/zjit-macos.yml
index 8e58605fe1..7060d6a252 100644
--- a/.github/workflows/zjit-macos.yml
+++ b/.github/workflows/zjit-macos.yml
@@ -125,6 +125,7 @@ jobs:
           ../src/bootstraptest/test_literal_suffix.rb \
           ../src/bootstraptest/test_load.rb \
           ../src/bootstraptest/test_marshal.rb \
+          ../src/bootstraptest/test_massign.rb \
           ../src/bootstraptest/test_method.rb \
           ../src/bootstraptest/test_objectspace.rb \
           ../src/bootstraptest/test_string.rb \
@@ -136,7 +137,6 @@ jobs:
           ../src/bootstraptest/test_yjit_rust_port.rb
         # ../src/bootstraptest/test_eval.rb \
         # ../src/bootstraptest/test_insns.rb \
-        # ../src/bootstraptest/test_massign.rb \
         # ../src/bootstraptest/test_proc.rb \
         # ../src/bootstraptest/test_ractor.rb \
         # ../src/bootstraptest/test_yjit.rb \
diff --git a/zjit/src/asm/arm64/mod.rs b/zjit/src/asm/arm64/mod.rs
index 1e1b125eaa..ef477821aa 100644
--- a/zjit/src/asm/arm64/mod.rs
+++ b/zjit/src/asm/arm64/mod.rs
@@ -936,7 +936,7 @@ pub fn stur(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
     let bytes: [u8; 4] = match (rt, rn) {
         (A64Opnd::Reg(rt), A64Opnd::Mem(rn)) => {
             assert!(rn.num_bits == 32 || rn.num_bits == 64);
-            assert!(mem_disp_fits_bits(rn.disp), "Expected displacement to be 9 bits or less");
+            assert!(mem_disp_fits_bits(rn.disp), "Expected displacement {} to be 9 bits or less", rn.disp);
 
             LoadStore::stur(rt.reg_no, rn.base_reg_no, rn.disp as i16, rn.num_bits).into()
         },
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 9fa088c0d1..419fc50983 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -72,6 +72,14 @@ pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, _ec: EcPtr) -> *co
         return std::ptr::null();
     }
 
+    // Reject ISEQs with very large temp stacks.
+    // We cannot encode too large offsets to access locals in arm64.
+    let stack_max = unsafe { rb_get_iseq_body_stack_max(iseq) };
+    if stack_max >= i8::MAX as u32 {
+        debug!("ISEQ stack too large: {stack_max}");
+        return std::ptr::null();
+    }
+
     // Take a lock to avoid writing to ISEQ in parallel with Ractors.
     // with_vm_lock() does nothing if the program doesn't use Ractors.
     let code_ptr = with_vm_lock(src_loc!(), || {