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