Sign in Sign up
kretrod/lodestone Public
Branches
master
Kretrod Dear ImGui Enigma menu, Linux support, silent-aim fix, crystal and mace c385e97 · 4 hours ago History
310 lines (278 loc) · 11.1 KB Raw 
//! A minimal inline hook.
//!
//! Overwrite the first instructions of a function with an absolute jump to our
//! code, and keep the bytes we clobbered in a trampoline that jumps back — so
//! the original function is still callable.
//!
//!   target:      ff 25 00 00 00 00 <8-byte detour>   (jmp qword ptr [rip+0])
//!   trampoline:  <stolen instructions> ff 25 00 00 00 00 <target+n>
//!
//! The jump is 14 bytes and fully absolute, which avoids having to find free
//! memory within ±2 GB of the target.
//!
//! Only the three memory primitives differ between platforms — making the page
//! writable, reserving a page near the target, and flushing the instruction
//! cache — so they live behind `mem` and everything else is shared.

use std::ffi::c_void;

use iced_x86::{
    BlockEncoder, BlockEncoderOptions, Decoder, DecoderOptions, FlowControl, Instruction,
    InstructionBlock,
};

const JMP_LEN: usize = 14;

#[cfg(windows)]
mod mem {
    use std::ffi::c_void;

    use windows_sys::Win32::System::Memory::{
        VirtualAlloc, VirtualProtect, MEM_COMMIT, MEM_RESERVE, PAGE_EXECUTE_READWRITE,
    };
    use windows_sys::Win32::System::SystemInformation::{GetSystemInfo, SYSTEM_INFO};

    pub unsafe fn granularity() -> u64 {
        let mut info: SYSTEM_INFO = std::mem::zeroed();
        GetSystemInfo(&mut info);
        info.dwAllocationGranularity.max(0x1000) as u64
    }

    /// Reserve executable memory at exactly `addr`, or null if it is taken.
    pub unsafe fn reserve_at(addr: u64, size: usize) -> *mut u8 {
        VirtualAlloc(
            addr as *const c_void,
            size,
            MEM_COMMIT | MEM_RESERVE,
            PAGE_EXECUTE_READWRITE,
        ) as *mut u8
    }

    /// Make `len` bytes writable and executable, handing back the old flags.
    pub unsafe fn unprotect(addr: *mut u8, len: usize) -> Option<u32> {
        let mut old = 0u32;
        if VirtualProtect(addr as *const c_void, len, PAGE_EXECUTE_READWRITE, &mut old) == 0 {
            None
        } else {
            Some(old)
        }
    }

    pub unsafe fn reprotect(addr: *mut u8, len: usize, old: u32) {
        let mut prev = 0u32;
        VirtualProtect(addr as *const c_void, len, old, &mut prev);
    }

    pub unsafe fn flush(addr: *mut u8, len: usize) {
        use windows_sys::Win32::System::Diagnostics::Debug::FlushInstructionCache;
        use windows_sys::Win32::System::Threading::GetCurrentProcess;
        FlushInstructionCache(GetCurrentProcess(), addr as *const c_void, len);
    }
}

#[cfg(unix)]
mod mem {
    use std::ffi::c_void;

    pub unsafe fn granularity() -> u64 {
        let p = libc::sysconf(libc::_SC_PAGESIZE);
        if p > 0 {
            p as u64
        } else {
            0x1000
        }
    }

    /// Reserve executable memory at exactly `addr`, or null if it is taken.
    /// MAP_FIXED_NOREPLACE is what makes this a probe rather than a demolition:
    /// without it a plain MAP_FIXED would silently unmap whatever lives there.
    pub unsafe fn reserve_at(addr: u64, size: usize) -> *mut u8 {
        let p = libc::mmap(
            addr as *mut c_void,
            size,
            libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC,
            libc::MAP_PRIVATE | libc::MAP_ANONYMOUS | libc::MAP_FIXED_NOREPLACE,
            -1,
            0,
        );
        if p == libc::MAP_FAILED {
            return std::ptr::null_mut();
        }
        // On kernels without MAP_FIXED_NOREPLACE the hint is merely advisory, so
        // an address we did not ask for is no use: give it straight back.
        if p as u64 != addr {
            libc::munmap(p, size);
            return std::ptr::null_mut();
        }
        p as *mut u8
    }

