Skip to main content

cradle_hooks/
engine.rs

1use crate::breakpoints::{Breakpoint, BreakpointMap, ModuleBaseCache};
2use crate::defs::{HookAction, HookContext, HookHandler};
3use crate::utils::{
4    get_thread_context, normalize_name, read_remote_or, read_single_byte, resolve_export,
5    set_thread_context, write_process_memory, write_single_byte,
6};
7use cradle_shared::{CradleError, CradleResult};
8use windows_sys::Win32::Foundation::HANDLE;
9
10/// Result of the hooking operation
11pub enum DispatchResult {
12    /// The hooked event was handled correctly
13    Handled,
14    /// No registered hook handled this event
15    Skip,
16}
17
18/// The core hooking engine
19/// Responsible for dispatching, registering, and removing hooks from the target process
20pub struct HookEngine {
21    breakpoints: BreakpointMap,
22    module_cache: ModuleBaseCache,
23    process: HANDLE,
24    pending_restore: Option<PendingRestore>,
25}
26
27struct PendingRestore {
28    addr: usize,
29    thread: HANDLE,
30}
31
32impl HookEngine {
33    /// Creates an empty [`HookEngine`] object
34    pub fn empty() -> HookEngine {
35        Self {
36            breakpoints: BreakpointMap::new(),
37            module_cache: ModuleBaseCache::new(),
38            process: HANDLE::default(),
39            pending_restore: None,
40        }
41    }
42}
43
44impl HookEngine {
45    /// Creates a new [`HookEngine`] object with the specified process handle
46    pub fn new(process: HANDLE) -> Self {
47        Self {
48            breakpoints: BreakpointMap::new(),
49            module_cache: ModuleBaseCache::new(),
50            process,
51            pending_restore: None,
52        }
53    }
54
55    /// Registers a module to hook given its name and base address
56    pub fn register_module(&mut self, name: &str, base: usize) {
57        self.module_cache.insert(normalize_name(name), base);
58    }
59
60    /// Dispatches a hook on the given handle
61    ///
62    /// # Safety
63    /// Calls `get_thread_context` to get the current registers. This function is unsafe and users should be aware of it
64    pub unsafe fn dispatch(&mut self, thread: HANDLE) -> CradleResult<DispatchResult> {
65        unsafe {
66            let mut ctx = get_thread_context(thread)?;
67            let bp_addr = (ctx.Rip - 1) as usize;
68            let bp = match self.breakpoints.get(&bp_addr) {
69                Some(bp) => bp,
70                None => return Ok(DispatchResult::Skip),
71            };
72            let mut hook_ctx = HookContext {
73                registers: &mut ctx,
74                process: self.process,
75                ret_value: 0,
76                export_name: &bp.export_name,
77                module_name: &bp.module_name,
78                #[cfg(feature = "unstable")]
79                return_hook: None,
80            };
81            let action = (bp.handler)(&mut hook_ctx);
82            let ret = hook_ctx.ret_value;
83            match action {
84                Ok(HookAction::Continue) | Ok(HookAction::Modify) => {
85                    write_process_memory(self.process, bp_addr, &[bp.original_byte])?;
86                    ctx.Rip = bp_addr as u64;
87                    ctx.EFlags |= 0x100;
88                    set_thread_context(thread, &ctx)?;
89                    self.pending_restore = Some(PendingRestore {
90                        addr: bp_addr,
91                        thread,
92                    });
93                }
94                Ok(HookAction::Block(u)) => {
95                    let ret_addr: u64 = read_remote_or(self.process, ctx.Rsp as usize, u);
96                    ctx.Rip = ret_addr;
97                    ctx.Rsp += 8;
98                    ctx.Rax = ret;
99                    set_thread_context(thread, &ctx)?;
100                }
101                Err(e) => return Err(e),
102            }
103            Ok(DispatchResult::Handled)
104        }
105    }
106
107    /// Dispatches a single breakpoint by writing it to the process
108    ///
109    /// # Safety
110    /// Calls `WriteProcessMemory` under the hood (via FFI). This may be unsafe
111    pub unsafe fn dispatch_single_step(&mut self, thread: HANDLE) -> CradleResult<DispatchResult> {
112        match self.pending_restore.take() {
113            Some(restore) if restore.thread == thread => {
114                if self.breakpoints.contains_key(&restore.addr) {
115                    unsafe {
116                        write_process_memory(self.process, restore.addr, &[0xCC])?;
117                    }
118                }
119                Ok(DispatchResult::Handled)
120            }
121            Some(restore) => {
122                self.pending_restore = Some(restore);
123                Ok(DispatchResult::Skip)
124            }
125            None => Ok(DispatchResult::Skip),
126        }
127    }
128
129    /// Attempts to hook an exported function in the specified module
130    /// Will return an error if a hook is already installed there
131    ///
132    /// # Safety
133    /// Reads and writes data to the process' memory, which is unsafe and should be treated as such
134    pub unsafe fn hook_export(
135        &mut self,
136        module: &str,
137        export: &str,
138        handler: HookHandler,
139    ) -> CradleResult {
140        if self.module_cache.is_empty() {
141            return Err(CradleError::ModuleNotFound(
142                "module cache is empty".to_owned(),
143            ));
144        }
145        let module_lower = normalize_name(module);
146        let base = self
147            .module_cache
148            .get(&module_lower)
149            .copied()
150            .ok_or_else(|| CradleError::ModuleNotFound(module.to_string()))?;
151        let addr = resolve_export(self.process, base, export)?;
152        if self.breakpoints.contains_key(&addr) {
153            return Err(CradleError::InvalidValue(format!(
154                "{module}!{export} already hooked at {addr:#x}",
155            )));
156        }
157        let mut orig = [0u8; 1];
158        unsafe {
159            read_single_byte(self.process, addr, &mut orig)?;
160            write_single_byte(self.process, addr, &[0xCC])?;
161        }
162        self.breakpoints.insert(
163            addr,
164            Breakpoint::new(orig[0], handler, module_lower, export.to_string()),
165        );
166
167        Ok(())
168    }
169    /// Attempts to unhook the breakpoint at the given address
170    ///
171    /// # Safety
172    /// Calls `WriteProcessMemory` under the hood (via FFI). This may be unsafe
173    pub unsafe fn unhook(&mut self, addr: usize) -> CradleResult {
174        if let Some(bp) = self.breakpoints.remove(&addr) {
175            unsafe {
176                write_process_memory(self.process, addr, &[bp.original_byte])?;
177            }
178        }
179        Ok(())
180    }
181
182    /// Attempts to unhook all existing hooks
183    ///
184    /// # Safety
185    /// Calls `self.unhook(addr)` on each hook. This function is unsafe
186    pub unsafe fn unhook_all(&mut self) -> CradleResult {
187        let addrs: Vec<usize> = self.breakpoints.keys().copied().collect();
188        for addr in addrs {
189            unsafe {
190                self.unhook(addr)?;
191            }
192        }
193        Ok(())
194    }
195
196    /// Returns the amount of breakpoints are currently installed
197    pub fn hooked_count(&self) -> usize {
198        self.breakpoints.len()
199    }
200}
201unsafe impl Send for HookEngine {}