From 20fc91df395b834ecaee0bdb29df8da224fdf89a Mon Sep 17 00:00:00 2001
From: Alan Wu <XrXr@users.noreply.github.com>
Date: Wed, 1 Oct 2025 23:53:48 -0400
Subject: [PATCH] YJIT: Prevent making a branch from a dead block to a live
 block

I'm seeing some memory corruption in the wild on blocks in
`IseqPayload::dead_blocks`. While I unfortunately can't recreate the
issue, (For all I know, it could be some external code corrupting YJIT's
memory.) establishing a link between dead blocks and live blocks seems
fishy enough that we ought to prevent it. When it did happen, it might've
had bad interacts with Code GC and the optimization to immediately
free empty blocks.
---
 yjit/src/core.rs | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index cfe55b8c76..2999f151bf 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -3591,6 +3591,13 @@ fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -
             return CodegenGlobals::get_stub_exit_code().raw_ptr(cb);
         }
 
+        // Bail if this branch is housed in an invalidated (dead) block.
+        // This only happens in rare invalidation scenarios and we need
+        // to avoid linking a dead block to a live block with a branch.
+        if branch.block.get().as_ref().iseq.get().is_null() {
+            return CodegenGlobals::get_stub_exit_code().raw_ptr(cb);
+        }
+
         (cfp, original_interp_sp)
     };
 
@@ -4297,11 +4304,9 @@ pub fn invalidate_block_version(blockref: &BlockRef) {
     incr_counter!(invalidation_count);
 }
 
-// We cannot deallocate blocks immediately after invalidation since there
-// could be stubs waiting to access branch pointers. Return stubs can do
-// this since patching the code for setting up return addresses does not
-// affect old return addresses that are already set up to use potentially
-// invalidated branch pointers. Example:
+// We cannot deallocate blocks immediately after invalidation since patching the code for setting
+// up return addresses does not affect outstanding return addresses that are on stack and will use
+// invalidated branch pointers when hit. Example:
 //   def foo(n)
 //     if n == 2
 //       # 1.times.each to create a cfunc frame to preserve the JIT frame
@@ -4309,13 +4314,16 @@ pub fn invalidate_block_version(blockref: &BlockRef) {
 //       return 1.times.each { Object.define_method(:foo) {} }
 //     end
 //
-//     foo(n + 1)
+//     foo(n + 1) # The block for this call houses the return branch stub
 //   end
 //   p foo(1)
 pub fn delayed_deallocation(blockref: BlockRef) {
     block_assumptions_free(blockref);
 
-    let payload = get_iseq_payload(unsafe { blockref.as_ref() }.iseq.get()).unwrap();
+    let block = unsafe { blockref.as_ref() };
+    // Set null ISEQ on the block to signal that it's dead.
+    let iseq = block.iseq.replace(ptr::null());
+    let payload = get_iseq_payload(iseq).unwrap();
     payload.dead_blocks.push(blockref);
 }