    /// mprotect works on whole pages, so the range is rounded outwards. There
    /// are no "previous flags" to read back, so the caller gets a placeholder.
    pub unsafe fn unprotect(addr: *mut u8, len: usize) -> Option<u32> {
        let (start, span) = page_span(addr, len);
        if libc::mprotect(
            start,
            span,
            libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC,
        ) == 0
        {
            Some(0)
        } else {
            None
        }
    }

    /// Back to read+execute: the patch is in place and should not stay writable.
    pub unsafe fn reprotect(addr: *mut u8, len: usize, _old: u32) {
        let (start, span) = page_span(addr, len);
        libc::mprotect(start, span, libc::PROT_READ | libc::PROT_EXEC);
    }

    unsafe fn page_span(addr: *mut u8, len: usize) -> (*mut c_void, usize) {
        let page = granularity() as usize;
        let start = (addr as usize) & !(page - 1);
        let end = ((addr as usize) + len + page - 1) & !(page - 1);
        (start as *mut c_void, end - start)
    }

    pub unsafe fn flush(_addr: *mut u8, _len: usize) {
        // x86-64 keeps its instruction cache coherent with stores, so there is
        // nothing to do here.
    }
}

pub struct Hook {
    target: *mut u8,
    original: [u8; JMP_LEN],
    /// The stolen prologue plus a jump back: call this to reach the real function.
    pub trampoline: *const c_void,
}

// SAFETY: a Hook is just three addresses; the code it points at is immutable
// once installed, and installation happens once on a single thread.
unsafe impl Send for Hook {}
unsafe impl Sync for Hook {}

/// Encode `jmp qword ptr [rip+0]; dq dest`.
fn abs_jmp(dest: u64) -> [u8; JMP_LEN] {
    let mut b = [0u8; JMP_LEN];
    b[0] = 0xFF;
    b[1] = 0x25;
    // disp32 = 0: the address follows immediately.
    b[6..14].copy_from_slice(&dest.to_le_bytes());
    b
}

