LibJS: Instrument HeapBlock cell allocation for ASAN

Mark the entirety of a heap block's storage poisoned at construction.
Unpoison all of a Cell's memory before allocating it, and re-poison as
much as possible on deallocation. Unfortunately, the entirety of the
FreelistEntry must be kept unpoisoned in order for reallocation to work
correctly.

Decreasing the size of FreelistEntry or adding a larger redzone to Cells
would make the instrumentation even better.

2 files changed, 31 insertions, 4 deletions

diff --git a/Userland/Libraries/LibJS/Heap/HeapBlock.cpp b/Userland/Libraries/LibJS/Heap/HeapBlock.cpp
index c2ae0b3470..4999be117f 100644
--- a/Userland/Libraries/LibJS/Heap/HeapBlock.cpp
+++ b/Userland/Libraries/LibJS/Heap/HeapBlock.cpp
@@ -6,11 +6,16 @@
 
 #include <AK/Assertions.h>
 #include <AK/NonnullOwnPtr.h>
+#include <AK/Platform.h>
 #include <LibJS/Heap/Heap.h>
 #include <LibJS/Heap/HeapBlock.h>
 #include <stdio.h>
 #include <sys/mman.h>
 
+#ifdef HAS_ADDRESS_SANITIZER
+#    include <sanitizer/asan_interface.h>
+#endif
+
 namespace JS {
 
 NonnullOwnPtr<HeapBlock> HeapBlock::create_with_cell_size(Heap& heap, size_t cell_size)
@@ -27,6 +32,7 @@ HeapBlock::HeapBlock(Heap& heap, size_t cell_size)
     , m_cell_size(cell_size)
 {
     VERIFY(cell_size >= sizeof(FreelistEntry));
+    ASAN_POISON_MEMORY_REGION(m_storage, block_size);
 }
 
 void HeapBlock::deallocate(Cell* cell)
@@ -35,11 +41,22 @@ void HeapBlock::deallocate(Cell* cell)
     VERIFY(!m_freelist || is_valid_cell_pointer(m_freelist));
     VERIFY(cell->state() == Cell::State::Live);
     VERIFY(!cell->is_marked());
+
     cell->~Cell();
     auto* freelist_entry = new (cell) FreelistEntry();
     freelist_entry->set_state(Cell::State::Dead);
     freelist_entry->next = m_freelist;
     m_freelist = freelist_entry;
+
+#ifdef HAS_ADDRESS_SANITIZER
+    auto dword_after_freelist = round_up_to_power_of_two(reinterpret_cast<uintptr_t>(freelist_entry) + sizeof(FreelistEntry), 8);
+    VERIFY((dword_after_freelist - reinterpret_cast<uintptr_t>(freelist_entry)) <= m_cell_size);
+    VERIFY(m_cell_size >= sizeof(FreelistEntry));
+    // We can't poision the cell tracking data, nor the FreeListEntry's vtable or next pointer
+    // This means there's sizeof(FreelistEntry) data at the front of each cell that is always read/write
+    // On x86_64, this ends up being 24 bytes due to the size of the FreeListEntry's vtable, while on x86, it's only 12 bytes.
+    ASAN_POISON_MEMORY_REGION(reinterpret_cast<void*>(dword_after_freelist), m_cell_size - sizeof(FreelistEntry));
+#endif
 }
 
 }
diff --git a/Userland/Libraries/LibJS/Heap/HeapBlock.h b/Userland/Libraries/LibJS/Heap/HeapBlock.h
index 6b2c1eff12..03b202317a 100644
--- a/Userland/Libraries/LibJS/Heap/HeapBlock.h
+++ b/Userland/Libraries/LibJS/Heap/HeapBlock.h
@@ -7,10 +7,15 @@
 #pragma once
 
 #include <AK/IntrusiveList.h>
+#include <AK/Platform.h>
 #include <AK/Types.h>
 #include <LibJS/Forward.h>
 #include <LibJS/Heap/Cell.h>
 
+#ifdef HAS_ADDRESS_SANITIZER
+#    include <sanitizer/asan_interface.h>
+#endif
+
 namespace JS {
 
 class HeapBlock {
@@ -27,13 +32,18 @@ public:
 
     ALWAYS_INLINE Cell* allocate()
     {
+        Cell* allocated_cell = nullptr;
         if (m_freelist) {
             VERIFY(is_valid_cell_pointer(m_freelist));
-            return exchange(m_freelist, m_freelist->next);
+            allocated_cell = exchange(m_freelist, m_freelist->next);
+        } else if (has_lazy_freelist()) {
+            allocated_cell = cell(m_next_lazy_freelist_index++);
+        }
+
+        if (allocated_cell) {
+            ASAN_UNPOISON_MEMORY_REGION(allocated_cell, m_cell_size);
         }
-        if (has_lazy_freelist())
-            return cell(m_next_lazy_freelist_index++);
-        return nullptr;
+        return allocated_cell;
     }
 
     void deallocate(Cell*);