memleak hook

Signed-off-by: zhengkaijie <2658880633@qq.com>

porting/linux/user/include/signal.h

@@ -202,6 +202,7 @@ int __libc_current_sigrtmax(void);
 #define MUSL_SIGNAL_NATIVE_LOCAL (SIGRTMIN + 3)
 #define MUSL_SIGNAL_JSHEAP (SIGRTMIN + 4)
 #define MUSL_SIGNAL_JSHEAP_PRIV (SIGRTMIN + 5)
+#define MUSL_SIGNAL_MEMCHECK (SIGRTMIN + 9)
 
 int kill(pid_t, int);
 int sigemptyset(sigset_t *);

porting/linux/user/src/hook/malloc_common.c

@@ -70,4 +70,14 @@ void* realloc(void *p, size_t n)
 		return MuslMalloc(realloc)(p, n);
 	}
 }
+
+size_t malloc_usable_size(void* addr)
+{
+	volatile const struct MallocDispatchType* dispatch_table = get_current_dispatch_table();
+	if (__predict_false(dispatch_table != NULL)) {
+		return dispatch_table->malloc_usable_size(addr);
+	} else {
+		return MuslMalloc(malloc_usable_size)(addr);
+	}
+}
 #endif

porting/linux/user/src/hook/musl_malloc_dispatch_table.h

@@ -7,7 +7,9 @@
 struct musl_libc_globals {
 	volatile atomic_llong current_dispatch_table;
 	volatile atomic_llong so_dispatch_table;
+	volatile atomic_llong memleak_tracker_so_dispatch_table;
 	struct MallocDispatchType malloc_dispatch_table;
+	struct MallocDispatchType memleak_tracker_malloc_dispatch_table;
 };
 
 #endif

porting/linux/user/src/hook/musl_preinit.c

@@ -47,10 +47,24 @@ static struct MallocDispatchType __ohos_malloc_hook_init_dispatch = {
 };
 #define MAX_SYM_NAME_SIZE 1000
 #define MAX_PROC_NAME_SIZE 256
+#define ADDR_NATIVE_ENABLE (1<<4)
+#define ADDR_NATIVE_DISABLE (1<<5)
+#define ADDR_NATIVE_SAVE (1<<6)
+#define ADDR_NATIVE_CLEAR (1<<7)
 static char *__malloc_hook_shared_lib = "libnative_hook.z.so";
 static char *__malloc_hook_function_prefix = "ohos_malloc_hook";
 volatile atomic_llong ohos_malloc_hook_shared_library;
+static char *kMemTrackSharedLib = "libmemleak_tracker.so";
+static char *kMemTrackPrefix = "track";
+static char *kMemTrackPropertyEnable = "const.hiview.memleak_tracker.enable";
+static char *kMemTrackSign = "true";
+bool checkLoadMallocMemTrack = false;
+unsigned int memLeakTypeContent = 0;
+volatile atomic_llong memleak_ever_shared_library_handle;
+volatile atomic_llong ohos_malloc_ever_shared_library_handle;
 void* function_of_shared_lib[LAST_FUNCTION];
+void* function_of_ohos_malloc_shared_lib[LAST_FUNCTION];
+void* function_of_memleak_shared_lib[LAST_FUNCTION];
 static enum EnumHookMode __hook_mode = STEP_HOOK_MODE;
 
 static void  get_native_hook_param(char *buf, unsigned int buf_len)
@@ -64,6 +78,21 @@ static void  get_native_hook_param(char *buf, unsigned int buf_len)
 #endif
 }
 
