Introduce rb_gc_remove_weak

If we're during incremental marking, then Ruby code can execute that
deallocates certain memory buffers that have been called with
rb_gc_mark_weak, which can cause use-after-free bugs.

gc.c

@@ -6924,6 +6924,27 @@ rb_gc_mark_weak(VALUE *ptr)
     objspace->profile.weak_references_count++;
 }
 
+void
+rb_gc_remove_weak(VALUE parent_obj, VALUE *ptr)
+{
+    rb_objspace_t *objspace = &rb_objspace;
+
+    /* If we're not incremental marking, then the state of the objects can't
+     * change so we don't need to do anything. */
+    if (!is_incremental_marking(objspace)) return;
+    /* If parent_obj has not been marked, then ptr has not yet been marked
+     * weak, so we don't need to do anything. */
+    if (!RVALUE_MARKED(parent_obj)) return;
+
+    VALUE **ptr_ptr;
+    rb_darray_foreach(objspace->weak_references, i, ptr_ptr) {
+        if (*ptr_ptr == ptr) {
+            *ptr_ptr = NULL;
+            break;
+        }
+    }
+}
+
 /* CAUTION: THIS FUNCTION ENABLE *ONLY BEFORE* SWEEPING.
  * This function is only for GC_END_MARK timing.
  */
@@ -8151,6 +8172,8 @@ gc_update_weak_references(rb_objspace_t *objspace)
     size_t retained_weak_references_count = 0;
     VALUE **ptr_ptr;
     rb_darray_foreach(objspace->weak_references, i, ptr_ptr) {
+        if (!ptr_ptr) continue;
+
         VALUE obj = **ptr_ptr;
 
         if (RB_SPECIAL_CONST_P(obj)) continue;

internal/gc.h

@@ -236,6 +236,7 @@ VALUE rb_define_finalizer_no_check(VALUE obj, VALUE block);
 void rb_gc_mark_and_move(VALUE *ptr);
 
 void rb_gc_mark_weak(VALUE *ptr);
+void rb_gc_remove_weak(VALUE parent_obj, VALUE *ptr);
 
 #define rb_gc_mark_and_move_ptr(ptr) do { \
     VALUE _obj = (VALUE)*(ptr); \

weakmap.c

@@ -380,7 +380,7 @@ wmap_aset_replace(st_data_t *key, st_data_t *val, st_data_t new_key_ptr, int exi
     VALUE new_val = *(((VALUE *)new_key_ptr) + 1);
 
     if (existing) {
-        assert(orig_pair[0] == *(VALUE *)new_key);
+        assert(*(VALUE *)*key == new_key);
     }
     else {
         VALUE *pair = xmalloc(sizeof(VALUE) * 2);
@@ -450,6 +450,9 @@ wmap_delete(VALUE self, VALUE key)
     if (st_delete(w->table, &orig_key_data, &orig_val_data)) {
         VALUE orig_val = *(VALUE *)orig_val_data;
 
+        rb_gc_remove_weak(self, (VALUE *)orig_key_data);
+        rb_gc_remove_weak(self, (VALUE *)orig_val_data);
+
         wmap_free_entry((VALUE *)orig_key_data, (VALUE *)orig_val_data);
 
         if (wmap_live_p(orig_val)) {
@@ -776,6 +779,8 @@ wkmap_delete(VALUE self, VALUE key)
     if (st_delete(w->table, &orig_key_data, &orig_val_data)) {
         VALUE orig_val = (VALUE)orig_val_data;
 
+        rb_gc_remove_weak(self, (VALUE *)orig_key_data);
+
         ruby_sized_xfree((VALUE *)orig_key_data, sizeof(VALUE));
 
         return orig_val;