From 4d000a87e0cf14cdb23c9bfb13d5de2ab8c398bd Mon Sep 17 00:00:00 2001
From: Stephen Seo <seo.disparate@gmail.com>
Date: Tue, 4 Jun 2024 15:40:57 +0900
Subject: [PATCH] Use pthreads to mutex-lock malloc/calloc/free

---
 CMakeLists.txt          |  1 +
 src/another_memcheck.cc | 31 +++++++++++++++++++++++++++++--
 src/another_memcheck.h  |  6 ++++++
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 471c9ef..2f8a1cd 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -28,3 +28,4 @@ set_target_properties(AnotherMemCheck PROPERTIES VERSION ${AnotherMemCheck_VERSI
                                                  SOVERSION ${AnotherMemCheck_SOVERSION})
 
 target_compile_features(AnotherMemCheck PUBLIC cxx_std_20)
+target_link_libraries(AnotherMemCheck PUBLIC pthread)
diff --git a/src/another_memcheck.cc b/src/another_memcheck.cc
index 0df3595..5286495 100644
--- a/src/another_memcheck.cc
+++ b/src/another_memcheck.cc
@@ -50,7 +50,7 @@ namespace SC_AM_Internal {
     return false;
   }
 
-  Stats::Stats() : malloced_list_head(nullptr), malloced_list_tail(nullptr) {
+  Stats::Stats() : malloced_list_head(nullptr), malloced_list_tail(nullptr), pthread_mutex{.__align=0} {
   }
 
   void Stats::initialize() {
@@ -63,13 +63,26 @@ namespace SC_AM_Internal {
     malloced_list_head->data = nullptr;
     malloced_list_tail->data = nullptr;
 
+    pthread_mutex_init(&pthread_mutex, nullptr);
+
     on_exit([] ([[maybe_unused]] int status, void *ptr) {
-      const Stats *stats = reinterpret_cast<const Stats*>(ptr);
+      Stats *stats = reinterpret_cast<Stats*>(ptr);
       stats->print_status();
+      stats->cleanup();
     }, this);
   }
 
+  void Stats::cleanup() {
+    pthread_mutex_destroy(&pthread_mutex);
+    // TODO maybe cleanup list, but it is the end of the program.
+  }
+
   void *Stats::do_malloc(std::size_t size) {
+    if(int ret = pthread_mutex_lock(&pthread_mutex); ret == EINVAL) {
+      std::clog << "ERROR: pthread mutex not properly initialized!\n";
+      return nullptr;
+    }
+
     void *address = real_malloc(size);
     if (address != nullptr) {
       Malloced *data = reinterpret_cast<Malloced*>(real_malloc(sizeof(Malloced)));
@@ -78,10 +91,16 @@ namespace SC_AM_Internal {
       ListNode::add_to_list(malloced_list_tail, data);
     }
 
+    pthread_mutex_unlock(&pthread_mutex);
     return address;
   }
 
   void *Stats::do_calloc(std::size_t n, std::size_t size) {
+    if(int ret = pthread_mutex_lock(&pthread_mutex); ret == EINVAL) {
+      std::clog << "ERROR: pthread mutex not properly initialized!\n";
+      return nullptr;
+    }
+
     void *address = real_calloc(n, size);
     if (address != nullptr) {
       Malloced *data = reinterpret_cast<Malloced*>(real_malloc(sizeof(Malloced)));
@@ -90,15 +109,23 @@ namespace SC_AM_Internal {
       ListNode::add_to_list(malloced_list_tail, data);
     }
 
+    pthread_mutex_unlock(&pthread_mutex);
     return address;
   }
 
   void Stats::do_free(void *ptr) {
+    if(int ret = pthread_mutex_lock(&pthread_mutex); ret == EINVAL) {
+      std::clog << "ERROR: pthread mutex not properly initialized!\n";
+      return;
+    }
+
     if (ptr) {
       if(!ListNode::remove_from_list(malloced_list_head, ptr)) {
         std::clog << "WARNING: Attempted free of unknown memory location!\n";
       }
     }
+
+    pthread_mutex_unlock(&pthread_mutex);
   }
 
   void Stats::print_status() const {
diff --git a/src/another_memcheck.h b/src/another_memcheck.h
index aab76d4..a28d400 100644
--- a/src/another_memcheck.h
+++ b/src/another_memcheck.h
@@ -1,6 +1,10 @@
 #ifndef SEODISPARATE_COM_ANOTHER_MEMCHECK_H
 #define SEODISPARATE_COM_ANOTHER_MEMCHECK_H
 
+extern "C" {
+#include <pthread.h>
+}
+
 namespace SC_AM_Internal {
   // Forward declaration.
   struct Stats;
@@ -42,8 +46,10 @@ namespace SC_AM_Internal {
 
     ListNode *malloced_list_head;
     ListNode *malloced_list_tail;
+    pthread_mutex_t pthread_mutex;
 
     void initialize();
+    void cleanup();
 
     void *do_malloc(std::size_t size);
     void *do_calloc(std::size_t n, std::size_t size);