vmprotect_macros/
lib.rs

1use std::ffi::CString;
2
3use proc_macro2::TokenStream;
4use quote::{format_ident, quote, quote_spanned};
5use syn::parse::{Parse, ParseStream};
6use syn::spanned::Spanned as _;
7use syn::visit_mut::VisitMut;
8use syn::{
9    braced, parse_macro_input, Attribute, FnArg, LitCStr, LitStr, Pat, PatIdent, Path, Signature,
10    Token, Visibility,
11};
12use twox_hash::XxHash3_64;
13
14struct ItemFnMini {
15    pub attrs: Vec<Attribute>,
16    pub vis: Visibility,
17    pub sig: Signature,
18    pub block: TokenStream,
19}
20impl Parse for ItemFnMini {
21    fn parse(input: ParseStream) -> Result<Self, syn::Error> {
22        let attrs = input.call(Attribute::parse_outer)?;
23        let vis: Visibility = input.parse()?;
24        let sig: Signature = input.parse()?;
25        let content;
26        let _brace = braced!(content in input);
27        let block = content.parse()?;
28        Ok(Self {
29            attrs,
30            vis,
31            sig,
32            block,
33        })
34    }
35}
36
37mod kw {
38    syn::custom_keyword!(mutate);
39    syn::custom_keyword!(virtualize);
40    syn::custom_keyword!(ultra);
41    syn::custom_keyword!(destroy);
42    syn::custom_keyword!(lock);
43}
44
45enum ProtectionType {
46    Mutate,
47    Virtualize,
48    Ultra,
49    Destroy,
50}
51impl ProtectionType {
52    fn name(&self) -> &'static str {
53        match self {
54            ProtectionType::Mutate => "mutate",
55            ProtectionType::Virtualize => "virtualize",
56            ProtectionType::Ultra => "ultra",
57            ProtectionType::Destroy => "destroy",
58        }
59    }
60}
61impl Parse for ProtectionType {
62    fn parse(input: ParseStream) -> syn::Result<Self> {
63        let lookahead = input.lookahead1();
64        Ok(if lookahead.peek(kw::mutate) {
65            input.parse::<kw::mutate>()?;
66            Self::Mutate
67        } else if lookahead.peek(kw::virtualize) {
68            input.parse::<kw::virtualize>()?;
69            Self::Virtualize
70        } else if lookahead.peek(kw::ultra) {
71            input.parse::<kw::ultra>()?;
72            Self::Ultra
73        } else if lookahead.peek(kw::destroy) {
74            input.parse::<kw::destroy>()?;
75            Self::Destroy
76        } else {
77            return Err(lookahead.error());
78        })
79    }
80}
81struct ProtectMacroAttr {
82    ty: ProtectionType,
83    lock: bool,
84}
85impl Parse for ProtectMacroAttr {
86    fn parse(input: ParseStream) -> syn::Result<Self> {
87        let ty: ProtectionType = input.parse()?;
88        let lookahead = input.lookahead1();
89        let lock = if lookahead.peek(Token![,]) {
90            input.parse::<Token![,]>()?;
91            input.parse::<kw::lock>()?;
92            true
93        } else if !input.is_empty() {
94            return Err(lookahead.error());
95        } else {
96            false
97        };
98        Ok(Self { ty, lock })
99    }
100}
101
102#[proc_macro]
103pub fn marker_name(s: proc_macro::TokenStream) -> proc_macro::TokenStream {
104    let s = parse_macro_input!(s as LitStr);
105    let Ok(cstring) = CString::new(s.value()) else {
106        return quote_spanned! {s.span() => compile_error!("marker name should not include internal NULs")}.into();
107    };
108    let s = LitCStr::new(cstring.as_c_str(), s.span());
109
110    quote! {::core::ffi::CStr::as_ptr(#s)}.into()
111}
112
113enum WrappedVisitorMode {
114    RetainNormal,
115    RetainWrapped,
116}
117struct WrappedAttrVisitor {
118    mode: WrappedVisitorMode,
119    found_normal: bool,
120    found_wrapped: bool,
121}
122impl WrappedAttrVisitor {
123    fn retain_wrapped() -> Self {
124        Self {
125            mode: WrappedVisitorMode::RetainWrapped,
126            found_normal: false,
127            found_wrapped: false,
128        }
129    }
130    fn retain_normal() -> Self {
131        Self {
132            mode: WrappedVisitorMode::RetainNormal,
133            found_normal: false,
134            found_wrapped: false,
135        }
136    }
137}
138impl VisitMut for WrappedAttrVisitor {
139    fn visit_attributes_mut(&mut self, i: &mut Vec<syn::Attribute>) {
140        i.retain(|a| {
141            let wrapped = a.path().is_ident("wrapped");
142            self.found_normal |= !wrapped;
143            self.found_wrapped |= wrapped;
144            match self.mode {
145                WrappedVisitorMode::RetainNormal => !wrapped,
146                WrappedVisitorMode::RetainWrapped => wrapped,
147            }
148        });
149    }
150}
151
152#[proc_macro_attribute]
153pub fn protected(
154    attr: proc_macro::TokenStream,
155    fn_ts: proc_macro::TokenStream,
156) -> proc_macro::TokenStream {
157    let attr = parse_macro_input!(attr as ProtectMacroAttr);
158
159    let ItemFnMini {
160        attrs,
161        vis,
162        sig,
163        block,
164    } = parse_macro_input!(fn_ts as ItemFnMini);
165
166    // For functions with the same name, we want a deterministic, yet unique identifier
167    // Until there is no way to look at span position, we can just use its debug repr.
168    // Next Rust release it will be possible to do that: https://github.com/rust-lang/rust/pull/140514
169    let s = format!("{:?}", sig.ident.span());
170    let id = XxHash3_64::oneshot(s.as_bytes());
171
172    let wrapped_ident = format_ident!(
173        "VMPROTECT_MARKER_{}{}_END_{id}",
174        attr.ty.name(),
175        if attr.lock { "_lock" } else { "" },
176    );
177
178    let mut wrapped_attrs = attrs.clone();
179    let mut wrapped_sig = sig.clone();
180    {
181        let mut visitor = WrappedAttrVisitor::retain_wrapped();
182        visitor.visit_attributes_mut(&mut wrapped_attrs);
183        visitor.visit_signature_mut(&mut wrapped_sig);
184    }
185    wrapped_sig.ident = wrapped_ident.clone();
186
187    let mut wrapper_attrs = attrs.clone();
188    let mut wrapper_sig = sig;
189    {
190        let mut visitor = WrappedAttrVisitor::retain_normal();
191        visitor.visit_attributes_mut(&mut wrapper_attrs);
192        // Signature is handled below
193    }
194
195    let should_inline_wrapper = !wrapper_attrs.iter().any(|v| {
196        let mut path = v.path().to_owned();
197        if path.is_ident("unsafe") {
198            if let Ok(ipath) = v.parse_args::<Path>() {
199                path = ipath;
200            }
201        }
202        path.is_ident("inline")
203            || path.is_ident("no_mangle")
204            || path.is_ident("export_name")
205            || path.is_ident("link_section")
206    });
207    let wrapper_inline_attr = should_inline_wrapper.then(|| quote!(#[inline(always)]));
208
209    if let Some(gt) = wrapped_sig.generics.gt_token {
210        // Maybe provide some macro for easier multi-specialization?
211        // TODO: Allow lifetime annotations
212        return quote_spanned! {gt.span() => compile_error!("Protected functions don't support generics.\nCreate manually monomorphized versions of the function, and then protect those.")}.into();
213    }
214    if let Some(asyncness) = wrapped_sig.asyncness {
215        // TODO: Document how does it work, limitations et cetera?
216        return quote_spanned! {asyncness.span() => compile_error!("VMProtect won't work as you would expect with async functions.")}.into();
217    }
218
219    if let Some(receiver) = wrapped_sig.receiver() {
220        // In theory it should be possible to do
221        //
222        // fn wrapper(self: &Ty) {
223        //    fn wrapped(this: &Ty) {}
224        // }
225        //
226        // I.e require user to specify full receiver type, and then rename `self` in code to `this`,
227        // but this it too complicated.
228        return quote_spanned! {receiver.span() => compile_error!("Receivers are not supported on protected functions.\nExtract your impl block function, and pass the call arguments manualy.")}.into();
229    }
230    let mut args = vec![];
231    for (i, ele) in wrapper_sig.inputs.iter_mut().enumerate() {
232        let mut visitor = WrappedAttrVisitor::retain_normal();
233        let FnArg::Typed(t) = ele else {
234            unreachable!("receivers are forbidden");
235        };
236        match &mut *t.pat {
237            Pat::Ident(PatIdent {
238                attrs,
239                ident,
240                subpat: None,
241                by_ref,
242                mutability,
243            }) => {
244                visitor.visit_attributes_mut(attrs);
245                *by_ref = None;
246                *mutability = None;
247                args.push(ident.clone());
248            }
249            p => {
250                visitor.visit_pat_mut(p);
251                if visitor.found_normal {
252                    return quote_spanned! {p.span() => compile_error!("Non-wrapped attributes are only supported on non-pattern arguments")}.into();
253                }
254                let ident = format_ident!("unnamed_{i}");
255                args.push(ident.clone());
256                t.pat = Box::new(Pat::Ident(PatIdent {
257                    attrs: vec![],
258                    by_ref: None,
259                    mutability: None,
260                    ident,
261                    subpat: None,
262                }))
263            }
264        }
265    }
266
267    (quote! {
268        #wrapper_inline_attr
269        #(#wrapper_attrs)*
270        #vis
271        #wrapper_sig
272        {
273            #[inline(never)]
274            #[doc(hidden)]
275            #(#wrapped_attrs)*
276            #wrapped_sig
277            { #block }
278            #wrapped_ident(#(#args),*)
279        }
280
281
282    })
283    .into()
284}
vmprotect_macros/lib.rs

vmprotect_macros/
lib.rs
