From 88e0ac35a3b6f66074307421b208db62ae5ffe39 Mon Sep 17 00:00:00 2001 From: Max Bernstein Date: Wed, 17 Sep 2025 16:04:49 -0400 Subject: [PATCH] ZJIT: Prevent custom allocator in ObjectAllocClass --- object.c | 9 +++++++++ test/ruby/test_zjit.rb | 23 +++++++++++++++++++++++ zjit.c | 19 +++++++++++++++++++ zjit/bindgen/src/main.rs | 2 ++ zjit/src/cruby_bindings.inc.rs | 2 ++ zjit/src/hir.rs | 30 +++++++++++++++++++++++------- 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); } }