1 files changed, 124 insertions, 105 deletions

diff --git a/kmtrace/ktrace.c b/kmtrace/ktrace.c
index 8b9eb2f4..e7eee47b 100644
--- a/kmtrace/ktrace.c
+++ b/kmtrace/ktrace.c
@@ -112,27 +112,36 @@ void kuntrace(void);
 
 static void addAllocationToTree(void);
 
-static void tr_freehook __P ((void*, const void*));
-static void* tr_reallochook __P ((void*, size_t,
-									const void*));
-static void* tr_mallochook __P ((size_t, const void*));
-/* Old hook values.  */
-static void (*tr_old_free_hook) __P ((void* ptr, const void*));
-static void* (*tr_old_malloc_hook) __P ((size_t size,
-										   const void*));
-static void* (*tr_old_realloc_hook) __P ((void* ptr,
-											size_t size,
-											const void*));
+// Pointers to real calls
+static void* (*real_malloc_ptr)  (size_t size) = NULL;
+static void* (*real_realloc_ptr) (void *ptr, size_t size) = NULL;
+static void  (*real_free_ptr)    (void *ptr) = NULL;
+
+// This is used to return "allocated" memory until the pointer
+// to the real malloc function becomes known.
+// It seems about 75k of memory are required before that happens,
+// but the buffer provides more space to cater for more if needed
+// in other OS.
+#define MALLOC_INIT_SIZE (1024 * 256)
+static char malloc_init_buf[MALLOC_INIT_SIZE];
+static size_t malloc_init_pos = 0;
+static pthread_mutex_t malloc_init_lock = PTHREAD_MUTEX_INITIALIZER;
 
 static FILE* mallstream;
 static char malloc_trace_buffer[TRACE_BUFFER_SIZE];
 
+static pthread_mutex_t malloc_lock = PTHREAD_MUTEX_INITIALIZER;
+static pthread_mutex_t free_lock = PTHREAD_MUTEX_INITIALIZER;
+static pthread_mutex_t realloc_lock = PTHREAD_MUTEX_INITIALIZER;
 
-/* Address to breakpoint on accesses to... */
-void* mallwatch;
-
-static pthread_mutex_t lock = PTHREAD_MUTEX_INITIALIZER;
+static int tr_malloc_hook_enabled = 0;
+static int tr_realloc_hook_enabled = 0;
+static int tr_free_hook_enabled = 0;
 
+// Hooks active in this thread
+static __thread int tr_malloc_hook_active = 0;
+static __thread int tr_realloc_hook_active = 0;
+static __thread int tr_free_hook_active = 0;
 
 static int bt_size;
 static void *bt[TR_BT_SIZE + 1];
@@ -159,14 +168,15 @@ static char* mallTreeFile = NULL;
 static FILE* mallTreeStream = NULL;
 static long mallThreshold = 2000;
 
-/* This function is called when the block being alloc'd, realloc'd, or
- * freed has an address matching the variable "mallwatch".  In a
- * debugger, set "mallwatch" to the address of interest, then put a
- * breakpoint on tr_break.  */
-void tr_break __P ((void));
-void
-tr_break()
+// Get real function pointers on loading
+static void __attribute__((constructor)) init(void)
 {
+	if (!real_malloc_ptr)
+	{
+		real_malloc_ptr = (void* (*)(size_t))(dlsym(RTLD_NEXT, "malloc"));
+		real_realloc_ptr = (void* (*)(void *, size_t))(dlsym(RTLD_NEXT, "realloc"));
+		real_free_ptr = (void (*)(void *))(dlsym(RTLD_NEXT, "free")); 
+	}
 }
 
 __inline__ static void 
@@ -373,104 +383,125 @@ tr_log(const void* caller, void* ptr, void* old,
 	}
 }
 
