vmprotect_macros/
lib.rs

1use proc_macro::{TokenStream, TokenTree};
2use syn::{punctuated::Punctuated, ItemFn};
3use twox_hash::xxh3::hash128;
4
5#[proc_macro_attribute]
6pub fn protected(attr: TokenStream, fn_ts: TokenStream) -> TokenStream {
7    println!("{:?}", attr);
8    let mut attr = attr.into_iter();
9    let prot_type = if let Some(prot_type @ TokenTree::Ident { .. }) = attr.next() {
10        prot_type.to_string()
11    } else {
12        panic!("missing protection type")
13    };
14
15    if !["mutate", "virtualize", "ultra"].contains(&prot_type.as_ref()) {
16        panic!("unknown protection type: {}", prot_type);
17    }
18
19    let lock = match attr.next() {
20        Some(TokenTree::Punct(..)) => {
21            if let Some(lock @ TokenTree::Ident { .. }) = attr.next() {
22                if lock.to_string() != "lock" {
23                    panic!("only possible variant is lock")
24                }
25                true
26            } else {
27                panic!("only possible variant is lock");
28            }
29        }
30
31        Some(_) => panic!("expected lock declaration"),
32        None => false,
33    };
34
35    let ItemFn {
36        attrs,
37        vis,
38        sig,
39        block,
40    } = syn::parse(fn_ts.clone()).expect("failed to parse as fn");
41    let mut name = "vmprotect_".to_owned();
42    name.push_str(&prot_type);
43    if lock {
44        name.push_str("_lock");
45    }
46    name.push('_');
47
48    name.push_str(&format!("{}", hash128(sig.ident.to_string().as_bytes())));
49    let wrapped_ident = syn::Ident::new(&name, sig.ident.span());
50    let wrapper = syn::ItemFn {
51        attrs: attrs.clone(),
52        vis: vis.clone(),
53        sig: sig.clone(),
54        block: {
55            let mut args: Punctuated<_, syn::token::Comma> = Punctuated::new();
56            for arg in &sig.inputs {
57                args.push(match arg {
58                    syn::FnArg::Receiver(_r) => panic!("not supported on trait/struct members"),
59                    syn::FnArg::Typed(t) => t.pat.clone(),
60                });
61            }
62
63            syn::parse(
64                (quote::quote! {
65                    {#wrapped_ident(#args)}
66                })
67                .into(),
68            )
69            .unwrap()
70        },
71    };
72    let mut wrapped_sig = sig.clone();
73    wrapped_sig.ident = wrapped_ident;
74    let wrapped = syn::ItemFn {
75        attrs,
76        vis: syn::Visibility::Inherited,
77        sig: wrapped_sig,
78        block,
79    };
80    (quote::quote! {
81        #[inline(never)]
82        #[no_mangle]
83        #[doc(hidden)]
84        #wrapped
85
86        #[inline(always)]
87        #wrapper
88    })
89    .into()
90}
vmprotect_macros/lib.rs

vmprotect_macros/
lib.rs
