sftrace_setup/
lib.rs

1use std::alloc::GlobalAlloc;
2
3unsafe extern "C" {
4    fn sftrace_setup(
5        entry_slot: unsafe extern "C" fn(),
6        exit_slot: unsafe extern "C" fn(),
7        tailcall_slot: unsafe extern "C" fn(),        
8    );
9
10    fn sftrace_alloc_event(
11        kind: u8,
12        size: usize,
13        align: usize,
14        ptr: *mut u8    
15    );    
16}
17
18#[cfg(target_arch = "x86_64")]
19macro_rules! build_slot {
20    ( $( $name:ident );* $( ; )? ) => {
21        $(
22            build_slot!(@$name);
23        )*
24    };
25    ( @ $name:ident ) => {
26        #[unsafe(naked)]
27        unsafe extern "C" fn $name() {
28            // ret + nop * 15
29            std::arch::naked_asm!(
30                "ret",
31                "nop",
32                "nop",
33                "nop",
34                "nop",
35                "nop",
36                "nop",
37                "nop",
38                "nop",
39                "nop",
40                "nop",
41                "nop",
42                "nop",
43                "nop",
44                "nop",
45                "nop",
46            );
47        }
48    };
49}
50
51#[cfg(target_arch = "aarch64")]
52macro_rules! build_slot {
53    ( $( $name:ident );* $( ; )? ) => {
54        $(
55            build_slot!(@$name);
56        )*
57    };
58    ( @ $name:ident ) => {
59        unsafe extern "C" fn $name() {
60            // empty
61        }
62    };
63}
64
65build_slot! {
66    sftrace_entry_slot;
67    sftrace_exit_slot;
68    sftrace_tailcall_slot;
69}
70
71#[inline(always)]
72pub unsafe fn setup() {
73    if std::env::var_os("SFTRACE_OUTPUT_FILE").is_none() {
74        // Not enabled, ignored
75        return;
76    }
77
78    ENABLE_ALLOCATOR_HOOK.store(true, std::sync::atomic::Ordering::Relaxed);
79    
80    unsafe {
81        sftrace_setup(
82            sftrace_entry_slot,
83            sftrace_exit_slot,
84            sftrace_tailcall_slot
85        );
86    }
87}
88
89static ENABLE_ALLOCATOR_HOOK: std::sync::atomic::AtomicBool =
90    std::sync::atomic::AtomicBool::new(false);
91
92pub struct SftraceAllocator<A: GlobalAlloc>(pub A);
93
94unsafe impl<A: GlobalAlloc> GlobalAlloc for SftraceAllocator<A> {
95    #[inline]
96    unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
97        let enable = ENABLE_ALLOCATOR_HOOK.load(std::sync::atomic::Ordering::Relaxed);
98        
99        unsafe {
100            let v = std::alloc::System.alloc(layout);
101            if enable {
102                sftrace_alloc_event(1, layout.size(), layout.align(), v);
103            }
104            v
105        }
106    }
107
108    #[inline]
109    unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
110        let enable = ENABLE_ALLOCATOR_HOOK.load(std::sync::atomic::Ordering::Relaxed);
111
112        unsafe {
113            if enable {
114                sftrace_alloc_event(2, layout.size(), layout.align(), ptr);
115            }
116
117            std::alloc::System.dealloc(ptr, layout);           
118        }
119    }
120
121    #[inline]
122    unsafe fn alloc_zeroed(&self, layout: std::alloc::Layout) -> *mut u8 {
123        let enable = ENABLE_ALLOCATOR_HOOK.load(std::sync::atomic::Ordering::Relaxed);
124
125        unsafe {
126            let v = std::alloc::System.alloc_zeroed(layout);
127            if enable {
128                sftrace_alloc_event(1, layout.size(), layout.align(), v);
129            }
130            v
131        }
132    }
133
134    #[inline]
135    unsafe fn realloc(&self, ptr: *mut u8, layout: std::alloc::Layout, new_size: usize) -> *mut u8 {
136        let enable = ENABLE_ALLOCATOR_HOOK.load(std::sync::atomic::Ordering::Relaxed);
137        unsafe {
138            if enable {
139                sftrace_alloc_event(4, layout.size(), layout.align(), ptr);
140            }
141            let v = std::alloc::System.realloc(ptr, layout, new_size);
142            if enable {
143                sftrace_alloc_event(3, new_size, layout.align(), v);
144            }
145            v            
146        }
147    }
148}
sftrace_setup/lib.rs

sftrace_setup/
lib.rs
