Skip to main content

hooking_macros/
lib.rs

1use std::ffi::CString;
2
3// lib.rs of proc macro crate
4use darling::FromMeta;
5use proc_macro::TokenStream;
6use proc_macro2::Literal;
7use quote::quote;
8use syn::{ItemFn, parse_macro_input};
9
10#[derive(FromMeta)]
11#[darling(derive_syn_parse)]
12struct HookArgs {
13    lib: Option<String>,
14    method: String,
15}
16
17#[proc_macro_attribute]
18pub fn hook(attr: TokenStream, item: TokenStream) -> TokenStream {
19    let args: HookArgs = match syn::parse(attr) {
20        Ok(v) => v,
21        Err(e) => {
22            return e.to_compile_error().into();
23        }
24    };
25
26    let func = parse_macro_input!(item as ItemFn);
27    let vis = &func.vis;
28
29    let func_name = &func.sig.ident;
30    let ctor_name = quote::format_ident!("__hook_register_{}", func_name);
31
32    let target_method = Literal::c_string(&CString::new(args.method).unwrap());
33    let target_lib = if let Some(lib) = args.lib {
34        let lib = Literal::c_string(&CString::new(lib).unwrap());
35        quote! { Some(#lib) }
36    } else {
37        quote! { None }
38    };
39
40    let macro_out = quote! {
41        #func
42
43        #vis mod #func_name {
44            use ::hooking::error::Result;
45            use ::hooking::__macro_support;
46
47            static __HOOK: __macro_support::StaticHook = __macro_support::StaticHook::new();
48
49            pub unsafe fn enable_hook() -> Result<()> {
50                unsafe {__macro_support::enable_hook(&__HOOK) }
51            }
52
53            pub unsafe fn disable_hook() -> Result<()> {
54                unsafe { __macro_support::disable_hook(&__HOOK) }
55            }
56
57            pub fn get_hook() -> Result<__macro_support::StaticHookGuard> {
58                unsafe { __macro_support::get_hook(&__HOOK) }
59            }
60
61
62            #[::ctor::ctor(unsafe)]
63            fn #ctor_name() {
64                let ptr = super::#func_name as *mut u8;
65
66                unsafe { ::hooking::__macro_support::init_hook(&__HOOK, #target_lib, #target_method, ptr) };
67            }
68        }
69
70
71    };
72
73    macro_out.into()
74}