From 0ce2bdc76dd17aa3d42a352a6244c87a51e7606d Mon Sep 17 00:00:00 2001
From: John Hawthorn <john@hawthorn.email>
Date: Wed, 12 Apr 2023 16:22:22 -0700
Subject: [PATCH] YJIT: Fix missing argc check in known cfuncs

Previously we were missing a compile-time check that the known cfuncs
receive the correct number of arguments.

We noticied this because in particular when using ARGS_SPLAT, which also
wasn't checked, YJIT would crash on code which was otherwise correct
(didn't raise exceptions in the VM).

This still supports vararg (argc == -1) cfuncs. I added an additional
assertion that when we use the specialized codegen for one of these
known functions that the argc are popped off the stack correctly, which
should help ensure they're implemented correctly (previously the crash
was usually observed on a future `leave` insn).

[Bug #19595]
---
 bootstraptest/test_yjit.rb | 32 ++++++++++++++++++++++++++++++++
 yjit/src/codegen.rs        |  4 +++-
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 34cb87ed88..ed8bd94e7c 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -1227,6 +1227,38 @@ assert_equal '42', %q{
   run
 }
 
+# splatting an empty array on a specialized method
+assert_equal 'ok', %q{
+  def run
+    "ok".to_s(*[])
+  end
+
+  run
+  run
+}
+
+# splatting an single element array on a specialized method
+assert_equal '[1]', %q{
+  def run
+    [].<<(*[1])
+  end
+
+  run
+  run
+}
+
+# specialized method with wrong args
+assert_equal 'ok', %q{
+  def run(x)
+    "bad".to_s(123) if x
+  rescue
+    :ok
+  end
+
+  run(false)
+  run(true)
+}
+
 # getinstancevariable on Symbol
 assert_equal '[nil, nil]', %q{
   # @foo to exercise the getinstancevariable instruction
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 7c302707db..e853f1308a 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -4987,10 +4987,12 @@ fn gen_send_cfunc(
     }
 
     // Delegate to codegen for C methods if we have it.
-    if kw_arg.is_null() && flags & VM_CALL_OPT_SEND == 0 {
+    if kw_arg.is_null() && flags & VM_CALL_OPT_SEND == 0 && flags & VM_CALL_ARGS_SPLAT == 0 && (cfunc_argc == -1 || argc == cfunc_argc) {
         let codegen_p = lookup_cfunc_codegen(unsafe { (*cme).def });
+        let expected_stack_after = asm.ctx.get_stack_size() as i32 - argc;
         if let Some(known_cfunc_codegen) = codegen_p {
             if known_cfunc_codegen(jit, asm, ocb, ci, cme, block, argc, recv_known_klass) {
+                assert_eq!(expected_stack_after, asm.ctx.get_stack_size() as i32);
                 // cfunc codegen generated code. Terminate the block so
                 // there isn't multiple calls in the same block.
                 jump_to_next_insn(jit, asm, ocb);