memtrace_lib/
lib.rs

1//! # memtrack-lib
2//!
3//! A dynamic library for collecting heap consumption.
4//!
5//! > **Platform support**: Currently tested on Linux (x86_64) and macOS (aarch64-apple-darwin)
6
7mod arch;
8mod trace;
9mod trace_tree;
10mod tracker;
11
12pub use memtrace_utils;
13
14use crate::tracker::Tracker;
15use fishhook::{register, Rebinding};
16use libc::{dlsym, size_t, RTLD_NEXT};
17use std::env;
18use std::ffi::c_void;
19use std::sync::{LazyLock, Mutex, Once};
20
21static INIT: Once = Once::new();
22static mut ORIGINAL_MALLOC: Option<unsafe extern "C" fn(size: size_t) -> *mut c_void> = None;
23static mut ORIGINAL_CALLOC: Option<unsafe extern "C" fn(num: size_t, size: size_t) -> *mut c_void> =
24    None;
25static mut ORIGINAL_REALLOC: Option<
26    unsafe extern "C" fn(ptr: *mut c_void, size: size_t) -> *mut c_void,
27> = None;
28static mut ORIGINAL_FREE: Option<unsafe extern "C" fn(ptr: *mut c_void)> = None;
29static TRACKER: LazyLock<Mutex<Option<Tracker>>> = LazyLock::new(|| Mutex::new(None));
30
31#[no_mangle]
32pub unsafe extern "C" fn my_malloc(size: size_t) -> *mut c_void {
33    let original_malloc = ORIGINAL_MALLOC.unwrap();
34    let ptr = original_malloc(size);
35
36    if let Ok(mut guard) = TRACKER.try_lock() {
37        if let Some(tracker) = guard.as_mut() {
38            tracker.on_malloc(size, ptr as usize);
39        }
40    };
41
42    ptr
43}
44
45#[no_mangle]
46pub unsafe extern "C" fn my_calloc(num: size_t, size: size_t) -> *mut c_void {
47    let original_calloc = ORIGINAL_CALLOC.unwrap();
48    let ptr = original_calloc(num, size);
49
50    if let Ok(mut guard) = TRACKER.try_lock() {
51        if let Some(tracker) = guard.as_mut() {
52            tracker.on_malloc(num * size, ptr as usize);
53        }
54    }
55
56    ptr
57}
58
59#[no_mangle]
60pub unsafe extern "C" fn my_realloc(ptr_in: *mut c_void, size: size_t) -> *mut c_void {
61    let original_realloc = ORIGINAL_REALLOC.unwrap();
62    let ptr_out = original_realloc(ptr_in, size);
63
64    if let Ok(mut guard) = TRACKER.try_lock() {
65        if let Some(tracker) = guard.as_mut() {
66            tracker.on_realloc(size, ptr_in as usize, ptr_out as usize);
67        }
68    }
69
70    ptr_out
71}
72
73#[no_mangle]
74pub unsafe extern "C" fn my_free(ptr: *mut c_void) {
75    let original_free = ORIGINAL_FREE.unwrap();
76    original_free(ptr);
77
78    if let Ok(mut guard) = TRACKER.try_lock() {
79        if let Some(tracker) = guard.as_mut() {
80            tracker.on_free(ptr as usize);
81        }
82    }
83}
84
85pub extern "C" fn my_exit() {
86    if let Ok(mut guard) = TRACKER.try_lock() {
87        if let Some(tracker) = guard.as_mut() {
88            tracker.on_exit();
89        }
90    }
91}
92
93unsafe fn init_functions() {
94    INIT.call_once(|| {
95        let symbol = b"malloc\0";
96        let malloc_ptr = dlsym(RTLD_NEXT, symbol.as_ptr() as *const _);
97        if !malloc_ptr.is_null() {
98            ORIGINAL_MALLOC = Some(std::mem::transmute(malloc_ptr));
99        } else {
100            eprintln!("Error: Could not locate original malloc!");
101        }
102
103        let symbol = b"calloc\0";
104        let calloc_ptr = dlsym(RTLD_NEXT, symbol.as_ptr() as *const _);
105        if !calloc_ptr.is_null() {
106            ORIGINAL_CALLOC = Some(std::mem::transmute(calloc_ptr));
107        } else {
108            eprintln!("Error: Could not locate original calloc!");
109        }
110
111        let symbol = b"realloc\0";
112        let realloc_ptr = dlsym(RTLD_NEXT, symbol.as_ptr() as *const _);
113        if !realloc_ptr.is_null() {
114            ORIGINAL_REALLOC = Some(std::mem::transmute(realloc_ptr));
115        } else {
116            eprintln!("Error: Could not locate original realloc!");
117        }
118
119        let symbol = b"free\0";
120        let free_ptr = dlsym(RTLD_NEXT, symbol.as_ptr() as *const _);
121        if !free_ptr.is_null() {
122            ORIGINAL_FREE = Some(std::mem::transmute(free_ptr));
123        } else {
124            eprintln!("Error: Could not locate original free!");
125        }
126
127        let pipe_filepath = env::var("PIPE_FILEPATH").expect("PIPE_FILEPATH must be set");
128
129        let mut tracker = Tracker::new(pipe_filepath);
130        tracker.init();
131
132        let mut lock = TRACKER.lock().unwrap();
133        *lock = Some(tracker);
134
135        libc::atexit(my_exit);
136    });
137}
138
139#[ctor::ctor]
140fn init() {
141    unsafe {
142        init_functions();
143
144        register(vec![
145            Rebinding {
146                name: "malloc".to_string(),
147                function: my_malloc as *const () as usize,
148            },
149            Rebinding {
150                name: "calloc".to_string(),
151                function: my_calloc as *const () as usize,
152            },
153            Rebinding {
154                name: "realloc".to_string(),
155                function: my_realloc as *const () as usize,
156            },
157            Rebinding {
158                name: "free".to_string(),
159                function: my_free as *const () as usize,
160            },
161        ]);
162    }
163}
memtrace_lib/lib.rs

memtrace_lib/
lib.rs
