From ac9447eba55e1d45a6a15446d3277921d32fd437 Mon Sep 17 00:00:00 2001
From: Squatch Stack <321440025+squatchlr@users.noreply.github.com>
Date: Fri, 9 Oct 2026 13:20:18 -0500
Subject: [PATCH] Count allocations through the sanitizer hooks when built with
 a sanitizer

Under ASan, TSan or LSan the runtime owns malloc, and replacing it through
dlsym(RTLD_NEXT) crashes run_tests before main (ASan: segfault, no output;
LSan: duplicate operator new at link time). With a sanitizer,
allocation_tracking now counts through __sanitizer_install_malloc_and_free_hooks
instead, so the real-time-safety tests keep their meaning there. Builds
without a sanitizer are unchanged.

Reproduce: configure with -DCMAKE_CXX_FLAGS=-fsanitize=address (or thread,
leak) and the matching linker flag; on main run_tests exits 139 with no
output, on this branch it passes.
---
 tools/test/allocation_tracking.cpp | 35 ++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/tools/test/allocation_tracking.cpp b/tools/test/allocation_tracking.cpp
index 983ed11d..9d92c17a 100644
--- a/tools/test/allocation_tracking.cpp
+++ b/tools/test/allocation_tracking.cpp
@@ -4,6 +4,18 @@
 
 #include "allocation_tracking.h"
 
+// Under a sanitizer (ASan, TSan, LSan) the runtime owns malloc, and replacing it here through dlsym(RTLD_NEXT) kills
+// run_tests before main. There we count through the sanitizer's own allocation hooks instead, so the allocation tests
+// still mean the same thing (only "zero" and "more than zero" are ever checked).
+#if defined(__SANITIZE_ADDRESS__) || defined(__SANITIZE_THREAD__)
+  #define NAM_TEST_UNDER_SANITIZER 1
+#elif defined(__has_feature)
+  #if __has_feature(address_sanitizer) || __has_feature(thread_sanitizer) || __has_feature(leak_sanitizer)             \
+    || __has_feature(memory_sanitizer)
+    #define NAM_TEST_UNDER_SANITIZER 1
+  #endif
+#endif
+
 // Allocation tracking globals - definitions
 namespace allocation_tracking
 {
@@ -17,6 +29,27 @@ void (*original_free)(void*) = nullptr;
 void* (*original_realloc)(void*, size_t) = nullptr;
 } // namespace allocation_tracking
 
+#ifdef NAM_TEST_UNDER_SANITIZER
+  #include <sanitizer/allocator_interface.h>
+
+namespace
+{
+void on_malloc(const volatile void* ptr, size_t)
+{
+  if (allocation_tracking::g_tracking_enabled && ptr != nullptr)
+    ++allocation_tracking::g_allocation_count;
+}
+
+void on_free(const volatile void* ptr)
+{
+  if (allocation_tracking::g_tracking_enabled && ptr != nullptr)
+    ++allocation_tracking::g_deallocation_count;
+}
+
+const int hooks_installed = __sanitizer_install_malloc_and_free_hooks(on_malloc, on_free);
+} // namespace
+
+#else
 // Override malloc/free to track Eigen allocations (Eigen uses malloc directly)
 extern "C" {
 void* malloc(size_t size)
@@ -88,3 +121,5 @@ void operator delete[](void* ptr) noexcept
     ++allocation_tracking::g_deallocation_count;
   std::free(ptr);
 }
+
+#endif // NAM_TEST_UNDER_SANITIZER