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

wraith/manipulation/inline_hook/hook/
mid.rs
