ZJIT: Prevent custom allocator in ObjectAllocClass

6 files changed, 78 insertions, 7 deletions

diff --git a/object.c b/object.c
index 56afc7e99b..c433e4bb17 100644
--- a/object.c
+++ b/object.c
@@ -2192,6 +2192,15 @@ class_get_alloc_func(VALUE klass)
     return allocator;
 }
 
+// Might return NULL.
+rb_alloc_func_t
+rb_zjit_class_get_alloc_func(VALUE klass)
+{
+    assert(RCLASS_INITIALIZED_P(klass));
+    assert(!RCLASS_SINGLETON_P(klass));
+    return rb_get_alloc_func(klass);
+}
+
 static VALUE
 class_call_alloc_func(rb_alloc_func_t allocator, VALUE klass)
 {
diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
index e530fb797f..a6296084ea 100644
--- a/test/ruby/test_zjit.rb
+++ b/test/ruby/test_zjit.rb
@@ -863,6 +863,29 @@ class TestZJIT < Test::Unit::TestCase
     }, insns: [:opt_new], call_threshold: 2
   end
 
+  def test_opt_new_with_custom_allocator
+    assert_compiles '"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"', %q{
+      require "digest"
+      def test = Digest::SHA256.new.hexdigest
+      test; test
+    }, insns: [:opt_new], call_threshold: 2
+  end
+
+  def test_opt_new_with_custom_allocator_raises
+    assert_compiles '[42, 42]', %q{
+      require "digest"
+      class C < Digest::Base; end
+      def test
+        begin
+          Digest::Base.new
+        rescue NotImplementedError
+          42
+        end
+      end
+      [test, test]
+    }, insns: [:opt_new], call_threshold: 2
+  end
+
   def test_new_hash_empty
     assert_compiles '{}', %q{
       def test = {}
diff --git a/zjit.c b/zjit.c
index 4b29578b4a..21618c39b1 100644
--- a/zjit.c
+++ b/zjit.c
@@ -175,6 +175,25 @@ bool rb_zjit_cme_is_cfunc(const rb_callable_method_entry_t *me, const void *func
 const struct rb_callable_method_entry_struct *
 rb_zjit_vm_search_method(VALUE cd_owner, struct rb_call_data *cd, VALUE recv);
 
+bool
+rb_zjit_class_initialized_p(VALUE klass)
+{
+    return RCLASS_INITIALIZED_P(klass);
+}
+
+rb_alloc_func_t rb_zjit_class_get_alloc_func(VALUE klass);
+
+VALUE rb_class_allocate_instance(VALUE klass);
+
+bool
+rb_zjit_class_has_default_allocator(VALUE klass)
+{
+    assert(RCLASS_INITIALIZED_P(klass));
+    assert(!RCLASS_SINGLETON_P(klass));
+    rb_alloc_func_t alloc = rb_zjit_class_get_alloc_func(klass);
+    return alloc == rb_class_allocate_instance;
+}
+
 // Primitives used by zjit.rb. Don't put other functions below, which wouldn't use them.
 VALUE rb_zjit_assert_compiles(rb_execution_context_t *ec, VALUE self);
 VALUE rb_zjit_stats(rb_execution_context_t *ec, VALUE self, VALUE target_key);
diff --git a/zjit/bindgen/src/main.rs b/zjit/bindgen/src/main.rs
index 6e9a5a529f..3bffdfd9ff 100644
--- a/zjit/bindgen/src/main.rs
+++ b/zjit/bindgen/src/main.rs
@@ -333,6 +333,8 @@ fn main() {
         .allowlist_function("rb_insn_name")
         .allowlist_function("rb_insn_len")
         .allowlist_function("rb_yarv_class_of")
+        .allowlist_function("rb_zjit_class_initialized_p")
+        .allowlist_function("rb_zjit_class_has_default_allocator")
         .allowlist_function("rb_get_ec_cfp")
         .allowlist_function("rb_get_cfp_iseq")
         .allowlist_function("rb_get_cfp_pc")
diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs
index dfa1be9b8f..4b83051a67 100644
--- a/zjit/src/cruby_bindings.inc.rs
+++ b/zjit/src/cruby_bindings.inc.rs
@@ -946,6 +946,8 @@ unsafe extern "C" {
         cd: *mut rb_call_data,
         recv: VALUE,
     ) -> *const rb_callable_method_entry_struct;
+    pub fn rb_zjit_class_initialized_p(klass: VALUE) -> bool;
+    pub fn rb_zjit_class_has_default_allocator(klass: VALUE) -> bool;
     pub fn rb_iseq_encoded_size(iseq: *const rb_iseq_t) -> ::std::os::raw::c_uint;
     pub fn rb_iseq_pc_at_idx(iseq: *const rb_iseq_t, insn_idx: u32) -> *mut VALUE;
     pub fn rb_iseq_opcode_at_pc(iseq: *const rb_iseq_t, pc: *const VALUE) -> ::std::os::raw::c_int;
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 7dcf1c6ba8..3a7e87293d 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -1918,14 +1918,30 @@ impl Function {
                     }
                     Insn::ObjectAlloc { val, state } => {
                         let val_type = self.type_of(val);
-                        if val_type.is_subtype(types::Class) && val_type.ruby_object_known() {
-                            let class = val_type.ruby_object().unwrap();
-                            let replacement = self.push_insn(block, Insn::ObjectAllocClass { class, state });
-                            self.insn_types[replacement.0] = self.infer_type(replacement);
-                            self.make_equal_to(insn_id, replacement);
-                        } else {
-                            self.push_insn_id(block, insn_id);
+                        if !val_type.is_subtype(types::Class) {
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+                        let Some(class) = val_type.ruby_object() else {
+                            self.push_insn_id(block, insn_id); continue;
+                        };
+                        // See class_get_alloc_func in object.c; if the class isn't initialized, is
+                        // a singleton class, or has a custom allocator, ObjectAlloc might raise an
+                        // exception or run arbitrary code.
+                        //
+                        // We also need to check if the class is initialized or a singleton before trying to read the allocator, otherwise it might raise.
+                        if !unsafe { rb_zjit_class_initialized_p(class) } {
+                            self.push_insn_id(block, insn_id); continue;
                         }
+                        if unsafe { rb_zjit_singleton_class_p(class) } {
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+                        if !unsafe { rb_zjit_class_has_default_allocator(class) } {
+                            // Custom or NULL allocator; could run arbitrary code.
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+                        let replacement = self.push_insn(block, Insn::ObjectAllocClass { class, state });
+                        self.insn_types[replacement.0] = self.infer_type(replacement);
+                        self.make_equal_to(insn_id, replacement);
                     }
                     _ => { self.push_insn_id(block, insn_id); }
                 }