+static void  get_memleak_hook_param()
+{
+#ifdef OHOS_ENABLE_PARAMETER
+	const char *key =  kMemTrackPropertyEnable;
+	unsigned int len = OHOS_PARAM_MAX_SIZE;
+	char memleak_param_value[OHOS_PARAM_MAX_SIZE + 1] = {0};
+	(void)SystemReadParam(key, memleak_param_value, &len);
+	if (strncmp(memleak_param_value, kMemTrackSign, strlen(kMemTrackSign)) == 0) {
+		checkLoadMallocMemTrack = true;
+	}
+#else
+	return;
+#endif
+}
+
 static int parse_hook_variable(enum EnumHookMode* mode, char* path, int size)
 {
 	if (!mode || !path || size <= 0) {
@@ -212,6 +241,17 @@ static bool init_realloc_function(void* malloc_shared_library_handler, MallocRea
 	return true;
 }
 
+static bool init_malloc_usable_size_function(void* malloc_shared_library_handler, MallocMallocUsableSizeType* func, const char* prefix)
+{
+	char symbol[MAX_SYM_NAME_SIZE];
+	snprintf(symbol, sizeof(symbol), "%s_%s", prefix, "malloc_usable_size");
+	*func = (MallocMallocUsableSizeType)(dlsym(malloc_shared_library_handler, symbol));
+	if (*func == NULL) {
+		return false;
+	}
+	return true;
+}
+
 static bool init_hook_functions(void* shared_library_handler, struct MallocDispatchType* table, const char* prefix)
 {
 	if (!init_malloc_function(shared_library_handler, &table->malloc, prefix)) {
@@ -235,6 +275,9 @@ static bool init_hook_functions(void* shared_library_handler, struct MallocDispa
 	if (!init_memorytag_function(shared_library_handler, prefix)) {
 		return false;
 	}
+	if (!init_malloc_usable_size_function(shared_library_handler, &table->malloc_usable_size, prefix)) {
+		return false;
+	}
 	return true;
 }
 
@@ -242,6 +285,11 @@ static void clear_function_table()
 {
 	for (size_t i = 0; i < LAST_FUNCTION; i++) {
 		function_of_shared_lib[i] = NULL;
+		if (__get_memleak_hook_flag()) {
+			function_of_memleak_shared_lib[i] = NULL;
+		} else if (__get_global_hook_flag()) {
+			function_of_ohos_malloc_shared_lib[i] = NULL;
+		}
 	}
 }
 
@@ -260,6 +308,11 @@ bool init_malloc_hook_shared_library(void* shared_library_handle, const char* sh
 		char symbol[MAX_SYM_NAME_SIZE];
 		snprintf(symbol, sizeof(symbol), "%s_%s", prefix, names[i]);
 		function_of_shared_lib[i] = dlsym(shared_library_handle, symbol);
+		if (__get_memleak_hook_flag()) {
+			function_of_memleak_shared_lib[i] = function_of_shared_lib[i];
+		} else if (__get_global_hook_flag()) {
+			function_of_ohos_malloc_shared_lib[i] = function_of_shared_lib[i];
+		}
 		if (function_of_shared_lib[i] == NULL) {
 			// __musl_log(__MUSL_LOG_ERROR, "%s: %s routine not found in %s\n", getprogname(), symbol, shared_lib);
 			clear_function_table();
@@ -318,14 +371,23 @@ bool finish_install_ohos_malloc_hooks(struct musl_libc_globals* globals, const c
 		return false;
 	}
 	on_start_func_t start_func = (on_start_func_t)(function_of_shared_lib[ON_START_FUNCTION]);
-	if (!start_func()) {
-		// __musl_log(__MUSL_LOG_ERROR, "%s: failed to start %s\n", getprogname(), prefix);
-		clear_function_table();
-		return false;
+	if (__get_global_hook_flag()) {
+		if (!start_func()) {
+			// __musl_log(__MUSL_LOG_ERROR, "%s: failed to start %s\n", getprogname(), prefix);
+			clear_function_table();
+			return false;
+		}
+		atomic_store_explicit(&__musl_libc_globals.so_dispatch_table, (volatile long long)&globals->malloc_dispatch_table, memory_order_seq_cst);
+		atomic_store_explicit(&__musl_libc_globals.current_dispatch_table, (volatile long long)&globals->malloc_dispatch_table, memory_order_seq_cst);
+	}
+	if (__get_memleak_hook_flag() && checkLoadMallocMemTrack) {
+		if (!start_func(memLeakTypeContent)) {
+			clear_function_table();
+			return false;
+		}
+		atomic_store_explicit(&__musl_libc_globals.memleak_tracker_so_dispatch_table, (volatile long long)&globals->memleak_tracker_malloc_dispatch_table, memory_order_seq_cst);
+		atomic_store_explicit(&__musl_libc_globals.current_dispatch_table, (volatile long long)&globals->memleak_tracker_malloc_dispatch_table, memory_order_seq_cst);
 	}
-	atomic_store_explicit(&__musl_libc_globals.so_dispatch_table, (volatile long long)&globals->malloc_dispatch_table, memory_order_seq_cst);
-	atomic_store_explicit(&__musl_libc_globals.current_dispatch_table, (volatile long long)&globals->malloc_dispatch_table, memory_order_seq_cst);
-
 	int ret_value = atexit(malloc_finalize);
 	if (ret_value != 0) {
 		// __musl_log(__MUSL_LOG_ERROR, "failed to set atexit cleanup function: %d\n", ret_value);
@@ -343,18 +405,27 @@ static bool is_empty_string(const char* str)
 	return true;
 }
 
-static void install_ohos_malloc_hook(struct musl_libc_globals* globals)
+static void install_ohos_malloc_hook(struct musl_libc_globals* globals, const char* shared_lib, const char* prefix)
 {
 	volatile void* shared_library_handle = (volatile void *)atomic_load_explicit(&ohos_malloc_hook_shared_library, memory_order_acquire);
 	assert(shared_library_handle == NULL || shared_library_handle == (volatile void*)-1);
-	shared_library_handle = (volatile void*)load_malloc_hook_shared_library(__malloc_hook_shared_lib, __malloc_hook_function_prefix, &globals->malloc_dispatch_table);
+	if (__get_memleak_hook_flag()) {
+		shared_library_handle = (volatile void*)load_malloc_hook_shared_library(shared_lib, prefix, &globals->memleak_tracker_malloc_dispatch_table);
+	} else if (__get_global_hook_flag()) {
+		shared_library_handle = (volatile void*)load_malloc_hook_shared_library(shared_lib, prefix, &globals->malloc_dispatch_table);
+	}
 	if (shared_library_handle == NULL) {
 		// __musl_log(__MUSL_LOG_ERROR, "Can't load shared library '%s'\n", __malloc_hook_shared_lib);
 		return;
 	}
 
-	if (finish_install_ohos_malloc_hooks(globals, NULL, __malloc_hook_function_prefix)) {
+	if (finish_install_ohos_malloc_hooks(globals, NULL, prefix)) {
 		atomic_store_explicit(&ohos_malloc_hook_shared_library, (volatile long long)shared_library_handle, memory_order_seq_cst);
+		if (strncmp(__malloc_hook_function_prefix, prefix, strlen(prefix)) == 0) {
+			atomic_store_explicit(&ohos_malloc_ever_shared_library_handle, (volatile long long)shared_library_handle, memory_order_seq_cst);
+		} else {
+			atomic_store_explicit(&memleak_ever_shared_library_handle, (volatile long long)shared_library_handle, memory_order_seq_cst);
+		}
 	} else {
 		// __musl_log(__MUSL_LOG_ERROR, "finish_install_ohos_malloc_hooks failed\n");
 		dlclose((void *)shared_library_handle);
@@ -364,7 +435,15 @@ static void install_ohos_malloc_hook(struct musl_libc_globals* globals)
 
 static void* init_ohos_malloc_hook()
 {
-	install_ohos_malloc_hook(&__musl_libc_globals);
+	if (__get_memleak_hook_flag()) {
+		get_memleak_hook_param();
+		if (checkLoadMallocMemTrack) {
+			install_ohos_malloc_hook(&__musl_libc_globals, kMemTrackSharedLib, kMemTrackPrefix);
+		}
+	}
+	if (__get_global_hook_flag()) {
+		install_ohos_malloc_hook(&__musl_libc_globals, __malloc_hook_shared_lib, __malloc_hook_function_prefix);
+	}
 	return NULL;
 }
 
@@ -388,10 +467,29 @@ static void __set_default_malloc()
 	atomic_store_explicit(&__musl_libc_globals.current_dispatch_table, (volatile const long long)NULL, memory_order_seq_cst);
 }
 
+static void __restore_hook_function_table()
+{
+	for (size_t i = 0; i < LAST_FUNCTION; i++) {
+		if (__get_memleak_hook_flag()) {
+			function_of_shared_lib[i] = function_of_memleak_shared_lib[i];
+		} else if (__get_global_hook_flag()) {
+			function_of_shared_lib[i] = function_of_ohos_malloc_shared_lib[i];
+		}
+	}
+}
+
 static void __install_malloc_hook()
 {
+	if (__get_memleak_hook_flag()) {
+		return;
+	}
 	atomic_store_explicit(&__hook_enable_hook_flag, (volatile bool)true, memory_order_seq_cst);
-
+	volatile void* ohos_malloc_ever_handle = (volatile void* )atomic_load_explicit(&ohos_malloc_ever_shared_library_handle, memory_order_acquire);
+	if (ohos_malloc_ever_handle != NULL) {
+		atomic_store_explicit(&ohos_malloc_hook_shared_library, (volatile long long)ohos_malloc_ever_handle, memory_order_seq_cst);
+	} else {
+		atomic_store_explicit(&ohos_malloc_hook_shared_library, (volatile long long)0, memory_order_seq_cst);
+	}
 	volatile void* shared_library_handle = (volatile void* )atomic_load_explicit(&ohos_malloc_hook_shared_library, memory_order_acquire);
 	if (shared_library_handle == NULL) {
 		if (__hook_mode == STEP_HOOK_MODE) {
@@ -401,17 +499,22 @@ static void __install_malloc_hook()
 			init_ohos_malloc_hook();
 		}
 	} else if (shared_library_handle != (void*)-1) {
+		__restore_hook_function_table();
+		volatile const struct MallocDispatchType* so_dispatch_value = (volatile const struct MallocDispatchType* )atomic_load_explicit(&__musl_libc_globals.so_dispatch_table, memory_order_acquire);
+		atomic_store_explicit(&__musl_libc_globals.current_dispatch_table, (volatile long long)so_dispatch_value, memory_order_seq_cst);
 		on_start_func_t start_func = (on_start_func_t)(function_of_shared_lib[ON_START_FUNCTION]);
 		if (start_func && !start_func()) {
 			// __musl_log(__MUSL_LOG_ERROR, "%s: failed to enable malloc\n", getprogname());
 		}
-		volatile const struct MallocDispatchType* so_dispatch_value = (volatile const struct MallocDispatchType* )atomic_load_explicit(&__musl_libc_globals.so_dispatch_table, memory_order_acquire);
-		atomic_store_explicit(&__musl_libc_globals.current_dispatch_table, (volatile long long)so_dispatch_value, memory_order_seq_cst);
+
 	}
 }
 
 static void __uninstal_malloc_hook()
 {
+	if (__get_memleak_hook_flag()) {
+		return;
+	}
 	atomic_store_explicit(&__hook_enable_hook_flag, (volatile bool)false, memory_order_seq_cst);
 	bool flag = __set_hook_flag(false);
 	__set_default_malloc();
@@ -421,6 +524,56 @@ static void __uninstal_malloc_hook()
 	}
 }
 
+static void __install_memleak_tracker_hook(int32_t sigNum, siginfo_t *info, void *ptr)
+{
+	if (__get_global_hook_flag()) {
+		return;
+	}
+	atomic_store_explicit(&__memleak_hook_flag, (volatile bool)true, memory_order_seq_cst);
+	volatile void* memleak_ever_handle = (volatile void* )atomic_load_explicit(&memleak_ever_shared_library_handle, memory_order_acquire);
+	if (memleak_ever_handle != NULL) {
+		atomic_store_explicit(&ohos_malloc_hook_shared_library, (volatile long long)memleak_ever_handle, memory_order_seq_cst);
+	} else {
+		atomic_store_explicit(&ohos_malloc_hook_shared_library, (volatile long long)0, memory_order_seq_cst);
+	}
+	volatile void* shared_library_handle = (volatile void* )atomic_load_explicit(&ohos_malloc_hook_shared_library, memory_order_acquire);
+	memLeakTypeContent = (unsigned int)(siginfo_t *)((info)->si_addr);
+	if (shared_library_handle == NULL) {
+		if (__hook_mode == STEP_HOOK_MODE) {
+			if (memLeakTypeContent & ADDR_NATIVE_ENABLE) {
+				atomic_store_explicit(&__musl_libc_globals.current_dispatch_table, (volatile const long long)&__ohos_malloc_hook_init_dispatch, memory_order_seq_cst);
+				atomic_store_explicit(&ohos_malloc_hook_shared_library, (volatile long long)-1, memory_order_seq_cst);
+			}
+		} else {
+			init_ohos_malloc_hook();
+		}
+	} else if (shared_library_handle != (void*)-1) {
+		if (checkLoadMallocMemTrack) {
+			__restore_hook_function_table();
+			on_start_func_t start_func = (on_start_func_t)(function_of_shared_lib[ON_START_FUNCTION]);
+			if (memLeakTypeContent & ADDR_NATIVE_ENABLE) {
+				volatile const struct MallocDispatchType* memleak_tracker_so_dispatch_value = (volatile const struct MallocDispatchType* )atomic_load_explicit(&__musl_libc_globals.memleak_tracker_so_dispatch_table, memory_order_acquire);
+				atomic_store_explicit(&__musl_libc_globals.current_dispatch_table, (volatile long long)memleak_tracker_so_dispatch_value, memory_order_seq_cst);
+			}
+			if (start_func && !start_func(memLeakTypeContent)) {
+				// __musl_log(__MUSL_LOG_ERROR, "%s: failed to enable malloc\n", getprogname());
+				clear_function_table();
+			}
+			if (memLeakTypeContent & ADDR_NATIVE_ENABLE) {
+				atomic_store_explicit(&__memleak_hook_flag, (volatile bool)false, memory_order_seq_cst);
+				bool flag = __set_hook_flag(false);
+				__set_default_malloc();
+				on_end_func_t end_func = (on_end_func_t)(function_of_shared_lib[ON_END_FUNCTION]);
+				if (end_func) {
+					end_func();
+				}
+				memLeakTypeContent = 0;
+			}
+		}
+	}
+}
+
+
 static void __install_malloc_hook_signal_handler()
 {
 	struct sigaction actionInstallHook = {};
@@ -434,6 +587,15 @@ static void __install_malloc_hook_signal_handler()
 	sigemptyset(&actionDef.sa_mask);
 	sigaddset(&actionDef.sa_mask, MUSL_SIGNAL_HOOK);
 	sigaction(MUSL_SIGNAL_UNHOOK, &actionDef, NULL);
+
+	struct sigaction actionInstallMemleakHook = {};
+	actionInstallMemleakHook.sa_handler = NULL;
+	actionInstallMemleakHook.sa_sigaction = __install_memleak_tracker_hook;
+	sigemptyset(&actionInstallMemleakHook.sa_mask);
+	sigaddset(&actionInstallMemleakHook.sa_mask, MUSL_SIGNAL_MEMCHECK);
+	actionInstallMemleakHook.sa_flags = SA_SIGINFO;
+	actionInstallMemleakHook.sa_restorer = NULL;
+	sigaction(MUSL_SIGNAL_MEMCHECK, &actionInstallMemleakHook, NULL);
 }
 
 static void __initialize_malloc()
@@ -444,6 +606,7 @@ static void __initialize_malloc()
 __attribute__((constructor(1))) static void __musl_initialize()
 {
 	atomic_store_explicit(&__hook_enable_hook_flag, (volatile bool)false, memory_order_seq_cst);
+	atomic_store_explicit(&__memleak_hook_flag, (volatile bool)false, memory_order_seq_cst);
 	__set_default_malloc();
 	char hook_process_path[MAX_PROC_NAME_SIZE + 1] = {0};
 	parse_hook_variable(&__hook_mode, hook_process_path, sizeof(hook_process_path) - 1);

porting/linux/user/src/hook/musl_preinit_common.c

@@ -14,8 +14,10 @@ struct MallocDispatchType __libc_malloc_default_dispatch = {
 	.munmap = MuslMalloc(munmap),
 	.calloc = MuslMalloc(calloc),
 	.realloc = MuslMalloc(realloc),
+	.malloc_usable_size = MuslMalloc(malloc_usable_size),
 };
 
 volatile atomic_bool __hook_enable_hook_flag;
+volatile atomic_bool __memleak_hook_flag;
 
 #endif

porting/linux/user/src/hook/musl_preinit_common.h

@@ -11,6 +11,9 @@ extern struct MallocDispatchType __libc_malloc_default_dispatch;
 
 extern volatile atomic_bool __hook_enable_hook_flag;
 
+extern volatile atomic_bool __memleak_hook_flag;
+extern bool checkLoadMallocMemTrack;
+
 enum EnumFunc {
 	INITIALIZE_FUNCTION,
 	FINALIZE_FUNCTION,
@@ -47,6 +50,18 @@ inline bool __get_global_hook_flag()
 #endif // HOOK_ENABLE
 }
 
+__attribute__((always_inline))
+inline bool __get_memleak_hook_flag()
+{
+#ifdef HOOK_ENABLE
+	volatile bool memleak_flag = atomic_load_explicit(&__memleak_hook_flag, memory_order_acquire);
+	return memleak_flag;
+#else
+	return false;
+#endif // HOOK_ENABLE
+}
+
+
 __attribute__((always_inline))
 inline bool __get_hook_flag()
 {
@@ -94,6 +109,9 @@ inline volatile const struct MallocDispatchType* get_current_dispatch_table()
 #ifdef HOOK_ENABLE
 	volatile const struct MallocDispatchType* ret = (struct MallocDispatchType *)atomic_load_explicit(&__musl_libc_globals.current_dispatch_table, memory_order_acquire);
 	if (ret != NULL) {
+		if (__get_memleak_hook_flag()) {
+			return ret;
+		}
 		if (!__get_global_hook_flag()) {
 			ret = NULL;
 		}

porting/linux/user/src/internal/malloc_impl.h

@@ -95,6 +95,8 @@ hidden void *internal_calloc(size_t m, size_t n);
 
 hidden void *internal_realloc(void *p, size_t n);
 
+hidden size_t internal_malloc_usable_size(void *p);
+
 #ifdef MALLOC_RED_ZONE
 hidden void chunk_checksum_set(struct chunk *c);
 

porting/linux/user/src/malloc/malloc_usable_size.c

@@ -11,7 +11,16 @@ extern size_t je_malloc_usable_size(void *p);
 
 hidden void *(*const __realloc_dep)(void *, size_t) = realloc;
 
+#ifdef HOOK_ENABLE
+size_t __libc_malloc_usable_size(void* p)
+#else
 size_t malloc_usable_size(void *p)
+#endif
+{
+	return internal_malloc_usable_size(p);
+}
+
+size_t internal_malloc_usable_size(void* p)
 {
 #ifdef USE_JEMALLOC
 	return je_malloc_usable_size(p);