cuda-interposer 0.2.4

use std::{env, ffi::CString, os::raw::c_void, sync::OnceLock};
use tracing::{debug, warn};

// Re-exports for macros
pub use libc;
pub use once_cell;
pub use paste;
pub use tracing;

// ─── Library Loading ─────────────────────────────────────────────────────────

struct DlHandle(*mut c_void);
unsafe impl Send for DlHandle {}
unsafe impl Sync for DlHandle {}

static CUDA_LIB: OnceLock<DlHandle> = OnceLock::new();
static CUDART_LIB: OnceLock<DlHandle> = OnceLock::new();

fn get_libcuda() -> *mut c_void {
    let handle_wrapper = CUDA_LIB.get_or_init(|| unsafe {
        let mut paths = vec![
            "/usr/local/cuda/compat/libcuda.so".to_string(),
            "/usr/lib/x86_64-linux-gnu/libcuda.so".to_string(),
            "/usr/lib64/libcuda.so".to_string(),
            "/usr/local/cuda/targets/x86_64-linux/lib/stubs/libcuda.so".to_string(),
        ];

        if let Some(cuda_home) = env::var_os("CUDA_HOME") {
            let path = format!("{}/compat/libcuda.so", cuda_home.to_string_lossy());
            paths.insert(0, path);
        }

        for path in paths.iter() {
            let s = CString::new(path.clone()).unwrap();
            let flags = libc::RTLD_NOW | libc::RTLD_LOCAL | libc::RTLD_NODELETE;
            let handle = libc::dlopen(s.as_ptr(), flags);
            if !handle.is_null() {
                debug!("Loaded real CUDA driver from: {}", path);
                return DlHandle(handle);
            }
        }

        panic!("Failed to find/load libcuda.so. Ensure it is in LD_LIBRARY_PATH.");
    });
    handle_wrapper.0
}

fn get_libcudart() -> *mut c_void {
    let handle_wrapper = CUDART_LIB.get_or_init(|| unsafe {
        let mut paths = vec![
            "/usr/local/cuda/targets/x86_64-linux/lib/libcudart.so".to_string(),
            "/usr/lib/x86_64-linux-gnu/libcudart.so".to_string(),
            "/usr/lib64/libcudart.so".to_string(),
        ];

        // If CUDA_PATH or CUDA_HOME are set, prioritize those
        if let Some(cuda_home) = env::var_os("CUDA_HOME") {
            let path = format!("{}/targets/x86_64-linux/lib/libcudart.so", cuda_home.to_string_lossy());
            paths.insert(0, path);
        }

        for path in paths.iter() {
            let s = CString::new(path.clone()).unwrap();
            let flags = libc::RTLD_NOW | libc::RTLD_LOCAL | libc::RTLD_NODELETE;
            let handle = libc::dlopen(s.as_ptr(), flags);
            if !handle.is_null() {
                debug!("Loaded real CUDA runtime from: {}", path);
                return DlHandle(handle);
            }
        }

        panic!("Failed to find/load libcudart.so. Ensure CUDA Toolkit is installed and in LD_LIBRARY_PATH.");
    });
    handle_wrapper.0
}

pub fn dlsym_next(symbol: &[u8]) -> *mut c_void {
    let sym_str = std::str::from_utf8(symbol).unwrap_or("");

    // Route to the correct library based on prefix
    let handle = if sym_str.starts_with("cuda") || sym_str.starts_with("__cuda") {
        get_libcudart()
    } else {
        get_libcuda()
    };

    let ptr = unsafe { libc::dlsym(handle, symbol.as_ptr() as *const _) };
    if ptr.is_null() {
        warn!("dlsym_next fail for symbol: {sym_str}");
    }
    ptr
}

// ─── Macros ──────────────────────────────────────────────────────────────────

