Sign in Sign up
kretrod/lodestone Public
Branches
master
lodestone Lodestone client: injected in-game menu driven through JNI a6eba02 · 3 days ago History
216 lines (196 loc) · 7.9 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.

use std::ffi::c_void;

use iced_x86::{
    BlockEncoder, BlockEncoderOptions, Decoder, DecoderOptions, FlowControl, Instruction,
    InstructionBlock,
};
use windows_sys::Win32::System::Memory::{
    VirtualAlloc, VirtualProtect, MEM_COMMIT, MEM_RESERVE, PAGE_EXECUTE_READWRITE,
};
use windows_sys::Win32::System::SystemInformation::{GetSystemInfo, SYSTEM_INFO};

const JMP_LEN: usize = 14;

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 mut old = 0u32;
    if VirtualProtect(target as *const c_void, JMP_LEN, PAGE_EXECUTE_READWRITE, &mut old) == 0 {
        return Err("VirtualProtect failed".into());
    }
    std::ptr::copy_nonoverlapping(jmp.as_ptr(), target, JMP_LEN);
    VirtualProtect(target as *const c_void, JMP_LEN, old, &mut old);
    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) {
        let mut old = 0u32;
        if VirtualProtect(
            self.target as *const c_void,
            JMP_LEN,
            PAGE_EXECUTE_READWRITE,
            &mut old,
        ) != 0
        {
            std::ptr::copy_nonoverlapping(self.original.as_ptr(), self.target, JMP_LEN);
            VirtualProtect(self.target as *const c_void, JMP_LEN, old, &mut old);
            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 mut info: SYSTEM_INFO = std::mem::zeroed();
    GetSystemInfo(&mut info);
    let granularity = info.dwAllocationGranularity.max(0x1000) as u64;
    let reach = 0x7FFF_0000u64;
    let base = target & !(granularity - 1);

    let try_at = |addr: u64| -> *mut u8 {
        VirtualAlloc(
            addr as *const c_void,
            size,
            MEM_COMMIT | MEM_RESERVE,
            PAGE_EXECUTE_READWRITE,
        ) as *mut u8
    };

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

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