main.rs

#![feature(link_llvm_intrinsics)]
#![feature(naked_functions)]
#![feature(asm)]

extern "C" {
    #[link_name = "llvm.returnaddress"]
    fn llvm_returnaddress(level: u32) -> *const u8;
}

macro_rules! return_address {
    ($level:literal) => { unsafe { llvm_returnaddress($level) } }
}

#[naked]
#[no_mangle]
#[allow(unused_variables)]
#[allow(no_mangle_generic_items)]
extern "C" fn call_with_fake_return_address_1<T>(caller: *const u8, function: extern "C" fn() -> T) -> T {
    unsafe {
        #[cfg(target_arch = "x86_64")]
        asm!(
            "push rbp",
            "sub rsp, 0x10",
            "mov [rsp + 0x8], rdi",
            "mov rbp, rsp",
            "call rsi",
            "add rsp, 0x10",
            "pop rbp",
            "ret",
            options(noreturn)
        );

        #[cfg(target_arch = "aarch64")]
        asm!(
            "stp x29, x30, [sp, #-32]!",
            "stp xzr, x0, [sp, #16]",
            "add x29, sp, #16",
            "blr x1",
            "ldp x29, x30, [sp], #32",
            "ret",
            options(noreturn)
        );
    }
}

#[naked]
#[no_mangle]
#[allow(unused_variables)]
#[allow(no_mangle_generic_items)]
extern "C" fn call_with_fake_return_address_n<T>(caller: *const u8, function: extern "C" fn() -> T) -> T {
    unsafe {
        #[cfg(target_arch = "x86_64")]
        asm!(
            "push rbp",
            "sub rsp, 0x10",
            "mov [rsp + 0x8], rdi",
            "mov [rsp], rsp",
            "mov rbp, rsp",
            "call rsi",
            "add rsp, 0x10",
            "pop rbp",
            "ret",
            options(noreturn)
        );

        #[cfg(target_arch = "aarch64")]
        asm!(
            "stp x29, x30, [sp, #-32]!",
            "add x29, sp, #16",
            "stp x29, x0, [sp, #16]",
            "blr x1",
            "ldp x29, x30, [sp], #32",
            "ret",
            options(noreturn)
        );
    }
}

extern "C" fn fn1() -> *const u8 {
    return_address!(1)
}
extern "C" fn fn2() -> *const u8 {
    return_address!(7)
}

fn main() {
    println!("{:x}", fn1() as u64);
    //println!("{:x}", fn2() as u64);
    println!("{:x}", call_with_fake_return_address_1(28 as *const u8, fn1 as extern "C" fn() -> *const u8) as u64);
    //println!("{:x}", call_with_fake_return_address_1(28 as *const u8, fn2 as extern "C" fn() -> *const u8) as u64);
    println!("{:x}", call_with_fake_return_address_n(28 as *const u8, fn1 as extern "C" fn() -> *const u8) as u64);
    println!("{:x}", call_with_fake_return_address_n(28 as *const u8, fn2 as extern "C" fn() -> *const u8) as u64);
}