-static void
-tr_freehook (ptr, caller)
-     void* ptr;
-     const void* caller;
+void free(void *ptr)
 {
-	if (ptr == NULL)
+	if (ptr == NULL ||
+	   (ptr >= (void*)malloc_init_buf && ptr < (void*)(malloc_init_buf + malloc_init_pos)))
+	{
 		return;
-	if (ptr == mallwatch)
-		tr_break ();
+	}
+
+	if (!real_free_ptr)
+	{
+		// Should never happen!
+		return;
+	}
+
+	if (!tr_free_hook_enabled || tr_free_hook_active)
+	{
+		// Real function is called if hooks are not enabled or if
+		// free is called recursively while logging data (should
+		// probably never happen!)
+		return (*real_free_ptr)(ptr);
+	}
+
+	pthread_mutex_lock(&free_lock);
+	tr_free_hook_active = 1;
 
-        pthread_mutex_lock(&lock);
 #ifdef PROFILE
 	tr_frees++;
 	tr_current_mallocs--;
 #endif
 
-	__free_hook = tr_old_free_hook;
-
-	if (tr_old_free_hook != NULL)
-		(*tr_old_free_hook) (ptr, caller);
-	else
-		free(ptr);
+	void *caller = __builtin_return_address(0);
+	(*real_free_ptr)(ptr);
 	tr_log(caller, ptr, 0, 0, TR_FREE);
 
-	__free_hook = tr_freehook;
-        pthread_mutex_unlock(&lock);
+	tr_free_hook_active = 0;
+	pthread_mutex_unlock(&free_lock);
 }
 
-static void*
-tr_mallochook (size, caller)
-     size_t size;
-     const void* caller;
+void *malloc(size_t size)
 {
-	void* hdr;
+	if (!real_malloc_ptr)
+	{
+		// Return memory from static buffer if available
+		pthread_mutex_lock(&malloc_init_lock);
+		if ((malloc_init_pos + size) > MALLOC_INIT_SIZE)
+		{
+			pthread_mutex_unlock(&malloc_init_lock);
+			return NULL;
+		}
+		else
+		{
+			void *ptr = (void*)(malloc_init_buf + malloc_init_pos);
+			malloc_init_pos += size;
+			pthread_mutex_unlock(&malloc_init_lock);
+			return ptr;
+		}
+	}
 
-        pthread_mutex_lock(&lock);
+	if (!tr_malloc_hook_enabled || tr_malloc_hook_active)
+	{
+		// Real function is called if hooks are not enabled or if
+		// malloc is called recursively while logging data
+		return (*real_malloc_ptr)(size);
+	}
 
-	__malloc_hook = tr_old_malloc_hook;
-	__realloc_hook = tr_old_realloc_hook;
-	__free_hook = tr_old_free_hook;
+	pthread_mutex_lock(&malloc_lock);
+	tr_malloc_hook_active = 1;
 
-	if (tr_old_malloc_hook != NULL)
-		hdr = (void*) (*tr_old_malloc_hook) (size, caller);
-	else
-		hdr = (void*) malloc(size);
+	void *caller = __builtin_return_address(0);
+	void *hdr = (*real_malloc_ptr)(size);
 	tr_log(caller, hdr, 0, size, TR_MALLOC);
+
 	/* We only build the allocation tree if mallTreeFile has been set. */
 	if (mallTreeFile)
 		addAllocationToTree();
 
-	__malloc_hook = tr_mallochook;
-	__realloc_hook = tr_reallochook;
-	__free_hook = tr_freehook;
-
 #ifdef PROFILE
 	tr_mallocs++;
 	tr_current_mallocs++;
 	if (tr_current_mallocs > tr_max_mallocs)
 		tr_max_mallocs = tr_current_mallocs;
 #endif
-        pthread_mutex_unlock(&lock);
 
-	if (hdr == mallwatch)
-		tr_break ();
+	tr_malloc_hook_active = 0;
+	pthread_mutex_unlock(&malloc_lock);
 
 	return hdr;
 }
 
