prehook-macros 1.1.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::{FnArg, ItemFn, Lit, Token, parse_macro_input, punctuated::Punctuated, spanned::Spanned};

mod interpose;

const SUPPORTED_HOOK_TYPES: &str = "retn (FunctionHook), jmp_back (InlineHook), jmp_to_ret (DynamicRedirect) and jmp_to_addr (Static Redirect)";

impl TargetKind {}

/// 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,ignore
/// use prehook::hook;
///
/// #[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,ignore
/// use prehook::hook;
///
/// #[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 is_void = match &func.sig.output {
        syn::ReturnType::Default => true,
        syn::ReturnType::Type(_, ty) => {
            if let syn::Type::Tuple(t) = &**ty {
                t.elems.is_empty()
            } else {
                false
            }
        }
    };

    let target_expr = args.target.emit();
    let handler_name = syn::Ident::new(&format!("__hook_{}", func_name), func.sig.span());

    let hook_name_str = func_name.to_string();

    let result = match args.kind {
        HookKind::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;

            let result = if arg_names.is_empty() {
                quote! {
                    let result = #func_name(orig_func_ptr);
                }
            } else {
                quote! {
                   let result = #func_name(#(#arg_names),*, orig_func_ptr);
                }
            };

            let return_val = if is_void {
                quote! { 0 }
            } else {
                quote! { result as usize }
            };
            let guard_name = format!("__guard_{}", hook_name_str);

            let m_orig_call = if is_void {
                quote! {
                    orig(#(#arg_names),*);
                    return 0;
                }
            } else {
                quote! {
                    return orig(#(#arg_names),*) as usize;
                }
            };
            let hook_guard_entry = if args.guard {
                quote! {
                    if !prehook::hook::HookGuard::try_enter(#guard_name) {
                        let orig: unsafe extern "C" fn(#(#arg_types),*) -> #return_type = std::mem::transmute(orig_func_ptr);
                        #m_orig_call
                    }
                }
            } else {
                quote! {}
            };

            let hook_guard_exit = if args.guard {
                quote! {
                    prehook::hook::HookGuard::exit(#guard_name);
                }
            } else {
                quote! {}
            };

            quote! {
                #(#func_attrs)*
                #func_vis #sig {
                    #[allow(unused_macros)]
                    macro_rules! call_original {
                        ($($arg:expr),*) => {
                            unsafe {
                                let orig: unsafe extern "C" 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 )*
                            #hook_guard_entry

                            #result

                            #hook_guard_exit
                            #return_val
                        }
                        __hook_handler
                    };

                    let hook = HookInstance::new_retn(
                        #target_expr,
                        stringify!(#func_name).to_string(),
                        handler,
                    );
                    HOOK_REGISTRY.register(hook);
                }
            }
        }
        HookKind::JmpBack => {
            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(
                        #target_expr,
                        stringify!(#func_name).to_string(),
                        __hook_handler,
                    );
                    HOOK_REGISTRY.register(hook);
                }
            }
        }
        HookKind::JmpToRet => {
            quote! {
                #func

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

                    static HOOK_NAME: &str = #hook_name_str;

                    unsafe extern "win64" fn __hook_handler(
                        reg: *mut prehook::hook::Registers,
                        _orig_func_ptr: usize,
                        _user_data: usize,
                    ) -> usize {
                        let base = prehook::init::MAIN_EXE_BASE.load(std::sync::atomic::Ordering::SeqCst);
                        let hook_point = #target_expr.resolve(base, 0).expect("Failed to resolve hook point address");
                        let result: prehook::hook::Target = #func_name(&mut *reg, _orig_func_ptr);

                        result.resolve(base, hook_point).expect("Failed to resolve return address")
                    }

                    let hook = HookInstance::new_jmptoret(
                        #target_expr,
                        stringify!(#func_name).to_string(),
                        __hook_handler,
                    );
                    HOOK_REGISTRY.register(hook);
                }
            }
        }
        HookKind::JmpToAddr => {
            let target_expr = args.target.emit();
            let dest_expr = args
                .dest
                .expect("dest attribute is required for jmp_to_addr hook")
                .emit();
            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(
                        #target_expr,
                        #dest_expr,
                        stringify!(#func_name).to_string(),
                        __hook_handler,
                    );
                    HOOK_REGISTRY.register(hook);
                }
            }
        }
    };

    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,ignore
/// use prehook::interpose;
///
/// #[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)
}