/// Install a hook on `target`, routing calls to `detour`.
///
/// # Safety
/// `target` must be the entry point of a real function and `detour` must have a
/// compatible signature. The caller must keep the returned `Hook` alive for as
/// long as the hook is installed.
pub unsafe fn install(target: *mut c_void, detour: *const c_void) -> Result<Hook, String> {
    let mut target = target as *mut u8;

    // An import thunk (`jmp rel32`) is too short to hold our jump; hook the
    // real function it points at instead.
    for _ in 0..4 {
        if *target == 0xE9 {
            let rel = i32::from_le_bytes([
                *target.add(1),
                *target.add(2),
                *target.add(3),
                *target.add(4),
            ]);
            target = (target as i64 + 5 + rel as i64) as *mut u8;
            continue;
        }
        break;
    }

    // Someone (probably an earlier copy of us) is already here.
    if *target == 0xFF && *target.add(1) == 0x25 {
        return Err("this function is already hooked".into());
    }

    // Copy whole instructions: never cut one in half.
    let window = std::slice::from_raw_parts(target, 64);
    let mut decoder = Decoder::with_ip(64, window, target as u64, DecoderOptions::NONE);
    let mut instructions: Vec<Instruction> = Vec::new();
    let mut stolen = 0usize;
    while stolen < JMP_LEN {
        let insn = decoder.decode();
        if insn.is_invalid() {
            return Err("could not decode the function prologue".into());
        }
        // The encoder rewrites RIP-relative displacements and branch targets
        // for the new address, so both relocate cleanly. What cannot be moved
        // is a branch *into* the bytes we are stealing — its target would end
        // up in the middle of our jump — or a function that simply ends inside
        // the range.
        match insn.flow_control() {
            FlowControl::Return | FlowControl::Interrupt | FlowControl::Exception => {
                return Err(format!("function ends inside the first {JMP_LEN} bytes ({insn})"));
            }
            FlowControl::IndirectBranch | FlowControl::IndirectCall => {
                // Fine unless it is RIP-relative, which the encoder handles.
            }
            _ => {}
        }
        if insn.is_jcc_short_or_near()
            || insn.is_jmp_short_or_near()
            || insn.is_call_near()
        {
            let dest = insn.near_branch_target();
            let base = target as u64;
            if dest >= base && dest < base + JMP_LEN as u64 {
                return Err(format!(
                    "a branch at +{stolen} jumps into the bytes we replace ({insn})"
                ));
            }
        }
        instructions.push(insn);
        stolen += insn.len();
    }

    // The trampoline must sit within ±2 GB of the original code: a relocated
    // RIP-relative operand still addresses its target with a 32-bit
    // displacement, and that displacement is now measured from here.
    let capacity = stolen * 2 + JMP_LEN + 32;
    let tramp = alloc_near(target as u64, capacity);
    if tramp.is_null() {
        return Err("no free page within 2 GB of the target".into());
    }

    // Re-encode at the new address so every displacement is corrected.
    let block = InstructionBlock::new(&instructions, tramp as u64);
    let encoded = BlockEncoder::encode(64, block, BlockEncoderOptions::NONE)
        .map_err(|e| format!("could not relocate the prologue: {e}"))?
        .code_buffer;
    if encoded.len() + JMP_LEN > capacity {
        return Err("relocated prologue does not fit the trampoline".into());
    }
    std::ptr::copy_nonoverlapping(encoded.as_ptr(), tramp, encoded.len());
    let back = abs_jmp(target as u64 + stolen as u64);
    std::ptr::copy_nonoverlapping(back.as_ptr(), tramp.add(encoded.len()), JMP_LEN);

    // Patch the target.
    let mut original = [0u8; JMP_LEN];
    std::ptr::copy_nonoverlapping(target, original.as_mut_ptr(), JMP_LEN);

    let jmp = abs_jmp(detour as u64);
    let Some(old) = mem::unprotect(target, JMP_LEN) else {
        return Err("could not make the target writable".into());
    };
    std::ptr::copy_nonoverlapping(jmp.as_ptr(), target, JMP_LEN);
    mem::reprotect(target, JMP_LEN, old);
    mem::flush(target, JMP_LEN);

    Ok(Hook { target, original, trampoline: tramp as *const c_void })
}

impl Hook {
    /// Put the original bytes back. Callers must then give in-flight threads a
    /// moment to leave the detour before unloading the code it lives in.
    ///
    /// # Safety
    /// Only valid if the target's first bytes are still our jump.
    pub unsafe fn remove(&self) {
        if let Some(old) = mem::unprotect(self.target, JMP_LEN) {
            std::ptr::copy_nonoverlapping(self.original.as_ptr(), self.target, JMP_LEN);
            mem::reprotect(self.target, JMP_LEN, old);
            mem::flush(self.target, JMP_LEN);
        }
    }
}

/// Reserve a page within ±2 GB of `target`, walking outwards in allocation
/// granularity steps until one is free.
unsafe fn alloc_near(target: u64, size: usize) -> *mut u8 {
    let granularity = mem::granularity();
    let reach = 0x7FFF_0000u64;
    let base = target & !(granularity - 1);

    let mut offset = granularity;
    while offset < reach {
        if let Some(addr) = base.checked_sub(offset) {
            let p = mem::reserve_at(addr, size);
            if !p.is_null() {
                return p;
            }
        }
        let p = mem::reserve_at(base + offset, size);
        if !p.is_null() {
            return p;
        }
        offset += granularity;
    }
    std::ptr::null_mut()
}