Fix use-after-free in WeakKeyMap#clear

[Bug #20691]

If the WeakKeyMap has been marked but sweeping hasn't started yet and we
cann WeakKeyMap#clear, then there could be a use-after-free because we do
not call rb_gc_remove_weak to remove the key from the GC.

For example, the following code triggers use-after-free errors in Valgrind:

    map = ObjectSpace::WeakKeyMap.new

    1_000.times do
      1_000.times do
        map[Object.new] = nil
      end

      map.clear
    end

Output from Valgrind:

    ==61230== Invalid read of size 8
    ==61230==    at 0x25CAF8: gc_update_weak_references (default.c:5593)
    ==61230==    by 0x25CAF8: gc_marks_finish (default.c:5641)
    ==61230==    by 0x26031C: gc_marks_continue (default.c:5987)
    ==61230==    by 0x26031C: gc_continue (default.c:2255)
    ==61230==    by 0x2605FC: newobj_cache_miss (default.c:2589)
    ==61230==    by 0x26111F: newobj_alloc (default.c:2622)
    ==61230==    by 0x26111F: rb_gc_impl_new_obj (default.c:2701)
    ==61230==    by 0x26111F: newobj_of (gc.c:890)
    ==61230==    by 0x26111F: rb_wb_protected_newobj_of (gc.c:917)
    ==61230==    by 0x2DE218: rb_class_allocate_instance (object.c:131)
    ==61230==    by 0x2E32A8: class_call_alloc_func (object.c:2141)
    ==61230==    by 0x2E32A8: rb_class_alloc (object.c:2113)
    ==61230==    by 0x2E32A8: rb_class_new_instance_pass_kw (object.c:2172)
    ==61230==    by 0x4296BC: vm_call_cfunc_with_frame_ (vm_insnhelper.c:3788)
    ==61230==    by 0x44A9CD: vm_sendish (vm_insnhelper.c:5955)
    ==61230==    by 0x44A9CD: vm_exec_core (insns.def:898)
    ==61230==    by 0x43A0E4: rb_vm_exec (vm.c:2564)
    ==61230==    by 0x2341B4: rb_ec_exec_node (eval.c:281)
    ==61230==    by 0x236258: ruby_run_node (eval.c:319)
    ==61230==    by 0x15D665: rb_main (main.c:43)
    ==61230==    by 0x15D665: main (main.c:62)
    ==61230==  Address 0x2159cb00 is 0 bytes inside a block of size 8 free'd
    ==61230==    at 0x4849B2C: free (vg_replace_malloc.c:989)
    ==61230==    by 0x248EF1: rb_gc_impl_free (default.c:8512)
    ==61230==    by 0x248EF1: rb_gc_impl_free (default.c:8493)
    ==61230==    by 0x248EF1: ruby_sized_xfree.constprop.0 (gc.c:4178)
    ==61230==    by 0x4627EC: wkmap_free_table_i (weakmap.c:652)
    ==61230==    by 0x3A54AF: apply_functor (st.c:1633)
    ==61230==    by 0x3A54AF: st_general_foreach (st.c:1543)
    ==61230==    by 0x3A54AF: rb_st_foreach (st.c:1640)
    ==61230==    by 0x46203C: wkmap_clear (weakmap.c:973)
    ==61230==    by 0x4296BC: vm_call_cfunc_with_frame_ (vm_insnhelper.c:3788)
    ==61230==    by 0x44A9CD: vm_sendish (vm_insnhelper.c:5955)
    ==61230==    by 0x44A9CD: vm_exec_core (insns.def:898)
    ==61230==    by 0x43A0E4: rb_vm_exec (vm.c:2564)
    ==61230==    by 0x2341B4: rb_ec_exec_node (eval.c:281)
    ==61230==    by 0x236258: ruby_run_node (eval.c:319)
    ==61230==    by 0x15D665: rb_main (main.c:43)
    ==61230==    by 0x15D665: main (main.c:62)
    ==61230==  Block was alloc'd at
    ==61230==    at 0x484680F: malloc (vg_replace_malloc.c:446)
    ==61230==    by 0x25C68E: rb_gc_impl_malloc (default.c:8527)
    ==61230==    by 0x4622E9: wkmap_aset_replace (weakmap.c:817)
    ==61230==    by 0x3A4D02: rb_st_update (st.c:1487)
    ==61230==    by 0x4623E4: wkmap_aset (weakmap.c:854)
    ==61230==    by 0x4296BC: vm_call_cfunc_with_frame_ (vm_insnhelper.c:3788)
    ==61230==    by 0x44A9CD: vm_sendish (vm_insnhelper.c:5955)
    ==61230==    by 0x44A9CD: vm_exec_core (insns.def:898)
    ==61230==    by 0x43A0E4: rb_vm_exec (vm.c:2564)
    ==61230==    by 0x2341B4: rb_ec_exec_node (eval.c:281)
    ==61230==    by 0x236258: ruby_run_node (eval.c:319)
    ==61230==    by 0x15D665: rb_main (main.c:43)
    ==61230==    by 0x15D665: main (main.c:62)
    ==61230==
    ==61230== Invalid write of size 8
    ==61230==    at 0x25CB3B: gc_update_weak_references (default.c:5598)
    ==61230==    by 0x25CB3B: gc_marks_finish (default.c:5641)
    ==61230==    by 0x26031C: gc_marks_continue (default.c:5987)
    ==61230==    by 0x26031C: gc_continue (default.c:2255)
    ==61230==    by 0x2605FC: newobj_cache_miss (default.c:2589)
    ==61230==    by 0x26111F: newobj_alloc (default.c:2622)
    ==61230==    by 0x26111F: rb_gc_impl_new_obj (default.c:2701)
    ==61230==    by 0x26111F: newobj_of (gc.c:890)
    ==61230==    by 0x26111F: rb_wb_protected_newobj_of (gc.c:917)
    ==61230==    by 0x2DE218: rb_class_allocate_instance (object.c:131)
    ==61230==    by 0x2E32A8: class_call_alloc_func (object.c:2141)
    ==61230==    by 0x2E32A8: rb_class_alloc (object.c:2113)
    ==61230==    by 0x2E32A8: rb_class_new_instance_pass_kw (object.c:2172)
    ==61230==    by 0x4296BC: vm_call_cfunc_with_frame_ (vm_insnhelper.c:3788)
    ==61230==    by 0x44A9CD: vm_sendish (vm_insnhelper.c:5955)
    ==61230==    by 0x44A9CD: vm_exec_core (insns.def:898)
    ==61230==    by 0x43A0E4: rb_vm_exec (vm.c:2564)
    ==61230==    by 0x2341B4: rb_ec_exec_node (eval.c:281)
    ==61230==    by 0x236258: ruby_run_node (eval.c:319)
    ==61230==    by 0x15D665: rb_main (main.c:43)
    ==61230==    by 0x15D665: main (main.c:62)
    ==61230==  Address 0x2159cb00 is 0 bytes inside a block of size 8 free'd
    ==61230==    at 0x4849B2C: free (vg_replace_malloc.c:989)
    ==61230==    by 0x248EF1: rb_gc_impl_free (default.c:8512)
    ==61230==    by 0x248EF1: rb_gc_impl_free (default.c:8493)
    ==61230==    by 0x248EF1: ruby_sized_xfree.constprop.0 (gc.c:4178)
    ==61230==    by 0x4627EC: wkmap_free_table_i (weakmap.c:652)
    ==61230==    by 0x3A54AF: apply_functor (st.c:1633)
    ==61230==    by 0x3A54AF: st_general_foreach (st.c:1543)
    ==61230==    by 0x3A54AF: rb_st_foreach (st.c:1640)
    ==61230==    by 0x46203C: wkmap_clear (weakmap.c:973)
    ==61230==    by 0x4296BC: vm_call_cfunc_with_frame_ (vm_insnhelper.c:3788)
    ==61230==    by 0x44A9CD: vm_sendish (vm_insnhelper.c:5955)
    ==61230==    by 0x44A9CD: vm_exec_core (insns.def:898)
    ==61230==    by 0x43A0E4: rb_vm_exec (vm.c:2564)
    ==61230==    by 0x2341B4: rb_ec_exec_node (eval.c:281)
    ==61230==    by 0x236258: ruby_run_node (eval.c:319)
    ==61230==    by 0x15D665: rb_main (main.c:43)
    ==61230==    by 0x15D665: main (main.c:62)
    ==61230==  Block was alloc'd at
    ==61230==    at 0x484680F: malloc (vg_replace_malloc.c:446)
    ==61230==    by 0x25C68E: rb_gc_impl_malloc (default.c:8527)
    ==61230==    by 0x4622E9: wkmap_aset_replace (weakmap.c:817)
    ==61230==    by 0x3A4D02: rb_st_update (st.c:1487)
    ==61230==    by 0x4623E4: wkmap_aset (weakmap.c:854)
    ==61230==    by 0x4296BC: vm_call_cfunc_with_frame_ (vm_insnhelper.c:3788)
    ==61230==    by 0x44A9CD: vm_sendish (vm_insnhelper.c:5955)
    ==61230==    by 0x44A9CD: vm_exec_core (insns.def:898)
    ==61230==    by 0x43A0E4: rb_vm_exec (vm.c:2564)
    ==61230==    by 0x2341B4: rb_ec_exec_node (eval.c:281)
    ==61230==    by 0x236258: ruby_run_node (eval.c:319)
    ==61230==    by 0x15D665: rb_main (main.c:43)
    ==61230==    by 0x15D665: main (main.c:62)

Co-authored-by: Jean Boussier <byroot@ruby-lang.org>

Author: peterzhu2118

test/ruby/test_weakkeymap.rb

@@ -61,6 +61,20 @@ def test_clear
     refute @wm[k]
   end
 
+  def test_clear_bug_20691
+    assert_normal_exit(<<~RUBY)
+      map = ObjectSpace::WeakKeyMap.new
+
+      1_000.times do
+        1_000.times do
+          map[Object.new] = nil
+        end
+
+        map.clear
+      end
+    RUBY
+  end
+
   def test_inspect
     x = Object.new
     k = Object.new

weakmap.c

@@ -958,6 +958,17 @@ wkmap_has_key(VALUE self, VALUE key)
     return RBOOL(!UNDEF_P(wkmap_lookup(self, key)));
 }
 
+static int
+wkmap_clear_i(st_data_t key, st_data_t val, st_data_t data)
+{
+    VALUE self = (VALUE)data;
+
+    /* This WeakKeyMap may have already been marked, so we need to remove the
+     * keys to prevent a use-after-free. */
+    rb_gc_remove_weak(self, (VALUE *)key);
+    return wkmap_free_table_i(key, val, 0);
+}
+
 /*
  *  call-seq:
  *    map.clear -> self
@@ -970,7 +981,7 @@ wkmap_clear(VALUE self)
     struct weakkeymap *w;
     TypedData_Get_Struct(self, struct weakkeymap, &weakkeymap_type, w);
 
-    st_foreach(w->table, wkmap_free_table_i, 0);
+    st_foreach(w->table, wkmap_clear_i, (st_data_t)self);
     st_clear(w->table);
 
     return self;