mallockit_macros/
lib.rs

1use std::collections::HashMap;
2
3use proc_macro::TokenStream;
4use quote::quote;
5
6fn construct_tests_from_config() -> Vec<proc_macro2::TokenStream> {
7    let meta = cargo_metadata::MetadataCommand::new().exec().unwrap();
8    let mut tests = HashMap::new();
9    let ws_meta = meta.workspace_metadata.as_object();
10    if let Some(v) = ws_meta.and_then(|v| v.get("malloc-tests")) {
11        for (name, cmd) in v.as_object().unwrap() {
12            let cmd = cmd.as_str().unwrap();
13            tests.insert(name.to_owned(), cmd.to_owned());
14        }
15    }
16    tests
17        .iter()
18        .map(|(k, v)| {
19            let s = format!(
20                r#"
21                    #[test]
22                    fn {}() {{
23                        ::mallockit::util::testing::malloc::test(env!("CARGO_CRATE_NAME"), {:?});
24                    }}
25                "#,
26                k, v,
27            );
28            s.parse().unwrap()
29        })
30        .collect()
31}
32
33#[proc_macro_attribute]
34pub fn plan(_attr: TokenStream, item: TokenStream) -> TokenStream {
35    let input = syn::parse_macro_input!(item as syn::DeriveInput);
36    let name = &input.ident;
37
38    let tests = construct_tests_from_config();
39
40    let result = quote! {
41        #input
42
43        mod __mallockit_plan {
44            type Plan = super::#name;
45
46            static PLAN: ::mallockit::util::Lazy<Plan> = ::mallockit::util::Lazy::new(|| <Plan as ::mallockit::Plan>::new());
47
48            #[cfg(feature = "malloc")]
49            #[::mallockit::ctor]
50            unsafe fn ctor() {
51                <<Plan as ::mallockit::Plan>::Mutator as ::mallockit::mutator::TLS>::current();
52                ::mallockit::util::sys::hooks::process_start(&*PLAN);
53            }
54
55            #[cfg(target_os = "macos")]
56            #[no_mangle]
57            extern "C" fn mallockit_initialize_macos_tls() -> *mut u8 {
58                use ::mallockit::mutator::TLS;
59                <Plan as ::mallockit::Plan>::Mutator::current() as *mut <Plan as ::mallockit::Plan>::Mutator as _
60            }
61
62            impl ::mallockit::plan::Singleton for super::#name {
63                fn singleton() -> &'static Self {
64                    unsafe { &PLAN }
65                }
66            }
67
68            ::mallockit::export_malloc_api!(PLAN, super::super::#name);
69            ::mallockit::export_rust_global_alloc_api!(super::super::#name);
70        }
71
72        pub use __mallockit_plan::__mallockit_rust_api::Global;
73
74        #[cfg(test)]
75        mod tests {
76            #[cfg(feature = "malloc")]
77            mod malloc {
78                #(#tests)*
79            }
80            ::mallockit::rust_allocator_tests!(crate::Global);
81        }
82    };
83    result.into()
84}
85
86#[proc_macro_attribute]
87pub fn mutator(_attr: TokenStream, item: TokenStream) -> TokenStream {
88    let input = syn::parse_macro_input!(item as syn::DeriveInput);
89    let name = &input.ident;
90    let result = quote! {
91        #[repr(align(256))]
92        #input
93
94
95        mod __mallockit_mutator {
96            #[cfg(not(target_os = "macos"))]
97            fn init() -> super::#name {
98                ::mallockit::mutator::init_pthread_specific();
99                <super::#name as ::mallockit::Mutator>::new()
100            }
101
102            #[cfg(not(target_os = "macos"))]
103            #[thread_local]
104            pub(super) static mut MUTATOR: ::mallockit::util::Lazy<super::#name, ::mallockit::util::Local> = ::mallockit::util::Lazy::new(init);
105
106            #[no_mangle]
107            #[cfg(not(target_os = "macos"))]
108            extern "C" fn mallockit_pthread_destructor() {
109                unsafe {
110                    MUTATOR.reset(init);
111                }
112            }
113
114            #[no_mangle]
115            #[cfg(target_os = "macos")]
116            extern "C" fn mallockit_pthread_destructor() {
117                use crate::mallockit::mutator::TLS;
118                <super::#name as ::mallockit::mutator::TLS>::current().reset();
119            }
120        }
121
122        impl ::mallockit::mutator::TLS for #name {
123            fn new() -> Self {
124                <Self as ::mallockit::Mutator>::new()
125            }
126
127            #[cfg(not(target_os = "macos"))]
128            fn current() -> &'static mut Self {
129                unsafe { &mut *__mallockit_mutator::MUTATOR }
130            }
131        }
132
133    };
134    result.into()
135}
136
137#[proc_macro_attribute]
138pub fn interpose(_attr: TokenStream, item: TokenStream) -> TokenStream {
139    let input = syn::parse_macro_input!(item as syn::ItemFn);
140    let name = &input.sig.ident;
141    let interpose_name = syn::Ident::new(&format!("_interpose_{}", name), name.span());
142    let result = quote! {
143        #[cfg(target_os = "macos")]
144        #[cfg(not(test))]
145        pub mod #name {
146            #[repr(C)]
147            pub struct Interpose {
148                _new: *const (),
149                _old: *const (),
150            }
151
152            #[no_mangle]
153            #[allow(non_upper_case_globals)]
154            #[link_section = "__DATA,__interpose"]
155            pub static mut #interpose_name: Interpose = Interpose {
156                _new: super::#name as *const (),
157                _old: #name as *const (),
158            };
159
160            extern {
161                pub fn #name();
162            }
163        }
164
165        #[cfg(target_os = "macos")]
166        #[cfg(not(test))]
167        #input
168
169        #[cfg(not(target_os = "macos"))]
170        #[cfg(not(test))]
171        #[no_mangle]
172        #input
173    };
174    result.into()
175}
176
177#[proc_macro_attribute]
178pub fn aligned_block(_attr: TokenStream, item: TokenStream) -> TokenStream {
179    let input = syn::parse_macro_input!(item as syn::DeriveInput);
180    let name = &input.ident;
181    let result = quote! {
182        #[repr(transparent)]
183        #input
184
185        mallockit::impl_aligned_block!(#name);
186    };
187    result.into()
188}
mallockit_macros/lib.rs

mallockit_macros/
lib.rs