-static void*
-tr_reallochook (ptr, size, caller)
-     void* ptr;
-     size_t size;
-     const void* caller;
+void* realloc(void *ptr, size_t size)
 {
-	void* hdr;
-
-	if (ptr == mallwatch)
-		tr_break ();
+	if (ptr >= (void*)malloc_init_buf && ptr < (void*)(malloc_init_buf + malloc_init_pos))
+	{
+		// Reallocating pointer from initial static buffer, nothing to free. Just call malloc
+		return malloc(size);
+	}
 
-        pthread_mutex_lock(&lock);
+	if (!real_realloc_ptr)
+	{
+		// Should never happen
+		return NULL;
+	}
 
-	__free_hook = tr_old_free_hook;
-	__malloc_hook = tr_old_malloc_hook;
-	__realloc_hook = tr_old_realloc_hook;
+	if (!tr_realloc_hook_enabled || tr_realloc_hook_active)
+	{
+		// Real function is called if hooks are not enabled or if
+		// realloc is called recursively while logging data
+		return (*real_realloc_ptr)(ptr, size);
+	}
 
-	if (tr_old_realloc_hook != NULL)
-		hdr = (void*) (*tr_old_realloc_hook) (ptr, size, caller);
-	else
-		hdr = (void*) realloc (ptr, size);
+	pthread_mutex_lock(&realloc_lock);
+	tr_realloc_hook_active = 1;
 
+	void *caller = __builtin_return_address(0);
+	void *hdr = (*real_realloc_ptr)(ptr, size);
 	tr_log(caller, hdr, ptr, size, TR_REALLOC);
 
-	__free_hook = tr_freehook;
-	__malloc_hook = tr_mallochook;
-	__realloc_hook = tr_reallochook;
-
 #ifdef PROFILE
-	/* If ptr is 0 there was no previos malloc of this location */
+	// If ptr is 0 there was no previos malloc of this location
 	if (ptr == NULL)
 	{
 		tr_mallocs++;
@@ -480,10 +511,8 @@ tr_reallochook (ptr, size, caller)
 	}
 #endif
 
-        pthread_mutex_unlock(&lock);
-
-	if (hdr == mallwatch)
-		tr_break ();
+	tr_realloc_hook_active = 0;
+	pthread_mutex_unlock(&realloc_lock);
 
 	return hdr;
 }
@@ -671,10 +700,7 @@ release_libc_mem (void)
 #endif
 
 /* We enable tracing if either the environment variable MALLOC_TRACE
- * or the variable MALLOC_TREE are set, or if the variable mallwatch
- * has been patched to an address that the debugging user wants us to
- * stop on.  When patching mallwatch, don't forget to set a breakpoint
- * on tr_break! */
+ * or the variable MALLOC_TREE are set */
 void
 ktrace()
 {
@@ -701,7 +727,7 @@ ktrace()
 	if( getenv("MALLOC_THRESHOLD") != NULL )
 	    mallThreshold = atol(getenv("MALLOC_THRESHOLD"));
 #endif
-	if (mallfile != NULL || mallTreeFile != NULL || mallwatch != NULL)
+	if (mallfile != NULL || mallTreeFile != NULL)
     {
 		mallstream = fopen (mallfile != NULL ? mallfile : "/dev/null", "w");
 		if (mallstream != NULL)
@@ -717,14 +743,10 @@ ktrace()
 			if(*buf)
 				fprintf (mallstream, "#%s\n", buf);
 
-			/* Save old hooks and hook in our own functions for all
-			 * malloc, realloc and free calls */
-			tr_old_free_hook = __free_hook;
-			__free_hook = tr_freehook;
-			tr_old_malloc_hook = __malloc_hook;
-			__malloc_hook = tr_mallochook;
-			tr_old_realloc_hook = __realloc_hook;
-			__realloc_hook = tr_reallochook;
+			// Enable hooks
+			tr_malloc_hook_enabled = 1;
+			tr_realloc_hook_enabled = 1;
+			tr_free_hook_enabled = 1;
 
 			tr_cache_pos = TR_CACHE_SIZE;
 			do
@@ -748,14 +770,14 @@ ktrace()
 void
 kuntrace()
 {
+	// Disable hooks
+	tr_malloc_hook_enabled = 0;
+	tr_realloc_hook_enabled = 0;
+	tr_free_hook_enabled = 0;
+
 	if (mallstream == NULL)
 		return;
 
-	/* restore hooks to original values */
-	__free_hook = tr_old_free_hook;
-	__malloc_hook = tr_old_malloc_hook;
-	__realloc_hook = tr_old_realloc_hook;
-
 	if (removeBranchesBelowThreshold(CallTree))
 		CallTree = 0;
 	if (mallTreeFile)
@@ -802,9 +824,6 @@ int fork()
     {
       close(fileno(mallstream));
       mallstream = NULL;
-      __free_hook = tr_old_free_hook;
-      __malloc_hook = tr_old_malloc_hook;
-      __realloc_hook = tr_old_realloc_hook;
     }
   }
   return result;