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
10pub enum DispatchResult {
12 Handled,
14 Skip,
16}
17
18pub 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 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 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 pub fn register_module(&mut self, name: &str, base: usize) {
57 self.module_cache.insert(normalize_name(name), base);
58 }
59
60 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 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 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 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 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 pub fn hooked_count(&self) -> usize {
198 self.breakpoints.len()
199 }
200}
201unsafe impl Send for HookEngine {}