diff --git a/Cargo.lock b/Cargo.lock index 2ce2088..8e15321 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -54,9 +54,16 @@ version = "0.1.0" dependencies = [ "pelite", "rayon", + "sharded-slab", "windows", ] +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + [[package]] name = "libc" version = "0.2.158" @@ -126,6 +133,15 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + [[package]] name = "syn" version = "2.0.76" diff --git a/Cargo.toml b/Cargo.toml index a48b4c0..3cab4d5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,6 +11,7 @@ description = "todo" [dependencies] rayon = "1.10.0" +sharded-slab = "0.1.7" [target.'cfg(windows)'.dependencies.pelite] version = "0.10.0" diff --git a/examples/sections.rs b/examples/sections.rs index 1e911c7..97f970e 100644 --- a/examples/sections.rs +++ b/examples/sections.rs @@ -1,5 +1,12 @@ -use inka::program; +use inka::Hook; fn main() { - dbg!(program()); + unsafe extern "C" fn hook_me() { + println!("Original"); + } + + let _guard = (hook_me as unsafe extern "C" fn()).hook(|| println!("Hooked")); + dbg!(_guard); + + unsafe { hook_me() }; } diff --git a/src/hook.rs b/src/hook.rs new file mode 100644 index 0000000..49dca72 --- /dev/null +++ b/src/hook.rs @@ -0,0 +1,55 @@ +use core::any::Any; +use sharded_slab::Slab; +use std::{ + ptr::NonNull, + sync::{LazyLock, Mutex}, +}; + +static HOOKS: LazyLock>>> = LazyLock::new(Slab::new); + +pub trait Hook: Copy { + fn hook(self, f: F) -> HookGuard; + fn as_ptr_u8(self) -> NonNull; +} + +#[derive(Debug)] +pub struct HookGuard { + ptr: NonNull, + bytes: [u8; 10], // used to restore original function prologue + index: usize, +} + +impl HookGuard { + pub fn unhook(&mut self) { + HOOKS.remove(self.index); + } +} + +impl Drop for HookGuard { + fn drop(&mut self) { + self.unhook(); + } +} + +impl Hook for unsafe extern "C" fn() +where + F: FnMut() + 'static + Send, +{ + fn hook(self, f: F) -> HookGuard { + let ptr = >::as_ptr_u8(self); + + let index = HOOKS + .insert(Mutex::new(Box::new(f))) + .expect("Max shards reached"); + + HookGuard { + ptr, + bytes: [0; 10], + index, + } + } + + fn as_ptr_u8(self) -> NonNull { + unsafe { NonNull::new_unchecked(self as *mut u8) } + } +} diff --git a/src/lib.rs b/src/lib.rs index 97ccf0d..4ce8cd9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,12 +2,14 @@ mod base; mod find; +mod hook; mod program; mod section; mod symbol; pub use base::Base; pub use find::Find; +pub use hook::{Hook, HookGuard}; pub use program::{program, Program}; pub use section::Section; pub use symbol::Symbol;