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

forge_macros/
lib.rs