#[derive(Clone, Debug)]
enum TargetKind {
    Raw(usize),
    Symbol(String),
    SymbolOffset(String, isize),
    Relative(isize),
}

impl TargetKind {
    fn parse_numeric(s: &str) -> isize {
        let s = s.trim();
        let (sign, s) = if let Some(stripped) = s.strip_prefix('+') {
            (1, stripped)
        } else if let Some(stripped) = s.strip_prefix('-') {
            (-1, stripped)
        } else {
            (1, s)
        };

        let val = if let Some(hex_val) = s.strip_prefix("0x") {
            usize::from_str_radix(hex_val, 16).unwrap_or(0)
        } else {
            s.parse::<usize>().unwrap_or(0)
        };

        (val as isize) * sign
    }
    fn emit(&self) -> proc_macro2::TokenStream {
        match self {
            TargetKind::Raw(v) => quote! { prehook::hook::Target::Raw(#v) },
            TargetKind::Symbol(s) => quote! { prehook::hook::Target::Symbol(#s.to_string()) },
            TargetKind::SymbolOffset(s, off) => {
                quote! { prehook::hook::Target::SymbolOffset(#s.to_string(), #off) }
            }
            TargetKind::Relative(off) => quote! { prehook::hook::Target::Relative(#off) },
        }
    }
}

impl syn::parse::Parse for TargetKind {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let lit: Lit = input.parse()?;

        match &lit {
            // Handle raw integer literals like 0x1234
            Lit::Int(i) => Ok(TargetKind::Raw(i.base10_parse()?)),

            // Handle complex string logic
            Lit::Str(s) => {
                let val = s.value();
                let input_str = val.trim();

                if input_str.is_empty() {
                    return Err(syn::Error::new(lit.span(), "Target string cannot be empty"));
                }

                // Handle Explicit Relative: "+0x10", "-16", "+50"
                // We check if it starts with a sign but ISN'T just a hex prefix (like +0x...)
                if (input_str.starts_with('+') || input_str.starts_with('-'))
                    && !input_str.starts_with("0x")
                    && !input_str.starts_with("-0x")
                {
                    return Ok(TargetKind::Relative(Self::parse_numeric(input_str)));
                }

                // Handle Hex String: "0xdeadbeef"
                if let Some(v) = input_str.strip_prefix("0x")
                    && let Ok(val) = usize::from_str_radix(v, 16)
                {
                    return Ok(TargetKind::Raw(val));
                }

                // Handle Symbol + Offset: "main+0x20"
                if let Some((sym, off)) = input_str.rsplit_once('+') {
                    return Ok(TargetKind::SymbolOffset(
                        sym.to_string(),
                        Self::parse_numeric(off),
                    ));
                }

                // Handle Symbol - Offset: "main-0x20"
                // We verify the right side looks like a number to avoid breaking mangled-names-with-hyphens
                if let Some((sym, off)) = input_str.rsplit_once('-') {
                    let off_trimmed = off.trim();
                    if off_trimmed.starts_with("0x")
                        || off_trimmed.chars().all(|c| c.is_ascii_digit())
                    {
                        return Ok(TargetKind::SymbolOffset(
                            sym.to_string(),
                            -Self::parse_numeric(off_trimmed),
                        ));
                    }
                }

                // Default to pure Symbol: "add"
                Ok(TargetKind::Symbol(input_str.to_string()))
            }
            _ => Err(syn::Error::new(
                lit.span(),
                "Expected string or integer for hook target",
            )),
        }
    }
}

enum HookArg {
    Kind(HookKind),
    Target(TargetKind),
    Dest(TargetKind),
    Offset(usize),
    Guard(bool),
}

impl syn::parse::Parse for HookArg {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let ident: syn::Ident = input.parse()?;
        let _: syn::Token![=] = input.parse()?;

        match ident.to_string().as_str() {
            "kind" => Ok(Self::Kind(input.parse()?)),
            "symbol" | "target" => Ok(Self::Target(input.parse()?)),
            "dest" => Ok(Self::Dest(input.parse()?)),
            "offset" => Ok(Self::Offset(input.parse::<syn::LitInt>()?.base10_parse()?)),
            "guard" => Ok(Self::Guard(input.parse::<syn::LitBool>()?.value())),
            _ => Err(syn::Error::new(
                ident.span(),
                format!("Unknown argument: {}", ident),
            )),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum HookKind {
    Retn,
    JmpBack,
    JmpToRet,
    JmpToAddr,
}

impl syn::parse::Parse for HookKind {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let lit: Lit = input.parse()?;

        let Lit::Str(s) = &lit else {
            return Err(syn::Error::new(
                lit.span(),
                "Hook kind only expects strings",
            ));
        };

        let hook_kind_str = s.value();

        Ok(match hook_kind_str.as_str() {
            "retn" | "FunctionHook" => Self::Retn,
            "jmp_back" | "InlineHook" => Self::JmpBack,
            "jmp_to_ret" | "DynamicRedirect" => Self::JmpToRet,
            "jmp_to_addr" | "StaticRedirect" => Self::JmpToAddr,
            _ => {
                return Err(syn::Error::new(
                    lit.span(),
                    format!("Unsupported hook type. Supported types are: {SUPPORTED_HOOK_TYPES}"),
                ));
            }
        })
    }
}

struct HookArgs {
    target: TargetKind,
    kind: HookKind,
    dest: Option<TargetKind>,
    guard: bool,
}

impl syn::parse::Parse for HookArgs {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let args: Punctuated<HookArg, Token![,]> =
            input.parse_terminated(HookArg::parse, Token![,])?;

        let mut kind = None;
        let mut target = None;
        let mut dest = None;
        let mut offset = None;
        let mut guard = None;

        // Distribute the parsed arguments into your struct fields
        for arg in args {
            match arg {
                HookArg::Kind(k) => kind = Some(k),
                HookArg::Target(t) => target = Some(t),
                HookArg::Dest(d) => dest = Some(d),
                HookArg::Offset(o) => offset = Some(o),
                HookArg::Guard(g) => guard = Some(g),
            }
        }

        if kind.is_none() {
            return Err(syn::Error::new(
                input.span(),
                format!(
                    "You need to specify a hook type. Supported types are: {SUPPORTED_HOOK_TYPES}"
                ),
            ));
        }
        let kind = kind.unwrap();

        let final_target = match (target, offset) {
            (Some(TargetKind::Symbol(s)), Some(off)) => {
                // Handle the "symbol + offset" case
                TargetKind::SymbolOffset(s, off as isize)
            }
            (Some(t), _) => t, // Use provided target or symbol
            (None, Some(off)) => TargetKind::Raw(off), // Fallback to raw offset
            (None, None) => {
                // This is the "Hard Guard" you wanted
                return Err(syn::Error::new(
                    input.span(),
                    "No hook target provided. You must specify 'symbol', 'offset', or 'target'.",
                ));
            }
        };

        // you can not use relative under target as it is likely to lead to segfault
        if let TargetKind::Relative(_) = final_target {
            return Err(syn::Error::new(
                input.span(),
                "You cannot use a relative offset on a hook point. This will likely lead to segfaults. Use raw offset instead",
            ));
        }

        if kind == HookKind::JmpToAddr && dest.is_none() {
            return Err(syn::Error::new(
                input.span(),
                "The 'jmp_to_addr' hook requires a 'dest' attribute (e.g., dest = 0x1c0 or dest = '+0x10').",
            ));
        }

        Ok(HookArgs {
            target: final_target,
            kind,
            dest,
            guard: guard.unwrap_or_default(),
        })
    }
}