prehook-macros 1.0.0

A library for hooking and overriding functions using LD_PRELOAD. Useful for binary modding.
Documentation
use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, FnArg, ItemFn, Ident, LitStr, Token};
use syn::parse::{Parse, ParseStream};

pub fn interpose_impl(attr: TokenStream, item: TokenStream) -> TokenStream {
    let args = parse_macro_input!(attr as InterposeArgs);
    let func = parse_macro_input!(item as ItemFn);

    let func_name = &func.sig.ident;
    let symbol_name = args.symbol.map(|s| s.value()).unwrap_or_else(|| func_name.to_string());
    let symbol_bytes = format!("{}\0", symbol_name);
    
    let func_args = &func.sig.inputs;
    let mut arg_types = Vec::new();
    let mut arg_names = Vec::new();

    for arg in func_args {
        match arg {
            FnArg::Typed(t) => {
                arg_types.push(&t.ty);
                arg_names.push(&t.pat);
            }
            FnArg::Receiver(_) => panic!("interpose does not support 'self'"),
        }
    }
    
    let ret_type = match &func.sig.output {
        syn::ReturnType::Default => quote! { () },
        syn::ReturnType::Type(_, ty) => quote! { #ty },
    };

    let body = &func.block;
    let vis = &func.vis;
    let attrs = &func.attrs;

    quote! {
        #[unsafe(no_mangle)]
        #(#attrs)*
        #vis unsafe extern "C" fn #func_name(#func_args) -> #ret_type {
            use std::sync::OnceLock;
            use std::mem::transmute;
            
            type __OrigFn = unsafe extern "C" fn(#(#arg_types),*) -> #ret_type;
            static __ORIGINAL: OnceLock<__OrigFn> = OnceLock::new();
            
            let __orig = __ORIGINAL.get_or_init(|| {
                unsafe {
                    let ptr = libc::dlsym(libc::RTLD_NEXT, #symbol_bytes.as_ptr() as *const libc::c_char);
                    if ptr.is_null() {
                        let ptr = libc::dlsym(libc::RTLD_DEFAULT, #symbol_bytes.as_ptr() as *const libc::c_char);
                        if ptr.is_null() {
                            panic!("Failed to find original symbol: {}", #symbol_name);
                        }
                        transmute(ptr)
                    } else {
                        transmute(ptr)
                    }
                }
            });

            #[allow(unused_macros)]
            macro_rules! call_original {
                ($($arg:expr),*) => {
                    unsafe { __orig($($arg),*) }
                }
            }

            #body
        }
    }.into()
}

pub struct InterposeArgs {
    pub symbol: Option<LitStr>,
}

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

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

            if ident == "symbol" {
                symbol = Some(input.parse()?);
            }

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

        Ok(InterposeArgs { symbol })
    }
}