wraith/manipulation/inline_hook/hook/
mid.rs

1//! Mid-function hooks
2//!
3//! Hooks at arbitrary locations within a function, not just the prologue.
4//! Requires saving and restoring CPU context around the detour call.
5
6use crate::error::{Result, WraithError};
7use crate::util::memory::ProtectionGuard;
8use crate::manipulation::inline_hook::arch::Architecture;
9use crate::manipulation::inline_hook::guard::HookGuard;
10use crate::manipulation::inline_hook::trampoline::ExecutableMemory;
11use crate::manipulation::inline_hook::asm::Encoder;
12use super::Hook;
13use core::marker::PhantomData;
14
15const PAGE_EXECUTE_READWRITE: u32 = 0x40;
16
17/// CPU context passed to mid-function hook handlers (x64)
18#[cfg(target_arch = "x86_64")]
19#[repr(C)]
20#[derive(Debug, Clone, Copy)]
21pub struct HookContext {
22    pub rflags: u64,
23    pub r15: u64,
24    pub r14: u64,
25    pub r13: u64,
26    pub r12: u64,
27    pub r11: u64,
28    pub r10: u64,
29    pub r9: u64,
30    pub r8: u64,
31    pub rdi: u64,
32    pub rsi: u64,
33    pub rbp: u64,
34    pub rsp: u64,
35    pub rbx: u64,
36    pub rdx: u64,
37    pub rcx: u64,
38    pub rax: u64,
39}
40
41/// CPU context passed to mid-function hook handlers (x86)
42#[cfg(target_arch = "x86")]
43#[repr(C)]
44#[derive(Debug, Clone, Copy)]
45pub struct HookContext {
46    pub eflags: u32,
47    pub edi: u32,
48    pub esi: u32,
49    pub ebp: u32,
50    pub esp: u32,
51    pub ebx: u32,
52    pub edx: u32,
53    pub ecx: u32,
54    pub eax: u32,
55}
56
57/// mid-function hook type signature
58#[cfg(target_arch = "x86_64")]
59pub type MidHookFn = extern "C" fn(ctx: *mut HookContext);
60
61#[cfg(target_arch = "x86")]
62pub type MidHookFn = extern "cdecl" fn(ctx: *mut HookContext);
63
64/// mid-function hook
65///
66/// hooks at an arbitrary location within a function.
67/// saves all registers, calls the detour with a context pointer,
68/// then restores registers and continues execution.
69pub struct MidFunctionHook<A: Architecture> {
70    /// address to hook (not necessarily function start)
71    address: usize,
72    /// detour function receiving context pointer
73    detour: usize,
74    _arch: PhantomData<A>,
75}
76
77impl<A: Architecture> MidFunctionHook<A> {
78    /// create a new mid-function hook
79    ///
80    /// # Arguments
81    /// * `address` - the exact address to place the hook
82    /// * `detour` - function pointer of type `MidHookFn`
83    pub fn new(address: usize, detour: MidHookFn) -> Self {
84        Self {
85            address,
86            detour: detour as usize,
87            _arch: PhantomData,
88        }
89    }
90
91    /// create from raw address
92    pub fn from_raw(address: usize, detour: usize) -> Self {
93        Self {
94            address,
95            detour,
96            _arch: PhantomData,
97        }
98    }
99
100    /// install the mid-function hook
101    #[cfg(target_arch = "x86_64")]
102    pub fn install(self) -> Result<HookGuard<A>> {
103        let min_hook_size = A::MIN_HOOK_SIZE;
104
105        // read original bytes at hook location
106        let original_bytes = unsafe {
107            let target_bytes = core::slice::from_raw_parts(self.address as *const u8, 64);
108            let boundary = A::find_instruction_boundary(target_bytes, min_hook_size)
109                .ok_or_else(|| WraithError::HookDetectionFailed {
110                    function: format!("{:#x}", self.address),
111                    reason: "failed to find instruction boundary".into(),
112                })?;
113            core::slice::from_raw_parts(self.address as *const u8, boundary).to_vec()
114        };
115
116        let prologue_size = original_bytes.len();
117
118        // allocate memory for the wrapper stub
119        let mut stub_memory = ExecutableMemory::allocate_near(self.address, 256)?;
120        let stub_base = stub_memory.base();
121
122        // build the context-saving wrapper stub
123        let wrapper_code = build_x64_wrapper(
124            stub_base,
125            self.detour,
126            self.address,
127            &original_bytes,
128        )?;
129
130        stub_memory.write(&wrapper_code)?;
131        stub_memory.flush_icache()?;
132
133        // write jump to wrapper at hook location
134        let hook_stub = A::encode_jmp_rel(self.address, stub_base)
135            .or_else(|| Some(A::encode_jmp_abs(stub_base)))
136            .unwrap();
137
138        let mut padded_stub = hook_stub;
139        if padded_stub.len() < prologue_size {
140            let padding = A::encode_nop_sled(prologue_size - padded_stub.len());
141            padded_stub.extend_from_slice(&padding);
142        }
143
144        // write the hook
145        {
146            let _guard = ProtectionGuard::new(
147                self.address,
148                prologue_size,
149                PAGE_EXECUTE_READWRITE,
150            )?;
151
152            unsafe {
153                core::ptr::copy_nonoverlapping(
154                    padded_stub.as_ptr(),
155                    self.address as *mut u8,
156                    prologue_size,
157                );
158            }
159        }
160
161        flush_icache(self.address, prologue_size)?;
162
163        Ok(HookGuard::new(
164            self.address,
165            self.detour,
166            original_bytes,
167            Some(stub_memory),
168        ))
169    }
170
171    #[cfg(target_arch = "x86")]
172    pub fn install(self) -> Result<HookGuard<A>> {
173        let min_hook_size = A::MIN_HOOK_SIZE;
174
175        let original_bytes = unsafe {
176            let target_bytes = core::slice::from_raw_parts(self.address as *const u8, 32);
177            let boundary = A::find_instruction_boundary(target_bytes, min_hook_size)
178                .ok_or_else(|| WraithError::HookDetectionFailed {
179                    function: format!("{:#x}", self.address),
180                    reason: "failed to find instruction boundary".into(),
181                })?;
182            core::slice::from_raw_parts(self.address as *const u8, boundary).to_vec()
183        };
184
185        let prologue_size = original_bytes.len();
186
187        let mut stub_memory = ExecutableMemory::allocate_near(self.address, 128)?;
188        let stub_base = stub_memory.base();
189
190        let wrapper_code = build_x86_wrapper(
191            stub_base,
192            self.detour,
193            self.address,
194            &original_bytes,
195        )?;
196
197        stub_memory.write(&wrapper_code)?;
198        stub_memory.flush_icache()?;
199
200        let hook_stub = A::encode_jmp_rel(self.address, stub_base)
201            .or_else(|| Some(A::encode_jmp_abs(stub_base)))
202            .unwrap();
203
204        let mut padded_stub = hook_stub;
205        if padded_stub.len() < prologue_size {
206            let padding = A::encode_nop_sled(prologue_size - padded_stub.len());
207            padded_stub.extend_from_slice(&padding);
208        }
209
210        {
211            let _guard = ProtectionGuard::new(
212                self.address,
213                prologue_size,
214                PAGE_EXECUTE_READWRITE,
215            )?;
216
217            unsafe {
218                core::ptr::copy_nonoverlapping(
219                    padded_stub.as_ptr(),
220                    self.address as *mut u8,
221                    prologue_size,
222                );
223            }
224        }
225
226        flush_icache(self.address, prologue_size)?;
227
228        Ok(HookGuard::new(
229            self.address,
230            self.detour,
231            original_bytes,
232            Some(stub_memory),
233        ))
234    }
235}
236
237impl<A: Architecture> Hook for MidFunctionHook<A> {
238    type Guard = HookGuard<A>;
239
240    fn install(self) -> Result<Self::Guard> {
241        MidFunctionHook::install(self)
242    }
243
244    fn target(&self) -> usize {
245        self.address
246    }
247
248    fn detour(&self) -> usize {
249        self.detour
250    }
251}
252
253/// build x64 context-saving wrapper
254#[cfg(target_arch = "x86_64")]
255fn build_x64_wrapper(
256    stub_base: usize,
257    detour: usize,
258    original_addr: usize,
259    original_bytes: &[u8],
260) -> Result<Vec<u8>> {
261    use crate::manipulation::inline_hook::arch::X64;
262
263    let mut code = Vec::with_capacity(256);
264
265    // save all registers (context structure)
266    // pushfq
267    code.push(0x9C);
268
269    // push r15-r8, rdi, rsi, rbp, rsp placeholder, rbx, rdx, rcx, rax
270    code.extend_from_slice(&[0x41, 0x57]); // push r15
271    code.extend_from_slice(&[0x41, 0x56]); // push r14
272    code.extend_from_slice(&[0x41, 0x55]); // push r13
273    code.extend_from_slice(&[0x41, 0x54]); // push r12
274    code.extend_from_slice(&[0x41, 0x53]); // push r11
275    code.extend_from_slice(&[0x41, 0x52]); // push r10
276    code.extend_from_slice(&[0x41, 0x51]); // push r9
277    code.extend_from_slice(&[0x41, 0x50]); // push r8
278    code.push(0x57); // push rdi
279    code.push(0x56); // push rsi
280    code.push(0x55); // push rbp
281    code.push(0x54); // push rsp (placeholder, will be fixed)
282    code.push(0x53); // push rbx
283    code.push(0x52); // push rdx
284    code.push(0x51); // push rcx
285    code.push(0x50); // push rax
286
287    // fix rsp in context (add offset for pushes)
288    // mov [rsp + 0x38], rsp; add qword ptr [rsp+0x38], 0x90
289    code.extend_from_slice(&[0x48, 0x89, 0x64, 0x24, 0x38]); // mov [rsp+0x38], rsp
290    code.extend_from_slice(&[0x48, 0x83, 0x44, 0x24, 0x38, 0x90]); // add [rsp+0x38], 0x90
291
292    // mov rcx, rsp (first arg = context pointer)
293    code.extend_from_slice(&[0x48, 0x89, 0xE1]);
294
295    // sub rsp, 0x28 (shadow space + alignment)
296    code.extend_from_slice(&[0x48, 0x83, 0xEC, 0x28]);
297
298    // call detour
299    // we'll use mov rax, imm64; call rax for simplicity
300    code.extend_from_slice(&[0x48, 0xB8]); // mov rax, imm64
301    code.extend_from_slice(&detour.to_le_bytes());
302    code.extend_from_slice(&[0xFF, 0xD0]); // call rax
303
304    // add rsp, 0x28
305    code.extend_from_slice(&[0x48, 0x83, 0xC4, 0x28]);
306
307    // restore registers
308    code.push(0x58); // pop rax
309    code.push(0x59); // pop rcx
310    code.push(0x5A); // pop rdx
311    code.push(0x5B); // pop rbx
312    code.extend_from_slice(&[0x48, 0x83, 0xC4, 0x08]); // add rsp, 8 (skip rsp)
313    code.push(0x5D); // pop rbp
314    code.push(0x5E); // pop rsi
315    code.push(0x5F); // pop rdi
316    code.extend_from_slice(&[0x41, 0x58]); // pop r8
317    code.extend_from_slice(&[0x41, 0x59]); // pop r9
318    code.extend_from_slice(&[0x41, 0x5A]); // pop r10
319    code.extend_from_slice(&[0x41, 0x5B]); // pop r11
320    code.extend_from_slice(&[0x41, 0x5C]); // pop r12
321    code.extend_from_slice(&[0x41, 0x5D]); // pop r13
322    code.extend_from_slice(&[0x41, 0x5E]); // pop r14
323    code.extend_from_slice(&[0x41, 0x5F]); // pop r15
324    code.push(0x9D); // popfq
325
326    // relocate and copy original instructions
327    // we need to track the current position in the stub for relocation calculations
328    let mut src_offset = 0;
329
330    while src_offset < original_bytes.len() {
331        let remaining = &original_bytes[src_offset..];
332
333        // find instruction length
334        let insn_len = X64::find_instruction_boundary(remaining, 1)
335            .ok_or_else(|| WraithError::HookDetectionFailed {
336                function: format!("{:#x}", original_addr + src_offset),
337                reason: "failed to decode instruction during relocation".into(),
338            })?;
339
340        let instruction = &original_bytes[src_offset..src_offset + insn_len];
341        let old_addr = original_addr + src_offset;
342        let new_addr = stub_base + code.len();
343
344        // check if instruction needs relocation
345        if X64::needs_relocation(instruction) {
346            match X64::relocate_instruction(instruction, old_addr, new_addr) {
347                Some(relocated) => {
348                    code.extend_from_slice(&relocated);
349                }
350                None => {
351                    // relocation failed - instruction target is too far
352                    // fall back to indirect: push target; ret pattern
353                    return Err(WraithError::RelocationFailed {
354                        rva: src_offset as u32,
355                        reason: format!(
356                            "instruction at {:#x} cannot be relocated to {:#x}",
357                            old_addr, new_addr
358                        ),
359                    });
360                }
361            }
362        } else {
363            // no relocation needed, copy as-is
364            code.extend_from_slice(instruction);
365        }
366
367        src_offset += insn_len;
368    }
369
370    // jump to continuation (original_addr + original_bytes.len())
371    let continuation = original_addr + original_bytes.len();
372    let jmp_location = stub_base + code.len();
373
374    if let Some(jmp) = X64::encode_jmp_rel(jmp_location, continuation) {
375        code.extend_from_slice(&jmp);
376    } else {
377        code.extend_from_slice(&X64::encode_jmp_abs(continuation));
378    }
379
380    Ok(code)
381}
382
383/// build x86 context-saving wrapper
384#[cfg(target_arch = "x86")]
385fn build_x86_wrapper(
386    stub_base: usize,
387    detour: usize,
388    original_addr: usize,
389    original_bytes: &[u8],
390) -> Result<Vec<u8>> {
391    use crate::manipulation::inline_hook::arch::X86;
392
393    let mut code = Vec::with_capacity(128);
394
395    // pushfd
396    code.push(0x9C);
397    // pushad (edi, esi, ebp, esp, ebx, edx, ecx, eax)
398    code.push(0x60);
399
400    // push esp (context pointer as argument)
401    code.push(0x54);
402
403    // call detour
404    code.push(0xB8); // mov eax, imm32
405    code.extend_from_slice(&(detour as u32).to_le_bytes());
406    code.extend_from_slice(&[0xFF, 0xD0]); // call eax
407
408    // add esp, 4 (pop argument)
409    code.extend_from_slice(&[0x83, 0xC4, 0x04]);
410
411    // popad
412    code.push(0x61);
413    // popfd
414    code.push(0x9D);
415
416    // relocate and copy original instructions
417    let mut src_offset = 0;
418
419    while src_offset < original_bytes.len() {
420        let remaining = &original_bytes[src_offset..];
421
422        // find instruction length
423        let insn_len = X86::find_instruction_boundary(remaining, 1)
424            .ok_or_else(|| WraithError::HookDetectionFailed {
425                function: format!("{:#x}", original_addr + src_offset),
426                reason: "failed to decode instruction during relocation".into(),
427            })?;
428
429        let instruction = &original_bytes[src_offset..src_offset + insn_len];
430        let old_addr = original_addr + src_offset;
431        let new_addr = stub_base + code.len();
432
433        // check if instruction needs relocation
434        if X86::needs_relocation(instruction) {
435            match X86::relocate_instruction(instruction, old_addr, new_addr) {
436                Some(relocated) => {
437                    code.extend_from_slice(&relocated);
438                }
439                None => {
440                    return Err(WraithError::RelocationFailed {
441                        rva: src_offset as u32,
442                        reason: format!(
443                            "instruction at {:#x} cannot be relocated to {:#x}",
444                            old_addr, new_addr
445                        ),
446                    });
447                }
448            }
449        } else {
450            // no relocation needed, copy as-is
451            code.extend_from_slice(instruction);
452        }
453
454        src_offset += insn_len;
455    }
456
457    // jump to continuation
458    let continuation = original_addr + original_bytes.len();
459    let jmp_location = stub_base + code.len();
460
461    if let Some(jmp) = X86::encode_jmp_rel(jmp_location, continuation) {
462        code.extend_from_slice(&jmp);
463    } else {
464        code.extend_from_slice(&X86::encode_jmp_abs(continuation as usize));
465    }
466
467    Ok(code)
468}
469
470fn flush_icache(address: usize, size: usize) -> Result<()> {
471    let result = unsafe {
472        FlushInstructionCache(
473            GetCurrentProcess(),
474            address as *const _,
475            size,
476        )
477    };
478
479    if result == 0 {
480        Err(WraithError::from_last_error("FlushInstructionCache"))
481    } else {
482        Ok(())
483    }
484}
485
486#[link(name = "kernel32")]
487extern "system" {
488    fn FlushInstructionCache(
489        hProcess: *mut core::ffi::c_void,
490        lpBaseAddress: *const core::ffi::c_void,
491        dwSize: usize,
492    ) -> i32;
493
494    fn GetCurrentProcess() -> *mut core::ffi::c_void;
495}
wraith/manipulation/inline_hook/hook/mid.rs

wraith/manipulation/inline_hook/hook/
mid.rs
