forge_macros/
lib.rs

1use crate::{
2    hook::{HookArgs, HookTransformer},
3    pure_virtual::{PureVirtualFn, VirtualArgs},
4};
5use proc_macro::TokenStream;
6use proc_macro_crate::{FoundCrate, crate_name};
7use proc_macro2::TokenStream as TokenStream2;
8use quote::quote;
9use syn::{DeriveInput, FnArg, ItemFn, parse_macro_input, visit_mut::VisitMut};
10
11fn forge_crate() -> TokenStream2 {
12    match crate_name("mhgu-forge") {
13        Ok(FoundCrate::Itself) => quote! { crate },
14        Ok(FoundCrate::Name(name)) => {
15            let ident = syn::Ident::new(&name, proc_macro2::Span::call_site());
16            quote! { ::#ident }
17        }
18        Err(_) => quote! { ::forge },
19    }
20}
21
22mod hook;
23mod pure_virtual;
24
25/// Marks a function as the plugins entry point.
26///
27/// ### Example
28/// ```ignore
29/// #[forge::entry]
30/// fn my_main_function() {
31///     // Your code here
32/// }
33/// ```
34#[proc_macro_attribute]
35pub fn entry(_attr: TokenStream, item: TokenStream) -> TokenStream {
36    let inner = parse_macro_input!(item as ItemFn);
37    let inner_name = &inner.sig.ident;
38
39    let expanded = quote! {
40        #[inline(always)]
41        #inner
42
43        #[unsafe(no_mangle)]
44        pub extern "C" fn forge_onLoad(params: *mut ::forge::sys::init::PluginInitParams) {
45            unsafe {
46                (*params).required_version = ::forge::REQUIRED_VERSION;
47            }
48        }
49
50        #[unsafe(no_mangle)]
51        pub extern "C" fn forge_onInit(params: *mut ::forge::sys::init::PluginInitParams) {
52            ::forge::log::init().expect("Failed to initialize logger");
53            #inner_name();
54        }
55    };
56
57    expanded.into()
58}
59
60/// Defines a function hook at a fixed offset from a base address.
61///
62/// The annotated function becomes a module of the same name that holds the
63/// hook's static state and can be installed via `forge::install_hook!`.
64///
65/// Inside the body the following pseudo-macros are available:
66/// - `original!(args)` — call the original function (zero or more args).
67/// - `original!()` — call the original function with no arguments.
68/// - `original_function!()` — obtain the raw function pointer without calling it.
69/// - `context!(T)` — borrow the context as `&mut T` (requires context variant of install).
70/// - `context!()` — obtain the raw `*const c_void` context pointer.
71///
72/// ### Example
73/// ```ignore
74/// #[forge::hook(offset = 0x1234)]
75/// fn my_hook(param: u32) -> u32 {
76///     let result = original!(param);
77///     result * 2
78/// }
79///
80/// // Context must outlive the hook (static or Box::leak with "allocator" feature).
81/// static mut MULTIPLIER: u32 = 2;
82///
83/// #[forge::hook(offset = 0x5678)]
84/// fn ctx_hook(value: u32) -> u32 {
85///     let m = context!(u32);
86///     original!(value) * *m
87/// }
88///
89/// #[forge::entry]
90/// fn main() {
91///     let base = forge::mem::text_addr();
92///     forge::install_hook!(base, my_hook);
93///     forge::install_hook!(base, ctx_hook, unsafe { &raw mut MULTIPLIER });
94/// }
95/// ```
96#[proc_macro_attribute]
97pub fn hook(attr: TokenStream, item: TokenStream) -> TokenStream {
98    let args = parse_macro_input!(attr as HookArgs);
99    let func = parse_macro_input!(item as ItemFn);
100
101    let offset = &args.offset;
102    let func_name = &func.sig.ident;
103    let inputs = &func.sig.inputs;
104    let output = &func.sig.output;
105
106    let param_types: Vec<TokenStream2> = inputs
107        .iter()
108        .map(|arg| match arg {
109            FnArg::Typed(pat_type) => {
110                let ty = &pat_type.ty;
111                quote! { #ty }
112            }
113            FnArg::Receiver(_) => {
114                panic!("#[forge::hook] does not support `self` parameters")
115            }
116        })
117        .collect();
118
119    let ret_type = match output {
120        syn::ReturnType::Default => quote! { () },
121        syn::ReturnType::Type(_, ty) => quote! { #ty },
122    };
123
124    let fn_ptr_type = quote! { unsafe extern "C" fn(#(#param_types),*) -> #ret_type };
125
126    let mut body = func.block.clone();
127    HookTransformer.visit_block_mut(&mut body);
128
129    let expanded = quote! {
130        pub mod #func_name {
131            #[allow(unused_imports)]
132            use super::*;
133
134            /// Offset of the hook target from the base address supplied to `install_hook!`.
135            pub const OFFSET: u32 = #offset as u32;
136
137            static mut __HOOK: ::core::mem::MaybeUninit<::forge::sys::hook::Hook> =
138                ::core::mem::MaybeUninit::uninit();
139            static mut __ORIGINAL: *const ::core::ffi::c_void = ::core::ptr::null();
140
141            pub unsafe extern "C" fn __detour(#inputs) #output {
142                let __forge_original: #fn_ptr_type = unsafe {
143                    ::core::mem::transmute(__ORIGINAL)
144                };
145                let __forge_context = unsafe { ::forge::sys::hook::forge_hook_getContext() };
146                #body
147            }
148
149            pub unsafe fn __install(base: u32) {
150                unsafe {
151                    __HOOK.write(::forge::sys::hook::forge_hook_create(
152                        (base + OFFSET) as *const ::core::ffi::c_void,
153                        __detour as *const ::core::ffi::c_void,
154                        ::core::ptr::addr_of_mut!(__ORIGINAL),
155                    ));
156                }
157            }
158
159            pub unsafe fn __install_with_ctx(base: u32, ctx: *const ::core::ffi::c_void) {
160                unsafe {
161                    __HOOK.write(::forge::sys::hook::forge_hook_createWithContext(
162                        (base + OFFSET) as *const ::core::ffi::c_void,
163                        __detour as *const ::core::ffi::c_void,
164                        ::core::ptr::addr_of_mut!(__ORIGINAL),
165                        ctx,
166                    ));
167                }
168            }
169
170            /// Update the context pointer for an already-installed hook.
171            pub unsafe fn __update_ctx(ctx: *const ::core::ffi::c_void) {
172                unsafe {
173                    let result = ::forge::sys::hook::forge_hook_updateContext(
174                        __HOOK.as_mut_ptr(),
175                        ctx,
176                    );
177                    debug_assert_eq!(result, 0, "forge_hook_updateContext failed");
178                }
179            }
180        }
181    };
182
183    expanded.into()
184}
185
186/// Marks a method as a virtual function at a given index in the vtable.
187/// Methods marked with this must be part of a type that implements `HasVtable`, and must take `&self` or `&mut self` as the first parameter.
188/// The method body is replaced with a call to the function pointer at the specified index in the vtable
189///
190/// ### Example
191/// ```ignore
192/// #[derive(forge::HasVtable)]
193/// pub struct MyStruct;
194///
195/// impl MyStruct {
196///     #[forge::pure_virtual(3)]
197///     pub fn my_virtual_func(&self) -> i32 {}
198/// }
199/// ```
200/// Note the lack of an actual implementation of the function.
201#[proc_macro_attribute]
202pub fn pure_virtual(attr: TokenStream, item: TokenStream) -> TokenStream {
203    let args = parse_macro_input!(attr as VirtualArgs);
204    let func = parse_macro_input!(item as PureVirtualFn);
205
206    let func_name = &func.sig.ident;
207    let inputs = &func.sig.inputs;
208    let output = &func.sig.output;
209    let visibility = &func.vis;
210
211    let has_self = !inputs.is_empty()
212        && match &inputs[0] {
213            FnArg::Receiver(receiver) => receiver.reference.is_some(),
214            _ => false,
215        };
216
217    if !has_self {
218        panic!("Functions marked with #[forge::pure_virtual] must have `&self` as their first parameter");
219    }
220
221    let param_types: Vec<TokenStream2> = inputs
222        .iter()
223        .map(|arg| match arg {
224            FnArg::Receiver(receiver) => {
225                let ty = &receiver.ty;
226                quote! { #ty }
227            }
228            FnArg::Typed(pat_type) => {
229                let ty = &pat_type.ty;
230                quote! { #ty }
231            }
232        })
233        .collect();
234
235    let ret_type = match output {
236        syn::ReturnType::Default => quote! { () },
237        syn::ReturnType::Type(_, ty) => quote! { #ty },
238    };
239
240    let fn_ptr_type = quote! { unsafe extern "C" fn(#(#param_types),*) -> #ret_type };
241
242    let index = &args.index;
243    let param_names: Vec<TokenStream2> = inputs
244        .iter()
245        .map(|arg| match arg {
246            FnArg::Receiver(_) => quote! { self },
247            FnArg::Typed(pat_type) => {
248                let pat = &pat_type.pat;
249                quote! { #pat }
250            }
251        })
252        .collect();
253
254    let forge = forge_crate();
255    let expanded = quote! {
256        #visibility fn #func_name(#inputs) #output {
257            let vtable = #forge::sys::cpp::HasVtable::vtable_ptr(self);
258            let addr = unsafe { #forge::sys::cpp::HasVtable::get_virtual_function(self, #index) };
259            let func: #fn_ptr_type = unsafe {
260                ::core::mem::transmute(addr)
261            };
262            unsafe { func(#(#param_names),*) }
263        }
264    };
265
266    expanded.into()
267}
268
269#[proc_macro_derive(HasVtable)]
270pub fn has_vtable_derive(input: TokenStream) -> TokenStream {
271    let input = parse_macro_input!(input as DeriveInput);
272    let type_name = &input.ident;
273    let forge = forge_crate();
274
275    let expanded = quote! {
276        impl #forge::sys::cpp::HasVtable for #type_name {
277            fn vtable_ptr(&self) -> *const *const ::core::ffi::c_void {
278                unsafe {
279                    let ptr = self as *const Self as *const *const *const ::core::ffi::c_void;
280                    *ptr
281                }
282            }
283        }
284    };
285
286    expanded.into()
287}
288
289#[proc_macro_derive(Object)]
290pub fn mt_object_derive(input: TokenStream) -> TokenStream {
291    let input = parse_macro_input!(input as DeriveInput);
292    let type_name = &input.ident;
293    let forge = forge_crate();
294
295    let expanded = quote! {
296        impl #forge::sys::cpp::HasVtable for #type_name {
297            fn vtable_ptr(&self) -> *const *const ::core::ffi::c_void {
298                unsafe {
299                    let ptr = self as *const Self as *const *const *const ::core::ffi::c_void;
300                    *ptr
301                }
302            }
303        }
304
305        impl #forge::mt::object::Object for #type_name {}
306    };
307
308    expanded.into()
309}
310
311#[proc_macro_derive(CacheDti)]
312pub fn cache_dti_derive(input: TokenStream) -> TokenStream {
313    let input = parse_macro_input!(input as DeriveInput);
314    let type_name = &input.ident;
315    let type_name_str = type_name.to_string();
316    let forge = forge_crate();
317
318    let expanded = quote! {
319        impl #forge::mt::dti::CacheDti for #type_name {
320            fn dti() -> Option<&'static #forge::mt::dti::MtDti> {
321                static DTI: core::sync::atomic::AtomicPtr<#forge::mt::dti::MtDti> = core::sync::atomic::AtomicPtr::new(core::ptr::null_mut());
322                let mut ptr = DTI.load(core::sync::atomic::Ordering::Relaxed);
323                if ptr.is_null() {
324                    ptr = #forge::mt::dti::MtDti::find(#type_name_str).map_or(core::ptr::null_mut(), |d| d as *const _ as *mut _);
325                    DTI.store(ptr, core::sync::atomic::Ordering::Relaxed);
326                }
327
328                if ptr.is_null() {
329                    None
330                } else {
331                    Some(unsafe { &*ptr })
332                }
333            }
334        }
335    };
336
337    expanded.into()
338}
forge_macros/lib.rs

forge_macros/
lib.rs
