memleak:support the malloc and free uprobe of third-party allocation libraries

tools/memleak.py

@@ -117,6 +117,10 @@ def run_command_get_pid(command):
         help="report sorted in given key; available key list: size, count")
 parser.add_argument("--symbols-prefix", type=str,
         help="memory allocator symbols prefix")
+parser.add_argument("--costom-malloc-uprobe", type=str,
+        help="costom memory allocator uprobe")
+parser.add_argument("--costom-free-uprobe", type=str,
+        help="costom memory free uprobe")
 
 args = parser.parse_args()
 
@@ -438,6 +442,26 @@ def run_command_get_pid(command):
 }
 """
 
+costom_malloc_uprobe_source = '''
+
+int COSTOM_MALLOC_UPROBE_enter(struct pt_regs *ctx, size_t size) {
+        return gen_alloc_enter(ctx, size, MALLOC);
+}
+
+int COSTOM_MALLOC_UPROBE_exit(struct pt_regs *ctx) {
+        return gen_alloc_exit2(ctx, PT_REGS_RC(ctx), MALLOC);
+}
+
+'''
+
+costom_free_uprobe_source = '''
+
+int COSTOM_FREE_UPROBE_enter(struct pt_regs *ctx, void *address) {
+        return gen_free_enter(ctx, address);
+}
+
+'''
+
 if kernel_trace:
         if args.percpu:
                 bpf_source += bpf_source_percpu
@@ -469,6 +493,15 @@ def run_command_get_pid(command):
         stack_flags += "|BPF_F_USER_STACK"
 bpf_source = bpf_source.replace("STACK_FLAGS", stack_flags)
 
+
+if args.costom_malloc_uprobe is not None:
+        costom_malloc_uprobe_source = costom_malloc_uprobe_source.replace("COSTOM_MALLOC_UPROBE",args.costom_malloc_uprobe)
+        bpf_source = bpf_source + costom_malloc_uprobe_source
+
+if args.costom_free_uprobe is not None:
+        costom_free_uprobe_source = costom_free_uprobe_source.replace("COSTOM_FREE_UPROBE",args.costom_free_uprobe)
+        bpf_source = bpf_source + costom_free_uprobe_source
+
 if args.ebpf:
     print(bpf_source)
     exit()
@@ -508,6 +541,10 @@ def attach_probes(sym, fn_prefix=None, can_fail=False, need_uretprobe=True):
         attach_probes("aligned_alloc", can_fail=True)  # added in C11
         attach_probes("free", need_uretprobe=False)
         attach_probes("munmap", can_fail=True, need_uretprobe=False) # failed on jemalloc
+        if args.costom_malloc_uprobe is not None:
+                attach_probes(args.costom_malloc_uprobe)
+        if args.costom_free_uprobe is not None:
+                attach_probes(args.costom_free_uprobe, need_uretprobe=False)
 
 else:
         print("Attaching to kernel allocators, Ctrl+C to quit.")

tools/memleak_example.txt

@@ -181,6 +181,23 @@ When using the --symbols-prefix argument, memleak can trace the third-party memo
 allocations, such as jemalloc whose symbols are usually identified by the "je_" prefix
 in redis project.
 
+# ./memleak.py  -p $(pgrep leak_test) -O /usr/local/lib/libjemalloc.so.2 -t 
+        --costom-malloc-uprobe=je_malloc_default --costom-free-uprobe=je_free_default
+Attaching to pid 2619464, Ctrl+C to quit.
+(b'leak_test', 2619464, 7, b'....', 501375.308951, b'alloc entered, size = 8192')
+(b'leak_test', 2619464, 7, b'....', 501375.309025, b'alloc exited, size = 8192, result = 7fb87b42b000')
+(b'leak_test', 2619464, 7, b'....', 501388.122185, b'alloc entered, size = 16384')
+(b'leak_test', 2619464, 7, b'....', 501388.122223, b'alloc exited, size = 16384, result = 7fb87b4373c0')
+(b'leak_test', 2619464, 2, b'....', 501413.748229, b'alloc entered, size = 63')
+(b'leak_test', 2619464, 2, b'....', 501413.748265, b'alloc exited, size = 63, result = 7fb87b41b000')
+(b'leak_test', 2619464, 2, b'....', 501413.748272, b'alloc entered, size = 32768')
+(b'leak_test', 2619464, 2, b'....', 501413.74831, b'alloc exited, size = 32768, result = 7fb87b43cf00')
+(b'leak_test', 2619464, 2, b'....', 501413.748358, b'free entered, address = 7fb87b4373c0, size = 16384')
+(b'leak_test', 2619464, 2, b'....', 501413.748373, b'free entered, address = 7fb87b41b000, size = 63')
+
+When using the --costom-malloc-uprobe and --costom-free-uprobe= parameters, we can use a third-party custom malloc and free entry,
+such as jemalloc 5.3.0,--costom-malloc-uprobe=je_malloc_default --costom-free-uprobe=je_free_default
+
 USAGE message:
 
 # ./memleak -h