Fix use-after-free for WeakKeyMap

[Bug #20688]

We cannot free the key before the ST_DELETE because it could hash the
key which would read the key and would cause a use-after-free. Instead,
we store the key and free it on the next iteration.
This commit is contained in:
Peter Zhu 2024-08-21 09:42:30 -04:00
parent df9a6aa943
commit 56a34b5af5
Notes: git 2024-08-22 14:02:15 +00:00

View File

@ -612,18 +612,24 @@ struct weakkeymap {
}; };
static int static int
wkmap_mark_table_i(st_data_t key, st_data_t val_obj, st_data_t _) wkmap_mark_table_i(st_data_t key, st_data_t val_obj, st_data_t data)
{ {
VALUE key_obj = *(VALUE *)key; VALUE **dead_entry = (VALUE **)data;
if (dead_entry != NULL) {
ruby_sized_xfree(*dead_entry, sizeof(VALUE));
*dead_entry = NULL;
}
if (wmap_live_p(key_obj)) { VALUE *key_ptr = (VALUE *)key;
rb_gc_mark_weak((VALUE *)key);
if (wmap_live_p(*key_ptr)) {
rb_gc_mark_weak(key_ptr);
rb_gc_mark_movable((VALUE)val_obj); rb_gc_mark_movable((VALUE)val_obj);
return ST_CONTINUE; return ST_CONTINUE;
} }
else { else {
ruby_sized_xfree((VALUE *)key, sizeof(VALUE)); *dead_entry = key_ptr;
return ST_DELETE; return ST_DELETE;
} }
@ -634,7 +640,11 @@ wkmap_mark(void *ptr)
{ {
struct weakkeymap *w = ptr; struct weakkeymap *w = ptr;
if (w->table) { if (w->table) {
st_foreach(w->table, wkmap_mark_table_i, (st_data_t)0); VALUE *dead_entry = NULL;
st_foreach(w->table, wkmap_mark_table_i, (st_data_t)&dead_entry);
if (dead_entry != NULL) {
ruby_sized_xfree(dead_entry, sizeof(VALUE));
}
} }
} }
@ -668,22 +678,28 @@ wkmap_memsize(const void *ptr)
} }
static int static int
wkmap_compact_table_i(st_data_t key, st_data_t val_obj, st_data_t _data, int _error) wkmap_compact_table_i(st_data_t key, st_data_t val_obj, st_data_t data, int _error)
{ {
VALUE key_obj = *(VALUE *)key; VALUE **dead_entry = (VALUE **)data;
if (dead_entry != NULL) {
ruby_sized_xfree(*dead_entry, sizeof(VALUE));
*dead_entry = NULL;
}
if (wmap_live_p(key_obj)) { VALUE *key_ptr = (VALUE *)key;
if (key_obj != rb_gc_location(key_obj) || val_obj != rb_gc_location(val_obj)) {
if (wmap_live_p(*key_ptr)) {
if (*key_ptr != rb_gc_location(*key_ptr) || val_obj != rb_gc_location(val_obj)) {
return ST_REPLACE; return ST_REPLACE;
} }
return ST_CONTINUE;
} }
else { else {
ruby_sized_xfree((VALUE *)key, sizeof(VALUE)); *dead_entry = key_ptr;
return ST_DELETE; return ST_DELETE;
} }
return ST_CONTINUE;
} }
static int static int
@ -703,7 +719,11 @@ wkmap_compact(void *ptr)
struct weakkeymap *w = ptr; struct weakkeymap *w = ptr;
if (w->table) { if (w->table) {
st_foreach_with_replace(w->table, wkmap_compact_table_i, wkmap_compact_table_replace, (st_data_t)0); VALUE *dead_entry = NULL;
st_foreach_with_replace(w->table, wkmap_compact_table_i, wkmap_compact_table_replace, (st_data_t)&dead_entry);
if (dead_entry != NULL) {
ruby_sized_xfree(dead_entry, sizeof(VALUE));
}
} }
} }