From f244967b488e4e22280b1a441d3dd5c1d931775e Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 11 Feb 2026 03:08:10 +0000 Subject: [PATCH 1/3] Initial plan From c68900a67e034debbc6c14618a7d85142dddbbcc Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 11 Feb 2026 03:19:32 +0000 Subject: [PATCH 2/3] feat: add unhook() API with full restore and trampoline cleanup Review and implement unhook code from codex/unhook-smoke-ci branch: 1. Add `unhook()` public API to restore instrument and inline_hook patches 2. Add `HookNotFound` error variant for unhook on unknown addresses 3. Add `InlinePatchSlot` storage in state.rs for inline_hook original bytes 4. Add `free_original_trampoline` to properly munmap trampoline pages 5. Refactor `inline_hook` to save original bytes before patching (enables restore) 6. Improve `instrument_internal` x86_64: pad full instruction with int3+NOPs, derive original_opcode from original_bytes instead of redundant read 7. Remove now-unused `patch_far_jump` and `patch_u8` from memory.rs 8. Make `read_bytes` available on all platforms (needed by aarch64 inline_hook) 9. Add instrument_unhook_restore example with CI smoke tests Co-authored-by: YinMo19 <144041694+YinMo19@users.noreply.github.com> --- .github/workflows/ci.yml | 31 +++- Cargo.toml | 5 + examples/README.md | 1 + examples/instrument_unhook_restore/README.md | 45 +++++ examples/instrument_unhook_restore/main.rs | 85 ++++++++++ examples/instrument_unhook_restore/target.c | 67 ++++++++ src/error.rs | 2 + src/lib.rs | 164 +++++++++++++++---- src/memory.rs | 26 +-- src/state.rs | 107 ++++++++++++ src/trampoline.rs | 16 ++ 11 files changed, 487 insertions(+), 62 deletions(-) create mode 100644 examples/instrument_unhook_restore/README.md create mode 100644 examples/instrument_unhook_restore/main.rs create mode 100644 examples/instrument_unhook_restore/target.c diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 36ed94c..2654ee4 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -59,6 +59,13 @@ jobs: echo "$OUT" [[ "$OUT" == *"calc(4, 5) = 99"* ]] + cc -O0 -fno-inline examples/instrument_unhook_restore/target.c -o examples/instrument_unhook_restore/app + DYLIB="$PWD/target/aarch64-apple-darwin/debug/examples/libinstrument_unhook_restore.dylib" + OUT=$(DYLD_INSERT_LIBRARIES="$DYLIB" examples/instrument_unhook_restore/app) + echo "$OUT" + [[ "$OUT" == *"hooked_calc(3, 4) = 123"* ]] + [[ "$OUT" == *"calc(3, 4) = 7"* ]] + cc -O0 -fno-inline examples/inline_hook_far/target.c -o examples/inline_hook_far/app DYLIB="$PWD/target/aarch64-apple-darwin/debug/examples/libinline_hook_far.dylib" OUT=$(DYLD_INSERT_LIBRARIES="$DYLIB" examples/inline_hook_far/app) @@ -102,10 +109,11 @@ jobs: - name: Cargo clippy run: cargo clippy --all-targets --target x86_64-apple-darwin -- -D warnings - - name: Build four base examples (macOS x86_64) + - name: Build five base examples (macOS x86_64) run: | cargo build --example instrument_with_original --target x86_64-apple-darwin cargo build --example instrument_no_original --target x86_64-apple-darwin + cargo build --example instrument_unhook_restore --target x86_64-apple-darwin cargo build --example inline_hook_far --target x86_64-apple-darwin cargo build --example patchcode_add_to_mul --target x86_64-apple-darwin @@ -126,6 +134,13 @@ jobs: echo "$OUT" [[ "$OUT" == *"calc(4, 5) = 99"* ]] + cc -O0 -fno-inline examples/instrument_unhook_restore/target.c -o examples/instrument_unhook_restore/app + DYLIB="$PWD/target/x86_64-apple-darwin/debug/examples/libinstrument_unhook_restore.dylib" + OUT=$(DYLD_INSERT_LIBRARIES="$DYLIB" examples/instrument_unhook_restore/app) + echo "$OUT" + [[ "$OUT" == *"hooked_calc(3, 4) = 123"* ]] + [[ "$OUT" == *"calc(3, 4) = 7"* ]] + cc -O0 -fno-inline examples/inline_hook_far/target.c -o examples/inline_hook_far/app DYLIB="$PWD/target/x86_64-apple-darwin/debug/examples/libinline_hook_far.dylib" OUT=$(DYLD_INSERT_LIBRARIES="$DYLIB" examples/inline_hook_far/app) @@ -188,6 +203,13 @@ jobs: echo "$OUT" [[ "$OUT" == *"calc(4, 5) = 99"* ]] + cc -O0 -fno-inline -rdynamic examples/instrument_unhook_restore/target.c -o examples/instrument_unhook_restore/app + SO="$PWD/target/x86_64-unknown-linux-gnu/debug/examples/libinstrument_unhook_restore.so" + OUT=$(LD_PRELOAD="$SO" examples/instrument_unhook_restore/app) + echo "$OUT" + [[ "$OUT" == *"hooked_calc(3, 4) = 123"* ]] + [[ "$OUT" == *"calc(3, 4) = 7"* ]] + cc -O0 -fno-inline -rdynamic examples/inline_hook_far/target.c -o examples/inline_hook_far/app SO="$PWD/target/x86_64-unknown-linux-gnu/debug/examples/libinline_hook_far.so" OUT=$(LD_PRELOAD="$SO" examples/inline_hook_far/app) @@ -256,6 +278,13 @@ jobs: echo "$OUT" [[ "$OUT" == *"calc(4, 5) = 99"* ]] + cc -O0 -fno-inline -rdynamic examples/instrument_unhook_restore/target.c -o examples/instrument_unhook_restore/app + SO="$PWD/target/aarch64-unknown-linux-gnu/debug/examples/libinstrument_unhook_restore.so" + OUT=$(LD_PRELOAD="$SO" examples/instrument_unhook_restore/app) + echo "$OUT" + [[ "$OUT" == *"hooked_calc(3, 4) = 123"* ]] + [[ "$OUT" == *"calc(3, 4) = 7"* ]] + cc -O0 -fno-inline -rdynamic examples/instrument_adrp_no_original/target.c -o examples/instrument_adrp_no_original/app SO="$PWD/target/aarch64-unknown-linux-gnu/debug/examples/libinstrument_adrp_no_original.so" OUT=$(LD_PRELOAD="$SO" examples/instrument_adrp_no_original/app) diff --git a/Cargo.toml b/Cargo.toml index d3b5533..6f5969a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -58,6 +58,11 @@ name = "instrument_adrp_no_original" path = "examples/instrument_adrp_no_original/main.rs" crate-type = ["cdylib"] +[[example]] +name = "instrument_unhook_restore" +path = "examples/instrument_unhook_restore/main.rs" +crate-type = ["cdylib"] + [[example]] name = "inline_hook_far" path = "examples/inline_hook_far/main.rs" diff --git a/examples/README.md b/examples/README.md index 788ce9b..66b4b87 100644 --- a/examples/README.md +++ b/examples/README.md @@ -30,6 +30,7 @@ Available examples: - `instrument_with_original`: BRK instrumentation + execute original opcode - `instrument_no_original`: BRK instrumentation + skip original opcode - `instrument_adrp_no_original`: aarch64 `adrp` patch-point via `instrument_no_original` + manual callback emulation +- `instrument_unhook_restore`: instrument + unhook restore demo - `inline_hook_far`: function-entry detour with inline hook ## Coverage matrix diff --git a/examples/instrument_unhook_restore/README.md b/examples/instrument_unhook_restore/README.md new file mode 100644 index 0000000..71fb356 --- /dev/null +++ b/examples/instrument_unhook_restore/README.md @@ -0,0 +1,45 @@ +# instrument_unhook_restore + +Demonstrates `sighook::instrument` + `sighook::unhook`. + +Flow: + +1. install an instruction hook +2. call `calc(3, 4)` once while hook is active and verify hooked value `123` +3. call `unhook` on the same patchpoint +4. execute target function again and verify original behavior is restored (`7`) + +The callback would force result to `123` if triggered: + +- `aarch64`: set `x8 = 120`, `x9 = 3`, then original `add w0, w8, w9` runs +- `linux x86_64`: set `rax = 123` + +Expected runtime output proves both stages: + +- callback is reached before unhook (`hooked_calc(3, 4) = 123`) +- original behavior is restored after unhook (`calc(3, 4) = 7`) + +## Run (from repository root) + +macOS: + +```bash +cc -O0 -fno-inline examples/instrument_unhook_restore/target.c -o examples/instrument_unhook_restore/app +cargo build --example instrument_unhook_restore +DYLD_INSERT_LIBRARIES="$PWD/target/debug/examples/libinstrument_unhook_restore.dylib" examples/instrument_unhook_restore/app +``` + +Linux: + +```bash +cc -O0 -fno-inline -rdynamic examples/instrument_unhook_restore/target.c -o examples/instrument_unhook_restore/app +cargo build --example instrument_unhook_restore +LD_PRELOAD="$PWD/target/debug/examples/libinstrument_unhook_restore.so" examples/instrument_unhook_restore/app +``` + +Expected output: + +```text +hooked_calc(3, 4) = 123 +calc(3, 4) = 7 +``` diff --git a/examples/instrument_unhook_restore/main.rs b/examples/instrument_unhook_restore/main.rs new file mode 100644 index 0000000..a67ea19 --- /dev/null +++ b/examples/instrument_unhook_restore/main.rs @@ -0,0 +1,85 @@ +use sighook::{HookContext, instrument, unhook}; + +#[cfg(all(any(target_os = "macos", target_os = "ios"), target_arch = "aarch64"))] +const ADD_INSN_OFFSET: u64 = 0x14; + +#[cfg(all(target_arch = "x86_64", any(target_os = "linux", target_os = "macos")))] +const X86_PATCHPOINT_OFFSET: u64 = 0x4; + +extern "C" fn on_hit_should_not_run(_address: u64, ctx: *mut HookContext) { + unsafe { + #[cfg(target_arch = "aarch64")] + { + (*ctx).regs.named.x8 = 120; + (*ctx).regs.named.x9 = 3; + } + + #[cfg(all(target_arch = "x86_64", any(target_os = "linux", target_os = "macos")))] + { + (*ctx).rax = 123; + } + } +} + +#[used] +#[cfg_attr( + any(target_os = "macos", target_os = "ios"), + unsafe(link_section = "__DATA,__mod_init_func") +)] +#[cfg_attr( + any(target_os = "linux", target_os = "android"), + unsafe(link_section = ".init_array") +)] +static INIT_ARRAY: extern "C" fn() = init; + +extern "C" fn init() { + unsafe { + let calc_symbol = libc::dlsym(libc::RTLD_DEFAULT, c"calc".as_ptr()); + if calc_symbol.is_null() { + return; + } + let calc_fn: extern "C" fn(i32, i32) -> i32 = std::mem::transmute(calc_symbol); + + let target_address = { + #[cfg(all( + any(target_os = "linux", target_os = "android"), + target_arch = "aarch64" + ))] + { + let symbol = libc::dlsym(libc::RTLD_DEFAULT, c"calc_add_insn".as_ptr()); + if symbol.is_null() { + return; + } + symbol as u64 + } + + #[cfg(all(target_arch = "x86_64", any(target_os = "linux", target_os = "macos")))] + { + let symbol = libc::dlsym(libc::RTLD_DEFAULT, c"calc".as_ptr()); + if symbol.is_null() { + return; + } + symbol as u64 + X86_PATCHPOINT_OFFSET + } + + #[cfg(all(any(target_os = "macos", target_os = "ios"), target_arch = "aarch64"))] + { + let symbol = libc::dlsym(libc::RTLD_DEFAULT, c"calc".as_ptr()); + if symbol.is_null() { + return; + } + symbol as u64 + ADD_INSN_OFFSET + } + }; + + let _ = instrument(target_address, on_hit_should_not_run); + + let hooked = calc_fn(3, 4); + println!("hooked_calc(3, 4) = {hooked}"); + if hooked != 123 { + return; + } + + let _ = unhook(target_address); + } +} diff --git a/examples/instrument_unhook_restore/target.c b/examples/instrument_unhook_restore/target.c new file mode 100644 index 0000000..4108a57 --- /dev/null +++ b/examples/instrument_unhook_restore/target.c @@ -0,0 +1,67 @@ +#include + +#if defined(__linux__) && defined(__aarch64__) +int calc(int a, int b); + +__asm__( + ".text\n" + ".global calc\n" + ".global calc_add_insn\n" + ".type calc, %function\n" + "calc:\n" + " mov x8, x0\n" + " mov x9, x1\n" + " nop\n" + " nop\n" + " nop\n" + "calc_add_insn:\n" + " add w0, w8, w9\n" + " ret\n" + ".size calc, .-calc\n"); +#elif defined(__x86_64__) && defined(__APPLE__) +int calc(int a, int b); + +__asm__( + ".text\n" + ".globl _calc\n" + "_calc:\n" + " mov %edi, %eax\n" + " add %esi, %eax\n" + " nop\n" + " ret\n"); +#elif defined(__x86_64__) && defined(__linux__) +int calc(int a, int b); + +__asm__( + ".text\n" + ".global calc\n" + ".type calc, @function\n" + "calc:\n" + " mov %edi, %eax\n" + " add %esi, %eax\n" + " nop\n" + " ret\n" + ".size calc, .-calc\n"); +#elif defined(__aarch64__) +__attribute__((naked, noinline)) +int calc(int a, int b) { + __asm__ volatile( + "mov x8, x0\n" + "mov x9, x1\n" + "nop\n" + "nop\n" + "nop\n" + "add w0, w8, w9\n" + "ret\n"); +} +#else +__attribute__((noinline)) +int calc(int a, int b) { + return a + b; +} +#endif + +int main(void) { + printf("calc(3, 4) = %d\n", calc(3, 4)); + return 0; +} diff --git a/src/error.rs b/src/error.rs index 3fba8b6..b32cddd 100644 --- a/src/error.rs +++ b/src/error.rs @@ -4,6 +4,7 @@ use std::fmt; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum SigHookError { InvalidAddress, + HookNotFound, UnsupportedPlatform, UnsupportedArchitecture, UnsupportedOperation, @@ -78,6 +79,7 @@ impl fmt::Display for SigHookError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match *self { SigHookError::InvalidAddress => write!(f, "invalid address"), + SigHookError::HookNotFound => write!(f, "hook not found at address"), SigHookError::UnsupportedPlatform => write!(f, "unsupported platform"), SigHookError::UnsupportedArchitecture => write!(f, "unsupported architecture"), SigHookError::UnsupportedOperation => write!(f, "unsupported operation"), diff --git a/src/lib.rs b/src/lib.rs index 9c7c822..63a2676 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -196,6 +196,9 @@ pub fn patch_asm(address: u64, asm: &str) -> Result { /// If the callback does not redirect control flow (`pc`/`rip` unchanged), /// the original instruction runs through an internal trampoline, then execution continues. /// +/// On `x86_64`, trap patching writes `int3` at instruction start and pads the +/// remaining bytes of that decoded instruction with `NOP`. +/// /// # PC-relative note /// /// Across architectures, this API does **not** support patch points whose original @@ -267,7 +270,10 @@ fn instrument_internal( callback, execute_original, )?; - return state::original_opcode_by_address(address).ok_or(SigHookError::InvalidAddress); + let original_opcode = state::cached_original_opcode_by_address(address) + .or_else(|| state::original_opcode_by_address(address)) + .ok_or(SigHookError::InvalidAddress)?; + return Ok(original_opcode); } signal::ensure_handlers_installed()?; @@ -275,16 +281,23 @@ fn instrument_internal( let step_len: u8 = memory::instruction_width(address)?; #[cfg(target_arch = "aarch64")] - let original_bytes = { + let (original_bytes, original_opcode) = { let original = memory::patch_u32(address, constants::BRK_OPCODE)?; - original.to_le_bytes().to_vec() + (original.to_le_bytes().to_vec(), original) }; #[cfg(all(target_arch = "x86_64", any(target_os = "linux", target_os = "macos")))] - let original_bytes = { + let (original_bytes, original_opcode) = { let original_bytes = memory::read_bytes(address, step_len as usize)?; - let _ = memory::patch_u8(address, memory::int3_opcode())?; - original_bytes + + let mut trap_patch = vec![0x90u8; step_len as usize]; + trap_patch[0] = memory::int3_opcode(); + let _ = memory::patch_bytes_public(address, &trap_patch)?; + + let mut opcode = [0u8; 4]; + let copy_len = original_bytes.len().min(4); + opcode[..copy_len].copy_from_slice(&original_bytes[..copy_len]); + (original_bytes, u32::from_le_bytes(opcode)) }; let register_result = state::register_slot( @@ -312,13 +325,6 @@ fn instrument_internal( return Err(err); } - if original_bytes.len() < 4 { - return Err(SigHookError::InvalidAddress); - } - - let mut bytes = [0u8; 4]; - bytes.copy_from_slice(&original_bytes[..4]); - let original_opcode = u32::from_le_bytes(bytes); state::cache_original_opcode(address, original_opcode); Ok(original_opcode) @@ -349,48 +355,134 @@ fn instrument_internal( pub fn inline_hook(addr: u64, replace_fn: u64) -> Result { #[cfg(target_arch = "aarch64")] { - match memory::encode_b(addr, replace_fn) { - Ok(b_opcode) => patchcode(addr, b_opcode), + let patch = match memory::encode_b(addr, replace_fn) { + Ok(b_opcode) => b_opcode.to_le_bytes().to_vec(), Err(SigHookError::BranchOutOfRange) => { - let original = memory::patch_far_jump(addr, replace_fn)?; + let mut bytes = [0u8; 16]; + bytes[0..4].copy_from_slice(&constants::LDR_X16_LITERAL_8.to_le_bytes()); + bytes[4..8].copy_from_slice(&constants::BR_X16.to_le_bytes()); + bytes[8..16].copy_from_slice(&replace_fn.to_le_bytes()); + bytes.to_vec() + } + Err(err) => return Err(err), + }; + + let original = memory::read_bytes(addr, patch.len())?; + let inserted = unsafe { state::cache_inline_patch(addr, &original)? }; + if let Err(err) = memory::patch_bytes_public(addr, &patch) { + if inserted { unsafe { - state::cache_original_opcode(addr, original); + state::remove_inline_patch(addr); } - Ok(original) } - Err(err) => Err(err), + return Err(err); + } + + if original.len() < 4 { + return Err(SigHookError::InvalidAddress); + } + + let mut opcode = [0u8; 4]; + opcode.copy_from_slice(&original[..4]); + let original_opcode = u32::from_le_bytes(opcode); + unsafe { + state::cache_original_opcode(addr, original_opcode); } + Ok(original_opcode) } #[cfg(all(target_arch = "x86_64", any(target_os = "linux", target_os = "macos")))] { - if let Ok(jmp) = memory::encode_jmp_rel32(addr, replace_fn) { - let original = memory::patch_bytes_public(addr, &jmp)?; - let mut opcode = [0u8; 4]; - if original.len() >= 4 { - opcode.copy_from_slice(&original[..4]); - let original_opcode = u32::from_le_bytes(opcode); + let patch = if let Ok(jmp) = memory::encode_jmp_rel32(addr, replace_fn) { + jmp.to_vec() + } else { + memory::encode_absolute_jump(replace_fn).to_vec() + }; + + let original = memory::read_bytes(addr, patch.len())?; + let inserted = unsafe { state::cache_inline_patch(addr, &original)? }; + if let Err(err) = memory::patch_bytes_public(addr, &patch) { + if inserted { unsafe { - state::cache_original_opcode(addr, original_opcode); + state::remove_inline_patch(addr); } - return Ok(original_opcode); } + return Err(err); + } + + if original.len() < 4 { return Err(SigHookError::InvalidAddress); } - let abs = memory::encode_absolute_jump(replace_fn); - let original = memory::patch_bytes_public(addr, &abs)?; let mut opcode = [0u8; 4]; - if original.len() >= 4 { - opcode.copy_from_slice(&original[..4]); - let original_opcode = u32::from_le_bytes(opcode); - unsafe { - state::cache_original_opcode(addr, original_opcode); + opcode.copy_from_slice(&original[..4]); + let original_opcode = u32::from_le_bytes(opcode); + unsafe { + state::cache_original_opcode(addr, original_opcode); + } + Ok(original_opcode) + } +} + +/// Restores a previously installed hook at `address`. +/// +/// This API supports hook points created by [`instrument`], [`instrument_no_original`], +/// and [`inline_hook`]. On success, the patched instruction bytes are restored and +/// internal runtime state for that address is removed. +/// +/// Signal handlers stay installed once initialized, even after unhooking all addresses. +/// +/// # Example +/// +/// ```rust,ignore +/// use sighook::{instrument, unhook, HookContext}; +/// +/// extern "C" fn on_hit(_address: u64, _ctx: *mut HookContext) {} +/// +/// let addr = 0x1000_0000u64; +/// let _ = instrument(addr, on_hit)?; +/// unhook(addr)?; +/// # Ok::<(), sighook::SigHookError>(()) +/// ``` +pub fn unhook(address: u64) -> Result<(), SigHookError> { + if address == 0 { + return Err(SigHookError::InvalidAddress); + } + + unsafe { + if let Some(slot) = state::slot_by_address(address) { + if slot.original_len == 0 { + return Err(SigHookError::InvalidAddress); } - return Ok(original_opcode); + + memory::patch_bytes_public( + address, + &slot.original_bytes[..slot.original_len as usize], + )?; + + if let Some(removed_slot) = state::remove_slot(address) { + if removed_slot.trampoline_pc != 0 { + trampoline::free_original_trampoline(removed_slot.trampoline_pc); + } + } + + state::remove_cached_original_opcode(address); + return Ok(()); + } + + if let Some((bytes, len)) = state::inline_patch_by_address(address) { + if len == 0 { + return Err(SigHookError::InvalidAddress); + } + + memory::patch_bytes_public(address, &bytes[..len as usize])?; + state::remove_inline_patch(address); + state::remove_cached_original_opcode(address); + return Ok(()); } - Err(SigHookError::InvalidAddress) } + + Err(SigHookError::HookNotFound) } /// Returns the saved original 4-byte value for a previously patched address. diff --git a/src/memory.rs b/src/memory.rs index fe3ce56..d63c9a8 100644 --- a/src/memory.rs +++ b/src/memory.rs @@ -1,7 +1,7 @@ #[cfg(any(target_os = "macos", target_os = "ios"))] use crate::constants::VM_PROT_COPY; #[cfg(target_arch = "aarch64")] -use crate::constants::{BR_X16, BRK_MASK, BRK_OPCODE, LDR_X16_LITERAL_8}; +use crate::constants::{BRK_MASK, BRK_OPCODE}; use crate::error::SigHookError; use libc::{c_int, c_void}; @@ -248,13 +248,6 @@ fn patch_bytes(address: u64, bytes: &[u8]) -> Result, SigHookError> { Ok(original) } -#[cfg(all(target_arch = "x86_64", any(target_os = "linux", target_os = "macos")))] -pub(crate) fn patch_u8(address: u64, new_opcode: u8) -> Result { - let original = patch_bytes(address, &[new_opcode])?; - Ok(original[0]) -} - -#[cfg(all(target_arch = "x86_64", any(target_os = "linux", target_os = "macos")))] pub(crate) fn read_bytes(address: u64, len: usize) -> Result, SigHookError> { if address == 0 || len == 0 { return Err(SigHookError::InvalidAddress); @@ -326,23 +319,6 @@ pub(crate) fn encode_b(from_address: u64, to_address: u64) -> Result Result { - if (from_address & 0b11) != 0 { - return Err(SigHookError::InvalidAddress); - } - - let mut bytes = [0u8; 16]; - bytes[0..4].copy_from_slice(&LDR_X16_LITERAL_8.to_le_bytes()); - bytes[4..8].copy_from_slice(&BR_X16.to_le_bytes()); - bytes[8..16].copy_from_slice(&to_address.to_le_bytes()); - - let original = patch_bytes(from_address, &bytes)?; - let mut opcode_bytes = [0u8; 4]; - opcode_bytes.copy_from_slice(&original[0..4]); - Ok(u32::from_le_bytes(opcode_bytes)) -} - #[cfg(all(target_arch = "x86_64", any(target_os = "linux", target_os = "macos")))] pub(crate) fn encode_jmp_rel32( from_address: u64, diff --git a/src/state.rs b/src/state.rs index 95ada67..796899a 100644 --- a/src/state.rs +++ b/src/state.rs @@ -32,6 +32,26 @@ pub(crate) static mut HANDLERS_INSTALLED: bool = false; pub(crate) static mut SLOTS: [InstrumentSlot; MAX_INSTRUMENTS] = [InstrumentSlot::EMPTY; MAX_INSTRUMENTS]; +#[derive(Copy, Clone)] +pub(crate) struct InlinePatchSlot { + pub used: bool, + pub address: u64, + pub original_bytes: [u8; 16], + pub original_len: u8, +} + +impl InlinePatchSlot { + pub const EMPTY: Self = Self { + used: false, + address: 0, + original_bytes: [0u8; 16], + original_len: 0, + }; +} + +pub(crate) static mut INLINE_PATCH_SLOTS: [InlinePatchSlot; MAX_INSTRUMENTS] = + [InlinePatchSlot::EMPTY; MAX_INSTRUMENTS]; + #[derive(Copy, Clone)] pub(crate) struct OriginalOpcodeSlot { pub used: bool, @@ -68,6 +88,15 @@ pub(crate) unsafe fn slot_by_address(address: u64) -> Option { Some(unsafe { SLOTS[index] }) } +pub(crate) unsafe fn remove_slot(address: u64) -> Option { + let index = unsafe { find_slot_index(address) }?; + let slot = unsafe { SLOTS[index] }; + unsafe { + SLOTS[index] = InstrumentSlot::EMPTY; + } + Some(slot) +} + pub(crate) unsafe fn register_slot( address: u64, original_bytes: &[u8], @@ -206,3 +235,81 @@ pub(crate) unsafe fn cached_original_opcode_by_address(address: u64) -> Option bool { + let index = match unsafe { find_original_opcode_slot_index(address) } { + Some(index) => index, + None => return false, + }; + + unsafe { + ORIGINAL_OPCODE_SLOTS[index] = OriginalOpcodeSlot::EMPTY; + } + + true +} + +unsafe fn find_inline_patch_slot_index(address: u64) -> Option { + let mut index = 0; + while index < MAX_INSTRUMENTS { + let slot = unsafe { INLINE_PATCH_SLOTS[index] }; + if slot.used && slot.address == address { + return Some(index); + } + index += 1; + } + None +} + +pub(crate) unsafe fn cache_inline_patch( + address: u64, + original_bytes: &[u8], +) -> Result { + if original_bytes.is_empty() || original_bytes.len() > 16 { + return Err(SigHookError::InvalidAddress); + } + + if unsafe { find_inline_patch_slot_index(address) }.is_some() { + return Ok(false); + } + + let mut stored_bytes = [0u8; 16]; + stored_bytes[..original_bytes.len()].copy_from_slice(original_bytes); + + let mut index = 0; + while index < MAX_INSTRUMENTS { + if !(unsafe { INLINE_PATCH_SLOTS[index].used }) { + unsafe { + INLINE_PATCH_SLOTS[index] = InlinePatchSlot { + used: true, + address, + original_bytes: stored_bytes, + original_len: original_bytes.len() as u8, + }; + } + return Ok(true); + } + index += 1; + } + + Err(SigHookError::InstrumentSlotsFull) +} + +pub(crate) unsafe fn inline_patch_by_address(address: u64) -> Option<([u8; 16], u8)> { + let index = unsafe { find_inline_patch_slot_index(address) }?; + let slot = unsafe { INLINE_PATCH_SLOTS[index] }; + Some((slot.original_bytes, slot.original_len)) +} + +pub(crate) unsafe fn remove_inline_patch(address: u64) -> bool { + let index = match unsafe { find_inline_patch_slot_index(address) } { + Some(index) => index, + None => return false, + }; + + unsafe { + INLINE_PATCH_SLOTS[index] = InlinePatchSlot::EMPTY; + } + + true +} diff --git a/src/trampoline.rs b/src/trampoline.rs index b979509..7515adc 100644 --- a/src/trampoline.rs +++ b/src/trampoline.rs @@ -2,6 +2,7 @@ use crate::constants::{BR_X16, LDR_X16_LITERAL_8}; use crate::error::SigHookError; use crate::memory::{flush_instruction_cache, last_errno}; +use libc::c_void; use std::ptr::null_mut; pub(crate) fn create_original_trampoline( @@ -96,6 +97,21 @@ pub(crate) fn create_original_trampoline( Ok(base as u64) } +pub(crate) unsafe fn free_original_trampoline(trampoline_pc: u64) { + if trampoline_pc == 0 { + return; + } + + let page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) }; + if page_size <= 0 { + return; + } + + unsafe { + let _ = libc::munmap(trampoline_pc as *mut c_void, page_size as usize); + } +} + #[cfg(all(target_arch = "x86_64", any(target_os = "linux", target_os = "macos")))] fn encode_abs_jmp_indirect(to_address: u64) -> [u8; 14] { let mut bytes = [0u8; 14]; From 0799e97d6c7c79ce0c20cc12640116018f0107aa Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 11 Feb 2026 03:21:07 +0000 Subject: [PATCH 3/3] review: address code review feedback - remove redundant unsafe, rename callback Co-authored-by: YinMo19 <144041694+YinMo19@users.noreply.github.com> --- examples/instrument_unhook_restore/main.rs | 4 ++-- src/trampoline.rs | 4 +--- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/examples/instrument_unhook_restore/main.rs b/examples/instrument_unhook_restore/main.rs index a67ea19..b65d7a3 100644 --- a/examples/instrument_unhook_restore/main.rs +++ b/examples/instrument_unhook_restore/main.rs @@ -6,7 +6,7 @@ const ADD_INSN_OFFSET: u64 = 0x14; #[cfg(all(target_arch = "x86_64", any(target_os = "linux", target_os = "macos")))] const X86_PATCHPOINT_OFFSET: u64 = 0x4; -extern "C" fn on_hit_should_not_run(_address: u64, ctx: *mut HookContext) { +extern "C" fn hook_callback(_address: u64, ctx: *mut HookContext) { unsafe { #[cfg(target_arch = "aarch64")] { @@ -72,7 +72,7 @@ extern "C" fn init() { } }; - let _ = instrument(target_address, on_hit_should_not_run); + let _ = instrument(target_address, hook_callback); let hooked = calc_fn(3, 4); println!("hooked_calc(3, 4) = {hooked}"); diff --git a/src/trampoline.rs b/src/trampoline.rs index 7515adc..73d26bf 100644 --- a/src/trampoline.rs +++ b/src/trampoline.rs @@ -107,9 +107,7 @@ pub(crate) unsafe fn free_original_trampoline(trampoline_pc: u64) { return; } - unsafe { - let _ = libc::munmap(trampoline_pc as *mut c_void, page_size as usize); - } + unsafe { libc::munmap(trampoline_pc as *mut c_void, page_size as usize) }; } #[cfg(all(target_arch = "x86_64", any(target_os = "linux", target_os = "macos")))]