1 files changed, 36 insertions, 10 deletions

diff --git a/kmtrace/ktrace.c b/kmtrace/ktrace.c
index 699d1571..0a9cff4f 100644
--- a/kmtrace/ktrace.c
+++ b/kmtrace/ktrace.c
@@ -177,9 +177,9 @@ static void __attribute__((constructor)) init(void)
 {
 	if (!real_malloc_ptr)
 	{
+		real_free_ptr = (void (*)(void *))(dlsym(RTLD_NEXT, "free"));
 		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")); 
 	}
 }
 
@@ -355,10 +355,17 @@ tr_log(const void* caller, void* ptr, void* old,
 
 		tr_cache[i].ptr = ptr; 
 		tr_cache[i].size = size;
-		tr_cache[i].bt = (void**) malloc(TR_BT_SIZE * sizeof(void*));
-		tr_cache[i].bt_size = backtrace(
-			tr_cache[i].bt, TR_BT_SIZE);
-		tr_cache[i].bt = realloc(tr_cache[i].bt, tr_cache[i].bt_size * sizeof(void*));
+		int old_tr_malloc_hook_enabled = tr_malloc_hook_enabled;
+		int old_tr_free_hook_enabled = tr_free_hook_enabled;
+		tr_malloc_hook_enabled = 0;
+		tr_free_hook_enabled = 0;
+		void **tmp_bt = (void**)malloc(TR_BT_SIZE * sizeof(void*));
+		tr_cache[i].bt_size = backtrace(tmp_bt, TR_BT_SIZE);
+		tr_cache[i].bt = (void**)malloc(tr_cache[i].bt_size * sizeof(void*));
+		backtrace(tr_cache[i].bt, tr_cache[i].bt_size);
+		free(tmp_bt);
+		tr_malloc_hook_enabled = old_tr_malloc_hook_enabled;
+		tr_free_hook_enabled = old_tr_free_hook_enabled;
 		tr_cache_level++;
 
 		return;
@@ -431,13 +438,16 @@ void *malloc(size_t size)
 	{
 		// Return memory from static buffer if available
 		pthread_mutex_lock(&malloc_init_lock);
-		if ((malloc_init_pos + size) > MALLOC_INIT_SIZE)
+		if ((malloc_init_pos + size + sizeof(size_t)) > MALLOC_INIT_SIZE)
 		{
 			pthread_mutex_unlock(&malloc_init_lock);
 			return NULL;
 		}
 		else
 		{
+			size_t *sizeptr = (size_t*)(malloc_init_buf + malloc_init_pos);
+			malloc_init_pos += sizeof(size_t);
+			*sizeptr = size;
 			void *ptr = (void*)(malloc_init_buf + malloc_init_pos);
 			malloc_init_pos += size;
 			pthread_mutex_unlock(&malloc_init_lock);
@@ -480,14 +490,30 @@ void* realloc(void *ptr, size_t size)
 {
 	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);
+		// Reallocating pointer from initial static buffer
+		size_t old_size = *(size_t*)((char*)(ptr) - sizeof(size_t));
+		void *new_ptr = malloc(size);
+		size_t min_size = size;
+		if (old_size < min_size)
+		{
+			min_size = old_size;
+		}
+		memcpy(new_ptr, ptr, min_size);
+		return new_ptr;
 	}
 
 	if (!real_realloc_ptr)
 	{
-		// Should never happen
-		return NULL;
+		void *new_ptr = malloc(size);
+		if (new_ptr && ptr && size > 0)
+		{
+			memcpy(new_ptr, ptr, size);
+		}
+		if (ptr)
+		{
+			free(ptr);
+		}
+		return new_ptr;
 	}
 
 	if (!tr_realloc_hook_enabled || tr_realloc_hook_active)