prehook-macros 1.0.0

//! Procedural macros for the prehook library.
//! 
//! This crate provides the `#[hook]` and `#[interpose]` attribute macros, 
//! which are the primary ways to define interceptions in the `prehook` ecosystem.

use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, spanned::Spanned, FnArg, ItemFn};

mod interpose;

/// The primary macro for defining a hook.
/// 
/// This attribute macro transforms a standard Rust function into a 
/// hook handler and automatically generates a constructor to register 
/// the hook with the global `HOOK_REGISTRY`.
/// 
/// ### Parameters
/// - `kind`: (Optional) The type of hook to apply.
///     - `"retn"` (Default): A function-level hook. Injects `call_original!(...)`.
///     - `"jmp_back"`: An inline hook that returns to the next instruction.
///     - `"jmp_to_ret"`: An inline hook that jumps to a dynamic address.
///     - `"jmp_to_addr"`: An inline hook that jumps to a fixed address.
/// - `symbol`: (Optional) The name of the symbol to hook.
/// - `offset`: (Optional) The address offset from the binary's base address.
/// - `dest`: (Required for `"jmp_to_addr"`) The destination offset to jump to.
/// 
/// ### Examples
/// 
/// #### Function Hook (`retn`)
/// ```rust
/// #[hook(kind = "retn", symbol = "check_license")]
/// fn my_hook(key: i32) -> i32 {
///     let result = call_original!(key);
///     if result == 0 { 1 } else { result }
/// }
/// ```
/// 
/// #### Inline Hook (`jmp_back`)
/// ```rust
/// #[hook(kind = "jmp_back", offset = 0x1234)]
/// fn my_inline_hook(reg: &mut Registers) {
///     println!("RAX: 0x{:x}", reg.rax);
/// }
/// ```
#[proc_macro_attribute]
pub fn hook(attr: TokenStream, item: TokenStream) -> TokenStream {
    let args = parse_macro_input!(attr as HookArgs);
    let mut func = parse_macro_input!(item as ItemFn);

    let func_name = func.sig.ident.clone();
    let func_vis = &func.vis;
    let func_attrs = &func.attrs;

    let hook_kind_str = args.kind.as_ref().map(|s| s.value()).unwrap_or_else(|| "retn".to_string());

    let offset = args.offset.unwrap_or_else(|| syn::parse_quote! { 0 });
    let symbol_expr = match &args.symbol {
        Some(s) => quote! { Some(#s.to_string()) },
        None => quote! { None },
    };
    let handler_name = syn::Ident::new(&format!("__hook_{}", func_name), func.sig.span());

    let result = match hook_kind_str.as_str() {
        "retn" => {
            let func_args = func.sig.inputs.clone();
            let mut arg_decls = Vec::new();
            let mut arg_names = Vec::new();
            let mut arg_types = Vec::new();

            for (i, arg) in func_args.iter().enumerate() {
                match arg {
                    FnArg::Receiver(_) => continue,
                    FnArg::Typed(typed) => {
                        let name = &typed.pat;
                        let ty = &typed.ty;
                        arg_names.push(name);
                        arg_types.push(ty);

                        let extraction = match i {
                            0 => quote! { let #name = (*reg).rdi as #ty; },
                            1 => quote! { let #name = (*reg).rsi as #ty; },
                            2 => quote! { let #name = (*reg).rdx as #ty; },
                            3 => quote! { let #name = (*reg).rcx as #ty; },
                            4 => quote! { let #name = (*reg).r8 as #ty; },
                            5 => quote! { let #name = (*reg).r9 as #ty; },
                            _ => {
                                let stack_index = i - 6;
                                quote! { let #name = unsafe { (*reg).get_stack(#stack_index) } as #ty; }
                            },
                        };
                        arg_decls.push(extraction);
                    }
                }
            }

            let return_type = match &func.sig.output {
                syn::ReturnType::Default => quote! { () },
                syn::ReturnType::Type(_, ty) => quote! { #ty },
            };

            func.sig.inputs.push(syn::parse_quote! { __orig_ptr: usize });
            let body = &func.block;
            let sig = &func.sig;

            quote! {
                #(#func_attrs)*
                #func_vis #sig {
                    #[allow(unused_macros)]
                    macro_rules! call_original {
                        ($($arg:expr),*) => {
                            unsafe {
                                let orig: unsafe extern "win64" fn(#(#arg_types),*) -> #return_type = std::mem::transmute(__orig_ptr);
                                orig($($arg),*)
                            }
                        }
                    }
                    #body
                }

                #[ctor::ctor]
                fn #handler_name() {
                    use prehook::hook::{HOOK_REGISTRY, HookInstance};
                    use prehook::hook::RetnRoutine;

                    let handler: RetnRoutine = {
                        unsafe extern "win64" fn __hook_handler(
                            reg: *mut prehook::hook::Registers,
                            orig_func_ptr: usize,
                            _stack: usize,
                        ) -> usize {
                            #( #arg_decls )*
                            let result = #func_name(#(#arg_names),*, orig_func_ptr);
                            result as usize
                        }
                        __hook_handler
                    };

                    let hook = HookInstance::new_retn(
                        #offset,
                        #symbol_expr,
                        stringify!(#func_name).to_string(),
                        handler,
                    );
                    HOOK_REGISTRY.register(hook);
                }
            }
        }
        "jmp_back" => {
            quote! {
                #func

                #[ctor::ctor]
                fn #handler_name() {
                    use prehook::hook::{HOOK_REGISTRY, HookInstance};
                    use prehook::hook::JmpBackRoutine;

                    unsafe extern "win64" fn __hook_handler(
                        reg: *mut prehook::hook::Registers,
                        _user_data: usize,
                    ) {
                        #func_name(&mut *reg);
                    }

                    let hook = HookInstance::new_jmpback(
                        #offset,
                        #symbol_expr,
                        stringify!(#func_name).to_string(),
                        __hook_handler,
                    );
                    HOOK_REGISTRY.register(hook);
                }
            }
        }
        "jmp_to_ret" => {
            quote! {
                #func

                #[ctor::ctor]
                fn #handler_name() {
                    use prehook::hook::{HOOK_REGISTRY, HookInstance};
                    use prehook::hook::JmpToRetRoutine;

                    unsafe extern "win64" fn __hook_handler(
                        reg: *mut prehook::hook::Registers,
                        _orig_func_ptr: usize,
                        _user_data: usize,
                    ) -> usize {
                        #func_name(&mut *reg)
                    }

                    let hook = HookInstance::new_jmptoret(
                        #offset,
                        #symbol_expr,
                        stringify!(#func_name).to_string(),
                        __hook_handler,
                    );
                    HOOK_REGISTRY.register(hook);
                }
            }
        }
        "jmp_to_addr" => {
            let dest = args.dest.expect("jmp_to_addr requires a dest parameter");
            quote! {
                #func

                #[ctor::ctor]
                fn #handler_name() {
                    use prehook::hook::{HOOK_REGISTRY, HookInstance};
                    use prehook::hook::JmpToAddrRoutine;

                    unsafe extern "win64" fn __hook_handler(
                        reg: *mut prehook::hook::Registers,
                        _orig_func_ptr: usize,
                        _user_data: usize,
                    ) {
                        #func_name(&mut *reg);
                    }

                    let hook = HookInstance::new_jmptoaddr(
                        #offset,
                        #symbol_expr,
                        stringify!(#func_name).to_string(),
                        #dest,
                        __hook_handler,
                    );
                    HOOK_REGISTRY.register(hook);
                }
            }
        }
        _ => panic!("Unsupported hook kind: {}", hook_kind_str),
    };

    result.into()
}

/// The macro for defining a symbol interposer.
/// 
/// This attribute macro transforms a standard Rust function into an 
/// exported symbol (via `#[no_mangle]`) and automatically resolves 
/// the original function from the next available library using `dlsym(RTLD_NEXT, ...)`.
/// 
/// ### Parameters
/// - `symbol`: (Optional) The name of the symbol to interpose. Defaults to the function name.
/// 
/// ### Example
/// ```rust
/// #[interpose(symbol = "malloc")]
/// fn my_malloc(size: usize) -> *mut c_void {
///     println!("Allocating {} bytes", size);
///     call_original!(size)
/// }
/// ```
#[proc_macro_attribute]
pub fn interpose(attr: TokenStream, item: TokenStream) -> TokenStream {
    interpose::interpose_impl(attr, item)
}

struct HookArgs {
    offset: Option<syn::LitInt>,
    symbol: Option<syn::LitStr>,
    kind: Option<syn::LitStr>,
    dest: Option<syn::LitInt>,
}

impl syn::parse::Parse for HookArgs {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let mut offset = None;
        let mut symbol = None;
        let mut kind = None;
        let mut dest = None;

        while !input.is_empty() {
            let ident: syn::Ident = input.parse()?;
            let _: syn::Token![=] = input.parse()?;

            if ident == "offset" {
                offset = Some(input.parse()?);
            } else if ident == "symbol" {
                symbol = Some(input.parse()?);
            } else if ident == "kind" {
                kind = Some(input.parse()?);
            } else if ident == "dest" {
                dest = Some(input.parse()?);
            }

            if !input.is_empty() {
                let _: syn::Token![,] = input.parse()?;
            }
        }

        Ok(HookArgs { offset, symbol, kind, dest })
    }
}