diff --git a/test/ruby/test_weakkeymap.rb b/test/ruby/test_weakkeymap.rb index 6b3ffbb81f..91c1538076 100644 --- a/test/ruby/test_weakkeymap.rb +++ b/test/ruby/test_weakkeymap.rb @@ -61,6 +61,20 @@ class TestWeakKeyMap < Test::Unit::TestCase 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 diff --git a/weakmap.c b/weakmap.c index ff205ff5f0..a9d0a6a12b 100644 --- a/weakmap.c +++ b/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;