2 files changed, 22 insertions, 1 deletions
diff --git a/test/ruby/test_weakmap.rb b/test/ruby/test_weakmap.rb
index 9f8528050e..a2904776bc 100644
--- a/test/ruby/test_weakmap.rb
+++ b/test/ruby/test_weakmap.rb
@@ -258,4 +258,11 @@ class TestWeakMap < Test::Unit::TestCase
 
     assert_equal b, @wm[1]
   end
+
+  def test_use_after_free_bug_20688
+    assert_normal_exit(<<~RUBY)
+      weakmap = ObjectSpace::WeakMap.new
+      10_000.times { weakmap[Object.new] = Object.new }
+    RUBY
+  end
 end
diff --git a/weakmap.c b/weakmap.c
index fb2fd22875..adbcaa8882 100644
--- a/weakmap.c
+++ b/weakmap.c
@@ -43,6 +43,8 @@ wmap_live_p(VALUE obj)
 struct wmap_foreach_data {
     int (*func)(struct weakmap_entry *, st_data_t);
     st_data_t arg;
+
+    struct weakmap_entry *dead_entry;
 };
 
 static int
@@ -50,6 +52,11 @@ wmap_foreach_i(st_data_t key, st_data_t val, st_data_t arg)
 {
     struct wmap_foreach_data *data = (struct wmap_foreach_data *)arg;
 
+    if (data->dead_entry != NULL) {
+        ruby_sized_xfree(data->dead_entry, sizeof(struct weakmap_entry));
+        data->dead_entry = NULL;
+    }
+
     struct weakmap_entry *entry = (struct weakmap_entry *)key;
     RUBY_ASSERT(&entry->val == (VALUE *)val);
 
@@ -65,7 +72,11 @@ wmap_foreach_i(st_data_t key, st_data_t val, st_data_t arg)
         return ret;
     }
     else {
-        ruby_sized_xfree(entry, sizeof(struct weakmap_entry));
+        /* We cannot free the weakmap_entry here because the ST_DELETE could
+         * hash the key which would read the weakmap_entry and would cause a
+         * use-after-free. Instead, we store this entry and free it on the next
+         * iteration. */
+        data->dead_entry = entry;
 
         return ST_DELETE;
     }
@@ -77,9 +88,12 @@ wmap_foreach(struct weakmap *w, int (*func)(struct weakmap_entry *, st_data_t),
     struct wmap_foreach_data foreach_data = {
         .func = func,
         .arg = arg,
+        .dead_entry = NULL,
     };
 
     st_foreach(w->table, wmap_foreach_i, (st_data_t)&foreach_data);
+
+    ruby_sized_xfree(foreach_data.dead_entry, sizeof(struct weakmap_entry));
 }
 
 static int