From 9cb0ad863c0781062388fe1bb2a5d279dc29e734 Mon Sep 17 00:00:00 2001
From: Alan Wu <XrXr@users.noreply.github.com>
Date: Mon, 11 Dec 2023 17:41:27 -0500
Subject: [PATCH] YJIT: Fix missing arity check for splat calls to methods with
 optionals

Previously, for splat callsites that land in methods with optional
parameters, we didn't reject the case where the caller supplies too many
arguments. Accepting those calls previously caused YJIT to construct
corrupted control frames, which leads to crashes if the callee uses
certain stack walking methods such as Kernel#raise and String#gsub (for
setting up the frame-local `$~`).

Example crash in a debug build:

    Assertion Failed: ../vm_core.h:1375:VM_ENV_FLAGS:FIXNUM_P(flags)
---
 bootstraptest/test_yjit.rb | 26 ++++++++++++++++++++++++++
 yjit/src/codegen.rs        | 11 ++++++++---
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index ba288a7936b5f3..5e64c83b5c41c2 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -1,3 +1,29 @@
+# regression test for arity check with splat
+assert_equal '[:ae, :ae]', %q{
+  def req_one(a_, b_ = 1) = raise
+
+  def test(args)
+    req_one *args
+  rescue ArgumentError
+    :ae
+  end
+
+  [test(Array.new 5), test([])]
+}
+
+# regression test for arity check with splat and send
+assert_equal '[:ae, :ae]', %q{
+  def two_reqs(a, b_, _ = 1) = a.gsub(a, a)
+
+  def test(name, args)
+    send(name, *args)
+  rescue ArgumentError
+    :ae
+  end
+
+  [test(:two_reqs, ["g", nil, nil, nil]), test(:two_reqs, ["g"])]
+}
+
 # regression test for GC marking stubs in invalidated code
 assert_normal_exit %q{
   garbage = Array.new(10_000) { [] } # create garbage to cause iseq movement
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index a3bbb95cd805ee..a50bed3bc3a369 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -6116,9 +6116,14 @@ fn gen_send_iseq(
             unsafe { rb_yjit_array_len(array) as u32}
         };
 
-        if opt_num == 0 && required_num != array_length as i32 + argc - 1 && !iseq_has_rest {
-            gen_counter_incr(asm, Counter::send_iseq_splat_arity_error);
-            return None;
+        // Arity check accounting for size of the splat. When callee has rest parameters, we insert
+        // runtime guards later in copy_splat_args_for_rest_callee()
+        if !iseq_has_rest {
+            let supplying = argc - 1 + array_length as i32;
+            if (required_num..=required_num + opt_num).contains(&supplying) == false {
+                gen_counter_incr(asm, Counter::send_iseq_splat_arity_error);
+                return None;
+            }
         }
 
         if iseq_has_rest && opt_num > 0 {