libinject/

instrument.rs

//! **Shadow-stack instrumentation:** Exports the `instruction_event` symbol, which can be
//! registered with DynamoRIO to instrument basic blocks such that a _shadow_ callstack is
//! maintained while to binary executes. At any point in execution this shadow-stack can be queried
//! for the raw absolute frame addresses of the function in the current callstack.
//!

use crate::ffi::{bool_, dr_emit_flags_t, instr_t, instrlist_t};
use crate::utils;
use crate::{MOD_BASE_ADDR, MOD_SIZE, ffi};
use std::os::raw::c_void;

#[derive(Debug, Clone)]
pub struct RawCallStack {
    // Raw absolute frame addresses (captured on calls); we filter to in-module elsewhere.
    pub stack: Vec<Option<usize>>,
}

impl RawCallStack {
    pub fn new() -> Self {
        Self {
            stack: Vec::with_capacity(128),
        }
    }
    pub fn push(&mut self, addr: Option<usize>) {
        self.stack.push(addr);
    }
    pub fn pop(&mut self) {
        self.stack.pop();
    }
}

/// Return raw frame addresses from the shadow stack (order preserved).
///
/// Returns Some(frames) when the thread-local shadow stack is available,
/// or None if the TLS pointer is null (e.g., uninitialized or torn down).
pub fn stack_frames(drcontext: *mut ::std::os::raw::c_void) -> Option<Vec<Option<usize>>> {
    use crate::drcore::log;
    log(&format!("[stack frames inside] in1"));
    unsafe {
        let stack_ptr = ffi::drmgr_get_tls_field(drcontext, crate::tls_idx) as *mut RawCallStack;
        if stack_ptr.is_null() {
            return None;
        }
        Some((*stack_ptr).stack.clone())
    }
}

/// Returns (module_base, module_size).
pub fn module_bounds() -> (usize, usize) {
    let base = *MOD_BASE_ADDR
        .lock()
        .expect("Failed to lock MOD_BASE_ADDR")
        .as_ref()
        .expect("Module base should always be known");
    let size = *MOD_SIZE
        .lock()
        .expect("Failed to lock MOD_SIZE")
        .as_ref()
        .expect("Module size should always be known");
    (base, size)
}

fn in_target_module(addr: usize) -> bool {
    let (base, size) = module_bounds();
    addr >= base && addr < base + size
}

#[unsafe(no_mangle)]
pub unsafe extern "C" fn instruction_event(
    drcontext: *mut ::std::os::raw::c_void,
    _tag: *mut ::std::os::raw::c_void,
    bb: *mut instrlist_t,
    instr: *mut instr_t,
    _for_trace: bool_,
    _translating: bool_,
    _user_data: *mut ::std::os::raw::c_void,
) -> dr_emit_flags_t {
    unsafe {
        if ffi::instr_is_call_direct(instr) != 0 {
            let instr_addr = ffi::instr_get_app_pc(instr);
            let target_addr = ffi::instr_get_branch_target_pc(instr);
            let target_addr_as_opnd = utils::opnd_create_uintptr(target_addr as usize);
            let instr_addr_as_opnd = utils::opnd_create_uintptr(instr_addr as usize);
            ffi::dr_insert_clean_call(
                drcontext,
                bb,
                instr,
                on_direct_call as *mut c_void,
                false as bool_,
                2,
                instr_addr_as_opnd,
                target_addr_as_opnd,
            );
        } else if ffi::instr_is_call_indirect(instr) != 0 {
            ffi::dr_insert_mbr_instrumentation(
                drcontext,
                bb,
                instr,
                on_indirect_call as *mut c_void,
                ffi::dr_spill_slot_t_SPILL_SLOT_1,
            );
        }

        if ffi::instr_is_return(instr) != 0 {
            // Insert call to your shadow stack pop handler (on_ret)
            ffi::dr_insert_mbr_instrumentation(
                drcontext,
                bb,
                instr,
                on_return as *mut c_void,
                ffi::dr_spill_slot_t_SPILL_SLOT_1,
            );
        }
    }

    ffi::dr_emit_flags_t_DR_EMIT_DEFAULT
}

// 32-bit Windows: ensure cdecl
#[cfg(all(target_pointer_width = "32"))]
extern "cdecl" fn on_direct_call(instr_addr: *mut c_void, target_addr: *mut c_void) {
    on_direct_call_logic(instr_addr, target_addr)
}

// All other targets: standard C (maps to the platform C ABI)
#[cfg(not(all(target_pointer_width = "32")))]
extern "C" fn on_direct_call(instr_addr: *mut c_void, target_addr: *mut c_void) {
    on_direct_call_logic(instr_addr, target_addr)
}

fn on_direct_call_logic(instr_addr: *mut c_void, target_addr: *mut c_void) {
    // Only instrument calls from the target module.
    if instr_addr.is_null() | !in_target_module(instr_addr as usize) {
        return;
    }

    let addr = if target_addr.is_null() {
        None
    } else {
        Some(target_addr as usize)
    };

    unsafe {
        let drcontext = ffi::dr_get_current_drcontext();
        let tls = ffi::drmgr_get_tls_field(drcontext, crate::tls_idx) as *mut RawCallStack;
        if !tls.is_null() {
            // Record even null direct calls to preserve callsite information
            (*tls).push(addr);
        }
    }
}

// 32-bit Windows: ensure cdecl
#[cfg(all(target_pointer_width = "32"))]
extern "cdecl" fn on_indirect_call(branch_instr_addr: *const u8, target_addr: *const u8) {
    on_indirect_call_logic(branch_instr_addr, target_addr)
}

// All other targets: standard C (maps to the platform C ABI)
#[cfg(not(all(target_pointer_width = "32")))]
extern "C" fn on_indirect_call(branch_instr_addr: *const u8, target_addr: *const u8) {
    on_indirect_call_logic(branch_instr_addr, target_addr)
}

fn on_indirect_call_logic(branch_instr_addr: *const u8, target_addr: *const u8) {
    unsafe {
        if !in_target_module(branch_instr_addr as usize) {
            return;
        }
        let drcontext = ffi::dr_get_current_drcontext();
        let tls = ffi::drmgr_get_tls_field(drcontext, crate::tls_idx) as *mut RawCallStack;
        if !tls.is_null() {
            (*tls).push(Some(target_addr as usize));
        }
    }
}

// 32-bit Windows: ensure cdecl
#[cfg(all(target_pointer_width = "32"))]
extern "cdecl" fn on_return(branch_instr_addr: *const u8, target_addr: *const u8) {
    on_return_logic(branch_instr_addr, target_addr)
}

// All other targets: standard C (maps to the platform C ABI)
#[cfg(not(all(target_pointer_width = "32")))]
extern "C" fn on_return(branch_instr_addr: *const u8, target_addr: *const u8) {
    on_return_logic(branch_instr_addr, target_addr)
}

fn on_return_logic(_branch_instr_addr: *const u8, target_addr: *const u8) {
    if !in_target_module(target_addr as usize) {
        return;
    }
    unsafe {
        let drcontext = ffi::dr_get_current_drcontext();
        let tls = ffi::drmgr_get_tls_field(drcontext, crate::tls_idx) as *mut RawCallStack;
        if !tls.is_null() {
            (*tls).pop();
        }
    }
}