gh-583 Heap arena canary guards

src/internal_modules/roc_core/heap_arena.cpp

@@ -16,39 +16,53 @@
 namespace roc {
 namespace core {
 
-int HeapArena::enable_leak_detection_ = false;
+size_t HeapArena::flags_ = DefaultHeapArenaFlags;
 
 HeapArena::HeapArena()
-    : num_allocations_(0) {
+    : num_allocations_(0)
+    , num_guard_failures_(0) {
 }
 
 HeapArena::~HeapArena() {
     if (num_allocations_ != 0) {
-        if (AtomicOps::load_seq_cst(enable_leak_detection_)) {
+        if (AtomicOps::load_seq_cst(flags_) & HeapArenaFlag_EnableLeakDetection) {
             roc_panic("heap arena: detected leak(s): %d objects was not freed",
                       (int)num_allocations_);
         }
     }
 }
 
-void HeapArena::enable_leak_detection() {
-    AtomicOps::store_seq_cst(enable_leak_detection_, true);
+void HeapArena::set_flags(size_t flags) {
+    AtomicOps::store_seq_cst(flags_, flags);
 }
 
 size_t HeapArena::num_allocations() const {
     return (size_t)num_allocations_;
 }
 
+size_t HeapArena::num_guard_failures() const {
+    return num_guard_failures_;
+}
+
 void* HeapArena::allocate(size_t size) {
     num_allocations_++;
 
-    Chunk* chunk = (Chunk*)malloc(sizeof(Chunk) + size);
+    size_t chunk_size =
+        sizeof(ChunkHeader) + sizeof(ChunkCanary) + size + sizeof(ChunkCanary);
 
-    chunk->size = size;
+    ChunkHeader* chunk = (ChunkHeader*)malloc(chunk_size);
 
-    MemoryOps::poison_before_use(chunk->data, size);
+    char* canary_before = (char*)chunk->data;
+    char* memory = (char*)chunk->data + sizeof(ChunkCanary);
+    char* canary_after = (char*)chunk->data + sizeof(ChunkCanary) + size;
+
+    MemoryOps::prepare_canary(canary_before, sizeof(ChunkCanary));
+    MemoryOps::poison_before_use(memory, size);
+    MemoryOps::prepare_canary(canary_after, sizeof(ChunkCanary));
+
+    chunk->size = size;
 
-    return chunk->data;
+    return memory;
 }
 
 void HeapArena::deallocate(void* ptr) {
@@ -62,9 +76,29 @@ void HeapArena::deallocate(void* ptr) {
         roc_panic("heap arena: unpaired deallocate");
     }
 
-    Chunk* chunk = ROC_CONTAINER_OF(ptr, Chunk, data);
+    ChunkHeader* chunk =
+        ROC_CONTAINER_OF((char*)ptr - sizeof(ChunkCanary), ChunkHeader, data);
+
+    size_t size = chunk->size;
+
+    char* canary_before = (char*)chunk->data;
+    char* memory = (char*)chunk->data + sizeof(ChunkCanary);
+    char* canary_after = (char*)chunk->data + sizeof(ChunkCanary) + size;
+
+    const bool canary_before_ok =
+        MemoryOps::check_canary(canary_before, sizeof(ChunkCanary));
+    const bool canary_after_ok =
+        MemoryOps::check_canary(canary_after, sizeof(ChunkCanary));
+
+    if (!canary_before_ok || !canary_after_ok) {
+        num_guard_failures_++;
+        if (AtomicOps::load_seq_cst(flags_) & HeapArenaFlag_EnableGuards) {
+            roc_panic("heap arena: detected memory violation: ok_before=%d ok_after=%d",
+                      (int)canary_before_ok, (int)canary_after_ok);
+        }
+    }
 
-    MemoryOps::poison_after_use(chunk->data, chunk->size);
+    MemoryOps::poison_after_use(memory, chunk->size);
 
     free(chunk);
 }

src/internal_modules/roc_core/heap_arena.h

@@ -20,6 +20,17 @@
 namespace roc {
 namespace core {
 
+//! Heap arena flags.
+enum HeapArenaFlags {
+    //! Enable panic if leaks detected in arena destructor.
+    HeapArenaFlag_EnableLeakDetection = (1 << 0),
+    //! Enable panic if memory violation detected when deallocating chunk.
+    HeapArenaFlag_EnableGuards = (1 << 1),
+};
+
+//! Default heap arena flags.
+enum { DefaultHeapArenaFlags = (HeapArenaFlag_EnableGuards) };
+
 //! Heap arena implementation.
 //!
 //! Uses malloc() and free().
@@ -30,30 +41,41 @@ namespace core {
 //! The memory is always maximum aligned. Thread-safe.
 class HeapArena : public IArena, public NonCopyable<> {
 public:
+    //! Initialize.
     HeapArena();
     ~HeapArena();
 
-    //! Enable panic on leak in destructor, for all instances.
-    static void enable_leak_detection();
+    //! Set flags, for all instances.
+    //!
+    //! @b Parameters
+    //! - @p flags defines options to modify behaviour as indicated in HeapArenaFlags
+    static void set_flags(size_t flags);
 
     //! Get number of allocated blocks.
     size_t num_allocations() const;
 
+    //! Get number of guard failures.
+    size_t num_guard_failures() const;
+
     //! Allocate memory.
     virtual void* allocate(size_t size);
 
     //! Deallocate previously allocated memory.
     virtual void deallocate(void*);
 
 private:
-    struct Chunk {
+    struct ChunkHeader {
         size_t size;
         AlignMax data[];
     };
 
-    static int enable_leak_detection_;
+    typedef AlignMax ChunkCanary;
 
     Atomic<int> num_allocations_;
+
+    static size_t flags_;
+
+    size_t num_guard_failures_;
 };
 
 } // namespace core

src/tests/bench_main.cpp

@@ -15,7 +15,8 @@
 using namespace roc;
 
 int main(int argc, char** argv) {
-    core::HeapArena::enable_leak_detection();
+    core::HeapArena::set_flags(core::DefaultHeapArenaFlags
+                               | core::HeapArenaFlag_EnableLeakDetection);
 
     core::CrashHandler crash_handler;
 

src/tests/roc_core/test_heap_arena.cpp

@@ -0,0 +1,74 @@
+/*
+ * Copyright (c) 2023 Roc Streaming authors
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ */
+
+#include <CppUTest/TestHarness.h>
+
+#include "roc_core/heap_arena.h"
+#include "roc_core/memory_ops.h"
+
+namespace roc {
+namespace core {
+
+// clang-format off
+TEST_GROUP(heap_arena) {
+    void setup() {
+        core::HeapArena::set_flags(core::DefaultHeapArenaFlags
+                                    & ~core::HeapArenaFlag_EnableGuards);
+    }
+    void teardown() {
+        core::HeapArena::set_flags(core::DefaultHeapArenaFlags);
+    }
+};
+// clang-format on
+
+TEST(heap_arena, guard_object) {
+    HeapArena arena;
+    void* pointer = NULL;
+
+    pointer = arena.allocate(127);
+    CHECK(pointer);
+
+    char* data = (char*)pointer;
+    char* before_data = data - 1;
+    char* after_data = data + 127;
+    CHECK(*before_data == MemoryOps::Pattern_Canary);
+    CHECK(*after_data == MemoryOps::Pattern_Canary);
+
+    arena.deallocate(pointer);
+}
+
+TEST(heap_arena, guard_object_violations) {
+    HeapArena arena;
+
+    void* pointers[2] = {};
+
+    pointers[0] = arena.allocate(128);
+    CHECK(pointers[0]);
+
+    pointers[1] = arena.allocate(128);
+    CHECK(pointers[1]);
+
+    {
+        char* data = (char*)pointers[0];
+        data--;
+        *data = 0x00;
+    }
+    arena.deallocate(pointers[0]);
+    CHECK(arena.num_guard_failures() == 1);
+
+    {
+        char* data = (char*)pointers[1];
+        data += 128;
+        *data = 0x00;
+    }
+    arena.deallocate(pointers[1]);
+    CHECK(arena.num_guard_failures() == 2);
+}
+
+} // namespace core
+} // namespace roc

src/tests/test_main.cpp

@@ -18,7 +18,8 @@
 using namespace roc;
 
 int main(int argc, const char** argv) {
-    core::HeapArena::enable_leak_detection();
+    core::HeapArena::set_flags(core::DefaultHeapArenaFlags
+                               | core::HeapArenaFlag_EnableLeakDetection);
 
     core::CrashHandler crash_handler;
 

src/tools/roc_copy/main.cpp

@@ -25,7 +25,8 @@
 using namespace roc;
 
 int main(int argc, char** argv) {
-    core::HeapArena::enable_leak_detection();
+    core::HeapArena::set_flags(core::DefaultHeapArenaFlags
+                               | core::HeapArenaFlag_EnableLeakDetection);
 
     core::CrashHandler crash_handler;
 

src/tools/roc_recv/main.cpp

@@ -31,7 +31,8 @@
 using namespace roc;
 
 int main(int argc, char** argv) {
-    core::HeapArena::enable_leak_detection();
+    core::HeapArena::set_flags(core::DefaultHeapArenaFlags
+                               | core::HeapArenaFlag_EnableLeakDetection);
 
     core::CrashHandler crash_handler;
 

src/tools/roc_send/main.cpp

@@ -31,7 +31,8 @@
 using namespace roc;
 
 int main(int argc, char** argv) {
-    core::HeapArena::enable_leak_detection();
+    core::HeapArena::set_flags(core::DefaultHeapArenaFlags
+                               | core::HeapArenaFlag_EnableLeakDetection);
 
     core::CrashHandler crash_handler;