/// Installs the `cuGetProcAddress` hooks required for the interposer to function.
/// This macro automatically includes the `hook_map.rs` generated by `cuda-interposer-build`.
#[macro_export]
macro_rules! install_hooks {
    () => {
        use std::{
            ffi::CStr,
            os::raw::{c_char, c_int, c_void},
        };
        use tracing::debug;

        #[inline(always)]
        fn get_local_hook(name: &str) -> Option<*mut $crate::libc::c_void> {
            // Include returns the closure expression from hook_map.rs
            let hook_fn = include!(concat!(env!("OUT_DIR"), "/hook_map.rs"));
            hook_fn(name)
        }

        type CUresult = u32; // enum
        type CUdriverProcAddressQueryResult = u32; // enum

        $crate::cuda_hook! {
            pub unsafe extern "C" fn cuGetProcAddress_v2(
                symbol: *const $crate::libc::c_char,
                pfn: *mut *mut $crate::libc::c_void,
                cuda_version: $crate::libc::c_int,
                flags: u64,
                symbol_status: *mut CUdriverProcAddressQueryResult
            ) -> CUresult {
                let sym_name_c = unsafe { ::std::ffi::CStr::from_ptr(symbol) };
                let sym_name = sym_name_c.to_string_lossy();

                // A. Call real implementation
                let real_fn = *__real_cuGetProcAddress_v2;
                let ret = unsafe { real_fn(symbol, pfn, cuda_version, flags, symbol_status) };

                if let Some(our_ptr) = get_local_hook(&sym_name) {
                    $crate::tracing::debug!("Hooking symbol via cuGetProcAddress_v2: {}", sym_name);
                    unsafe { *pfn = our_ptr };
                    return 0; // CUDA_SUCCESS
                }

                ret
            }
        }

        $crate::cuda_hook! {
            pub unsafe extern "C" fn cuGetProcAddress(
                symbol: *const $crate::libc::c_char,
                pfn: *mut *mut $crate::libc::c_void,
                cuda_version: $crate::libc::c_int,
                flags: u64,
                symbol_status: *mut CUdriverProcAddressQueryResult
            ) -> CUresult {
               cuGetProcAddress_v2(symbol, pfn, cuda_version, flags, symbol_status)
            }
        }
    };
}

#[macro_export]
macro_rules! cuda_hook {
    (
        pub unsafe extern "C" fn $fname:ident( $($arg:ident : $arg_ty:ty),* $(,)? )
        -> $ret:ty
        $body:block
    ) => {
        $crate::paste::paste! {
            #[allow(non_upper_case_globals)]
            pub static [<__real_ $fname>]: $crate::once_cell::sync::Lazy<
                unsafe extern "C" fn($($arg_ty),*) -> $ret
            > = $crate::once_cell::sync::Lazy::new(|| {
                let name = concat!(stringify!($fname), "\0");
                let sym = $crate::dlsym_next(name.as_bytes());
                if sym.is_null() {
                    panic!("Missing symbol: {}", stringify!($fname));
                }
                unsafe { std::mem::transmute(sym) }
            });

            #[unsafe(no_mangle)]
            pub unsafe extern "C" fn $fname( $($arg : $arg_ty),* ) -> $ret {
                 $body
            }
        }
    };
}

#[macro_export]
macro_rules! generate_proxy {
    // Internal: Generate specific alias function
    (
        @generate_alias
        alias: $alias:ident,
        target_fn: $fname:ident,
        args: ( [ $( ($arg:ident : $arg_ty:ty) ),* ] ),
        ret: $ret:ty
    ) => {
        $crate::paste::paste! {
            #[unsafe(no_mangle)]
            pub unsafe extern "C" fn $alias( $( $arg : $arg_ty ),* ) -> $ret {
                let f = *[<__REAL_ $fname:upper>];
                f( $( $arg ),* )
            }
        }
    };

    // Internal: Recurse over aliases
    (
        @recurse_aliases
        target_fn: $fname:ident,
        ret: $ret:ty,
        args_tt: $args_tt:tt,
        aliases: [ $($alias:ident),* ]
    ) => {
        $(
            $crate::generate_proxy!(
                @generate_alias
                alias: $alias,
                target_fn: $fname,
                args: $args_tt,
                ret: $ret
            );
        )*
    };

    // Internal: Generate Main Function and Lazy static
    (
        @generate_main
        fn $fname:ident ( [ $( ($arg:ident : $arg_ty:ty) ),* ] ) -> $ret:ty;
        target_symbol: $real_sym:ident
    ) => {
        $crate::paste::paste! {
            static [<__REAL_ $fname:upper>]: $crate::once_cell::sync::Lazy<
                  unsafe extern "C" fn( $($arg_ty),* ) -> $ret
            > = $crate::once_cell::sync::Lazy::new(|| {
                let name = concat!(stringify!($real_sym), "\0");
                let ptr = $crate::dlsym_next(name.as_bytes());
                if ptr.is_null() {
                    eprintln!("fatal: symbol '{}' not found in underlying library", name);
                    std::process::abort();
                }
                unsafe { std::mem::transmute(ptr) }
            });

            #[unsafe(no_mangle)]
            pub unsafe extern "C" fn $fname( $( $arg : $arg_ty ),* ) -> $ret {
                let f = *[<__REAL_ $fname:upper>];
                f( $( $arg ),* )
            }
        }
    };

    // Entry Point
    (
        fn $fname:ident $args_tt:tt -> $ret:ty;
        name: $real_sym:ident
        $(, aliases: $($alias:ident),* )?
    ) => {
        $crate::generate_proxy!(
            @generate_main
            fn $fname $args_tt -> $ret;
            target_symbol: $real_sym
        );
        $(
            $crate::generate_proxy!(
                @recurse_aliases
                target_fn: $fname,
                ret: $ret,
                args_tt: $args_tt,
                aliases: [ $($alias),* ]
            );
        )?
    };
}