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