prehook-macros 1.1.0

use proc_macro::TokenStream;
use quote::quote;
use syn::parse::{Parse, ParseStream};
use syn::{FnArg, Ident, ItemFn, LitBool, LitStr, Token, parse_macro_input};

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.unwrap_or_else(|| func_name.to_string());
    let guard_name = format!("__guard_{}", symbol_name);

    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 hook_guard_entry = if args.guard {
        quote! {
            if !prehook::hook::HookGuard::try_enter(#guard_name) {
                call_original!(#(#arg_names),*);
            }
        }
    } else {
        quote! {}
    };

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

    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),*) }
                }
            }

            #hook_guard_entry
            let res =  #body;
            #hook_guard_exit
            return res;
        }
    }.into()
}

pub struct InterposeArgs {
    pub symbol: Option<String>,
    pub guard: bool,
}

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

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

            if ident == "symbol" {
                symbol = Some(input.parse::<LitStr>()?.value());
            }

            if ident == "guard" {
                guard = Some(input.parse::<LitBool>()?.value())
            }

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

        Ok(InterposeArgs {
            symbol,
            guard: guard.unwrap_or_default(),
        })
    }
}