From 7777e3f2dd5ee97695ddaad7519248eb2292516f Mon Sep 17 00:00:00 2001
From: Milian Wolff <mail@milianw.de>
Date: Sat, 6 Dec 2014 13:35:34 +0100
Subject: [PATCH] Cleanup code a bit.

---
 heaptrack_inject.cpp | 44 ++++++++++++++++++++++++++++++--------------
 1 file changed, 30 insertions(+), 14 deletions(-)

diff --git a/heaptrack_inject.cpp b/heaptrack_inject.cpp
index 2187355..925df76 100644
--- a/heaptrack_inject.cpp
+++ b/heaptrack_inject.cpp
@@ -31,22 +31,38 @@
 
 namespace {
 
-auto original_malloc = &malloc;
-auto original_free = &free;
+namespace hooks {
 
-void* overwrite_malloc(size_t size)
+struct malloc
 {
-    auto ptr = original_malloc(size);
-    heaptrack_malloc(ptr, size);
-    return ptr;
-}
+    static constexpr auto name = "malloc";
+    static constexpr auto original = &::malloc;
+    using Signature = void* (*) (size_t);
+
+    static void* hook(size_t size)
+    {
+        auto ptr = original(size);
+        heaptrack_malloc(ptr, size);
+        return ptr;
+    }
+};
 
-void overwrite_free(void* ptr)
+struct free
 {
-    heaptrack_free(ptr);
-    original_free(ptr);
+    static constexpr auto name = "free";
+    static constexpr auto original = &::free;
+    using Signature = void (*) (void*);
+
+    static void hook(void *ptr)
+    {
+        heaptrack_free(ptr);
+        original(ptr);
+    }
+};
+
 }
 
+
 template<typename T, ElfW(Sxword) AddrTag, ElfW(Sxword) SizeTag>
 struct elftable
 {
@@ -85,10 +101,10 @@ void try_overwrite_symbols(const ElfW(Dyn) *dyn, const ElfW(Addr) base)
         auto relsymidx = ELF64_R_SYM(rela->r_info);
         const char *relsymname = strings.table + symbols.table[relsymidx].st_name;
         auto addr = rela->r_offset + base;
-        if (!strcmp("malloc", relsymname)) {
-            *reinterpret_cast<void*(**)(size_t)>(addr) = &overwrite_malloc;
-        } else if (!strcmp("free", relsymname)) {
-            *reinterpret_cast<void(**)(void*)>(addr) = &overwrite_free;
+        if (!strcmp(hooks::malloc::name, relsymname)) {
+            *reinterpret_cast<hooks::malloc::Signature*>(addr) = &hooks::malloc::hook;
+        } else if (!strcmp(hooks::free::name, relsymname)) {
+            *reinterpret_cast<hooks::free::Signature*>(addr) = &hooks::free::hook;
         }
     }
 }
-- 
2.7.4