LibGC: Add live HeapBlock registry to fix weak container UAF

Maintain a HashTable<HeapBlock*> of live heap blocks in the Heap,
updated on block creation and destruction.

Weak containers (WeakMap, WeakSet, WeakRef, FinalizationRegistry)
now check block liveness before accessing cell memory in their
remove_dead_cells() methods. This prevents use-after-free when
blocks have been freed during incremental sweeping.

Libraries/LibGC/CellAllocator.cpp

@@ -56,6 +56,7 @@ Cell* CellAllocator::allocate_cell(Heap& heap)
 void CellAllocator::block_did_become_empty(Badge<Heap>, HeapBlock& block)
 {
     block.m_list_node.remove();
+    block.heap().m_live_heap_blocks.remove(&block);
     // NOTE: HeapBlocks are managed by the BlockAllocator, so we don't want to `delete` the block here.
     block.~HeapBlock();
     m_block_allocator.deallocate_block(&block);

Libraries/LibGC/Heap.h

@@ -8,6 +8,7 @@
 
 #include <AK/Badge.h>
 #include <AK/Function.h>
+#include <AK/HashTable.h>
 #include <AK/Noncopyable.h>
 #include <AK/NonnullOwnPtr.h>
 #include <AK/RefPtr.h>
@@ -98,6 +99,8 @@ public:
 
     void sweep_block(HeapBlock&);
 
+    bool is_live_heap_block(HeapBlock* block) const { return m_live_heap_blocks.contains(block); }
+
     void enqueue_post_gc_task(AK::Function<void()>);
 
     WeakImpl* create_weak_impl(void*);
@@ -106,6 +109,8 @@ public:
     void did_free_external_memory(size_t);
 
 private:
+    friend class CellAllocator;
+    friend class HeapBlock;
     friend class MarkingVisitor;
     friend class GraphConstructorVisitor;
     friend class DeferGC;
@@ -181,6 +186,8 @@ private:
     Vector<AK::Function<void()>> m_post_gc_tasks;
     Vector<AK::Function<void()>> m_sweep_callbacks;
 
+    HashTable<HeapBlock*> m_live_heap_blocks;
+
     WeakBlock::List m_usable_weak_blocks;
     WeakBlock::List m_full_weak_blocks;
 

Libraries/LibGC/HeapBlock.cpp

@@ -9,6 +9,7 @@
 #include <AK/Platform.h>
 #include <LibGC/CellAllocator.h>
 #include <LibGC/Forward.h>
+#include <LibGC/Heap.h>
 #include <LibGC/HeapBlock.h>
 
 #ifdef HAS_ADDRESS_SANITIZER
@@ -22,6 +23,7 @@ NonnullOwnPtr<HeapBlock> HeapBlock::create_with_cell_size(Heap& heap, CellAlloca
     char const* name = nullptr;
     auto* block = static_cast<HeapBlock*>(cell_allocator.block_allocator().allocate_block(name));
     new (block) HeapBlock(heap, cell_allocator, cell_size, overrides_must_survive_garbage_collection, overrides_finalize);
+    heap.m_live_heap_blocks.set(block);
     return NonnullOwnPtr<HeapBlock>(NonnullOwnPtr<HeapBlock>::Adopt, *block);
 }
 

Libraries/LibJS/Runtime/FinalizationRegistry.cpp

@@ -4,6 +4,7 @@
  * SPDX-License-Identifier: BSD-2-Clause
  */
 
+#include <LibGC/HeapBlock.h>
 #include <LibJS/Runtime/AbstractOperations.h>
 #include <LibJS/Runtime/FinalizationRegistry.h>
 
@@ -54,7 +55,10 @@ void FinalizationRegistry::remove_dead_cells(Badge<GC::Heap>)
 {
     auto any_cells_were_removed = false;
     for (auto& record : m_records) {
-        if (!record.target || (record.target->state() == Cell::State::Live && record.target->is_marked()))
+        if (!record.target)
+            continue;
+        auto* block = GC::HeapBlock::from_cell(record.target);
+        if (heap().is_live_heap_block(block) && record.target->state() == Cell::State::Live && record.target->is_marked())
             continue;
         record.target = nullptr;
         any_cells_were_removed = true;

Libraries/LibJS/Runtime/WeakMap.cpp

@@ -4,6 +4,7 @@
  * SPDX-License-Identifier: BSD-2-Clause
  */
 
+#include <LibGC/HeapBlock.h>
 #include <LibJS/Runtime/ExternalMemory.h>
 #include <LibJS/Runtime/WeakMap.h>
 
@@ -52,8 +53,9 @@ bool WeakMap::weak_map_remove(GC::Ptr<Cell> key)
 
 void WeakMap::remove_dead_cells(Badge<GC::Heap>)
 {
-    m_values.remove_all_matching([](Cell* key, Value) {
-        return key->state() != Cell::State::Live || !key->is_marked();
+    m_values.remove_all_matching([this](Cell* key, Value) {
+        auto* block = GC::HeapBlock::from_cell(key);
+        return !heap().is_live_heap_block(block) || key->state() != Cell::State::Live || !key->is_marked();
     });
 }
 

Libraries/LibJS/Runtime/WeakRef.cpp

@@ -4,6 +4,7 @@
  * SPDX-License-Identifier: BSD-2-Clause
  */
 
+#include <LibGC/HeapBlock.h>
 #include <LibJS/Runtime/WeakRef.h>
 
 namespace JS {
@@ -38,7 +39,13 @@ WeakRef::WeakRef(Symbol& value, Object& prototype)
 
 void WeakRef::remove_dead_cells(Badge<GC::Heap>)
 {
-    if (m_value.visit([](Cell* cell) -> bool { return cell->state() == Cell::State::Live && cell->is_marked(); }, [](Empty) -> bool { return true; }))
+    auto is_alive = m_value.visit(
+        [this](Cell* cell) -> bool {
+            auto* block = GC::HeapBlock::from_cell(cell);
+            return heap().is_live_heap_block(block) && cell->state() == Cell::State::Live && cell->is_marked();
+        },
+        [](Empty) -> bool { return true; });
+    if (is_alive)
         return;
 
     m_value = Empty {};

Libraries/LibJS/Runtime/WeakSet.cpp

@@ -4,6 +4,7 @@
  * SPDX-License-Identifier: BSD-2-Clause
  */
 
+#include <LibGC/HeapBlock.h>
 #include <LibJS/Runtime/ExternalMemory.h>
 #include <LibJS/Runtime/WeakSet.h>
 
@@ -45,8 +46,9 @@ bool WeakSet::weak_set_remove(GC::Ptr<Cell> value)
 
 void WeakSet::remove_dead_cells(Badge<GC::Heap>)
 {
-    m_values.remove_all_matching([](Cell* cell) {
-        return cell->state() != Cell::State::Live || !cell->is_marked();
+    m_values.remove_all_matching([this](Cell* cell) {
+        auto* block = GC::HeapBlock::from_cell(cell);
+        return !heap().is_live_heap_block(block) || cell->state() != Cell::State::Live || !cell->is_marked();
     });
 }