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    unsafe {
74        sftrace_setup(
75            sftrace_entry_slot,
76            sftrace_exit_slot,
77            sftrace_tailcall_slot
78        );
79    }
80}
81
82pub struct SftraceAllocator<A: GlobalAlloc>(pub A);
83
84unsafe impl<A: GlobalAlloc> GlobalAlloc for SftraceAllocator<A> {
85    #[inline]
86    unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
87        unsafe {
88            let v = std::alloc::System.alloc(layout);
89            sftrace_alloc_event(1, layout.size(), layout.align(), v);
90            v
91        }
92    }
93
94    #[inline]
95    unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
96        unsafe {
97            sftrace_alloc_event(2, layout.size(), layout.align(), ptr);
98            std::alloc::System.dealloc(ptr, layout);           
99        }
100    }
101
102    #[inline]
103    unsafe fn alloc_zeroed(&self, layout: std::alloc::Layout) -> *mut u8 {
104        unsafe {
105            let v = std::alloc::System.alloc_zeroed(layout);
106            sftrace_alloc_event(1, layout.size(), layout.align(), v);
107            v
108        }
109    }
110
111    #[inline]
112    unsafe fn realloc(&self, ptr: *mut u8, layout: std::alloc::Layout, new_size: usize) -> *mut u8 {
113        unsafe {
114            sftrace_alloc_event(4, layout.size(), layout.align(), ptr);
115            let v = std::alloc::System.realloc(ptr, layout, new_size);
116            sftrace_alloc_event(3, new_size, layout.align(), v);
117            v            
118        }
119    }
120}
sftrace_setup/lib.rs

sftrace_setup/
lib.rs
