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}