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    /// Returns true if the address belongs to a registered breakpoint.
61    pub fn has_breakpoint(&self, addr: usize) -> bool {
62        self.breakpoints.contains_key(&addr)
63    }
64
65    /// Dispatches a hook on the given handle
66    ///
67    /// # Safety
68    /// Calls `get_thread_context` to get the current registers. This function is unsafe and users should be aware of it
69    pub unsafe fn dispatch(&mut self, thread: HANDLE) -> CradleResult<DispatchResult> {
70        unsafe {
71            let mut ctx = get_thread_context(thread).map_err(|e| {
72                CradleError::InvalidValue(format!("get thread context failed: {e}"))
73            })?;
74            let bp_addr = (ctx.Rip - 1) as usize;
75            let bp = match self.breakpoints.get(&bp_addr) {
76                Some(bp) => bp,
77                None => return Ok(DispatchResult::Skip),
78            };
79            let mut hook_ctx = HookContext {
80                registers: &mut ctx,
81                process: self.process,
82                ret_value: 0,
83                export_name: &bp.export_name,
84                module_name: &bp.module_name,
85                #[cfg(feature = "unstable")]
86                return_hook: None,
87            };
88            let mut action = HookAction::Continue;
89            for handler in &bp.handlers {
90                match handler(&mut hook_ctx) {
91                    Ok(HookAction::Continue) => {}
92                    Ok(HookAction::Modify) => action = HookAction::Modify,
93                    Ok(block @ HookAction::Block(_)) => {
94                        action = block;
95                        break;
96                    }
97                    Err(_) => {}
98                }
99            }
100            let ret = hook_ctx.ret_value;
101            match action {
102                HookAction::Continue | HookAction::Modify => {
103                    write_process_memory(self.process, bp_addr, &[bp.original_byte]).map_err(
104                        |e| {
105                            CradleError::InvalidValue(format!(
106                                "{}!{} restore byte at {bp_addr:#x} failed: {e}",
107                                bp.module_name, bp.export_name
108                            ))
109                        },
110                    )?;
111                    ctx.Rip = bp_addr as u64;
112                    ctx.EFlags |= 0x100;
113                    set_thread_context(thread, &ctx).map_err(|e| {
114                        CradleError::InvalidValue(format!(
115                            "{}!{} set context failed at {bp_addr:#x}: {e}",
116                            bp.module_name, bp.export_name
117                        ))
118                    })?;
119                    self.pending_restore = Some(PendingRestore {
120                        addr: bp_addr,
121                        thread,
122                    });
123                }
124                HookAction::Block(u) => {
125                    let ret_addr: u64 = read_remote_or(self.process, ctx.Rsp as usize, u);
126                    ctx.Rip = ret_addr;
127                    ctx.Rsp += 8;
128                    ctx.Rax = ret;
129                    set_thread_context(thread, &ctx).map_err(|e| {
130                        CradleError::InvalidValue(format!(
131                            "{}!{} block set context failed at {bp_addr:#x}: {e}",
132                            bp.module_name, bp.export_name
133                        ))
134                    })?;
135                }
136            }
137            Ok(DispatchResult::Handled)
138        }
139    }
140
141    /// Dispatches a single breakpoint by writing it to the process
142    ///
143    /// # Safety
144    /// Calls `WriteProcessMemory` under the hood (via FFI). This may be unsafe
145    pub unsafe fn dispatch_single_step(&mut self, thread: HANDLE) -> CradleResult<DispatchResult> {
146        match self.pending_restore.take() {
147            Some(restore) if restore.thread == thread => {
148                if self.breakpoints.contains_key(&restore.addr) {
149                    unsafe {
150                        write_process_memory(self.process, restore.addr, &[0xCC])?;
151                    }
152                }
153                Ok(DispatchResult::Handled)
154            }
155            Some(restore) => {
156                self.pending_restore = Some(restore);
157                Ok(DispatchResult::Skip)
158            }
159            None => Ok(DispatchResult::Skip),
160        }
161    }
162
163    /// Attempts to hook an exported function in the specified module
164    /// Will return an error if a hook is already installed there
165    ///
166    /// # Safety
167    /// Reads and writes data to the process' memory, which is unsafe and should be treated as such
168    pub unsafe fn hook_export(
169        &mut self,
170        module: &str,
171        export: &str,
172        handler: HookHandler,
173    ) -> CradleResult {
174        if self.module_cache.is_empty() {
175            return Err(CradleError::ModuleNotFound(
176                "module cache is empty".to_owned(),
177            ));
178        }
179        let module_lower = normalize_name(module);
180        let base = self
181            .module_cache
182            .get(&module_lower)
183            .copied()
184            .ok_or_else(|| CradleError::ModuleNotFound(module.to_string()))?;
185        let addr = resolve_export(self.process, base, export)?;
186        if let Some(bp) = self.breakpoints.get_mut(&addr) {
187            bp.push_handler(handler);
188            return Ok(());
189        }
190        let mut orig = [0u8; 1];
191        unsafe {
192            read_single_byte(self.process, addr, &mut orig)?;
193            write_single_byte(self.process, addr, &[0xCC])?;
194        }
195        self.breakpoints.insert(
196            addr,
197            Breakpoint::new(orig[0], handler, module_lower, export.to_string()),
198        );
199
200        Ok(())
201    }
202    /// Attempts to unhook the breakpoint at the given address
203    ///
204    /// # Safety
205    /// Calls `WriteProcessMemory` under the hood (via FFI). This may be unsafe
206    pub unsafe fn unhook(&mut self, addr: usize) -> CradleResult {
207        if let Some(bp) = self.breakpoints.remove(&addr) {
208            unsafe {
209                write_process_memory(self.process, addr, &[bp.original_byte])?;
210            }
211        }
212        Ok(())
213    }
214
215    /// Attempts to unhook all existing hooks
216    ///
217    /// # Safety
218    /// Calls `self.unhook(addr)` on each hook. This function is unsafe
219    pub unsafe fn unhook_all(&mut self) -> CradleResult {
220        let addrs: Vec<usize> = self.breakpoints.keys().copied().collect();
221        for addr in addrs {
222            unsafe {
223                self.unhook(addr)?;
224            }
225        }
226        Ok(())
227    }
228
229    /// Returns the amount of breakpoints are currently installed
230    pub fn hooked_count(&self) -> usize {
231        self.breakpoints.len()
232    }
233}
234unsafe impl Send for HookEngine {}
cradle_hooks/engine.rs

cradle_hooks/
engine.rs
