1use core::ffi::c_void;
4use core::ptr::addr_of_mut;
5use core::sync::atomic::{Ordering, AtomicBool};
6
7use super::{
8 NtSetContextThread,
9 NtGetContextThread,
10 NtCurrentThread
11};
12use super::data::{
13 CONTEXT, CONTEXT_DEBUG_REGISTERS_AMD64, EXCEPTION_SINGLE_STEP,
14 EXCEPTION_CONTINUE_EXECUTION, EXCEPTION_CONTINUE_SEARCH,
15 EXCEPTION_POINTERS, HANDLE, OBJECT_ATTRIBUTES,
16 CONTEXT_DEBUG_REGISTERS_X86
17};
18
19pub static mut CURRENT_API: Option<WINAPI> = None;
21
22static USE_BREAKPOINT: AtomicBool = AtomicBool::new(false);
24
25#[inline(always)]
31pub fn set_use_breakpoint(enabled: bool) {
32 USE_BREAKPOINT.store(enabled, Ordering::SeqCst);
33}
34
35#[inline(always)]
41pub fn is_breakpoint_enabled() -> bool {
42 USE_BREAKPOINT.load(Ordering::SeqCst)
43}
44
45pub(crate) fn set_breakpoint<T: Into<u64>>(address: T) {
51 let mut ctx = CONTEXT {
52 ContextFlags: if cfg!(target_arch = "x86_64") { CONTEXT_DEBUG_REGISTERS_AMD64 } else { CONTEXT_DEBUG_REGISTERS_X86 },
53 ..Default::default()
54 };
55
56 NtGetContextThread(NtCurrentThread(), &mut ctx);
57
58 cfg_if::cfg_if! {
59 if #[cfg(target_arch = "x86_64")] {
60 ctx.Dr0 = address.into();
61 ctx.Dr6 = 0x00;
62 ctx.Dr7 = set_dr7_bits(ctx.Dr7, 0, 1, 1);
63 } else {
64 ctx.Dr0 = address.into() as u32;
65 ctx.Dr6 = 0x00;
66 ctx.Dr7 = set_dr7_bits(ctx.Dr7 as u64, 0, 1, 1) as u32;
67 }
68 }
69
70 NtSetContextThread(NtCurrentThread(), &ctx);
71}
72
73fn set_dr7_bits<T: Into<u64>>(current: T, start_bit: i32, nmbr_bits: i32, new_bit: u64) -> u64 {
86 let current = current.into();
87 let mask = (1u64 << nmbr_bits) - 1;
88 (current & !(mask << start_bit)) | (new_bit << start_bit)
89}
90
91#[derive(Debug)]
93pub enum WINAPI {
94 NtAllocateVirtualMemory {
96 ProcessHandle: HANDLE,
97 Protect: u32,
98 },
99
100 NtProtectVirtualMemory {
102 ProcessHandle: HANDLE,
103 NewProtect: u32,
104 },
105
106 NtCreateThreadEx {
108 ProcessHandle: HANDLE,
109 ThreadHandle: *mut HANDLE,
110 DesiredAccess: u32,
111 ObjectAttributes: *mut OBJECT_ATTRIBUTES
112 },
113
114 NtWriteVirtualMemory {
116 ProcessHandle: HANDLE,
117 Buffer: *mut c_void,
118 NumberOfBytesToWrite: *mut usize,
119 },
120}
121
122#[cfg(target_arch = "x86_64")]
134#[allow(unsafe_op_in_unsafe_fn)]
135pub unsafe extern "system" fn veh_handler(exceptioninfo: *mut EXCEPTION_POINTERS) -> i32 {
136 if !is_breakpoint_enabled() || (*(*exceptioninfo).ExceptionRecord).ExceptionCode != EXCEPTION_SINGLE_STEP {
137 return EXCEPTION_CONTINUE_SEARCH;
138 }
139
140 let context = (*exceptioninfo).ContextRecord;
141 if (*context).Rip == (*context).Dr0 && (*context).Dr7 & 1 == 1 {
142 if let Some(current) = (*addr_of_mut!(CURRENT_API)).take() {
143 match current {
144 WINAPI::NtAllocateVirtualMemory {
145 ProcessHandle,
146 Protect
147 } => {
148 (*context).R10 = ProcessHandle as u64;
149 *(((*context).Rsp + 0x30) as *mut u32) = Protect;
150 },
151
152 WINAPI::NtProtectVirtualMemory {
153 ProcessHandle,
154 NewProtect,
155 } => {
156 (*context).R10 = ProcessHandle as u64;
157 (*context).R9 = NewProtect as u64;
158 },
159
160 WINAPI::NtCreateThreadEx {
161 ProcessHandle,
162 ThreadHandle,
163 DesiredAccess,
164 ObjectAttributes
165 } => {
166 (*context).R10 = ThreadHandle as u64;
167 (*context).Rdx = DesiredAccess as u64;
168 (*context).R8 = ObjectAttributes as u64;
169 (*context).R9 = ProcessHandle as u64;
170 },
171
172 WINAPI::NtWriteVirtualMemory {
173 ProcessHandle,
174 Buffer,
175 NumberOfBytesToWrite,
176 } => {
177 (*context).R10 = ProcessHandle as u64;
178 (*context).R8 = Buffer as u64;
179 (*context).R9 = NumberOfBytesToWrite as u64;
180 }
181 }
182
183 (*context).Dr0 = 0x00;
184 (*context).Dr6 = 0x00;
185 (*context).Dr7 = set_dr7_bits((*context).Dr7, 0, 1, 0);
186 }
187
188 return EXCEPTION_CONTINUE_EXECUTION;
189 }
190
191 EXCEPTION_CONTINUE_SEARCH
192}
193
194#[cfg(target_arch = "x86")]
206#[allow(unsafe_op_in_unsafe_fn)]
207pub unsafe extern "system" fn veh_handler(exceptioninfo: *mut EXCEPTION_POINTERS) -> i32 {
208 if !is_breakpoint_enabled() || (*(*exceptioninfo).ExceptionRecord).ExceptionCode != EXCEPTION_SINGLE_STEP {
209 return EXCEPTION_CONTINUE_SEARCH;
210 }
211
212 let context = (*exceptioninfo).ContextRecord;
213 if (*context).Eip == (*context).Dr0 && (*context).Dr7 & 1 == 1 {
214 if let Some(current) = (*addr_of_mut!(CURRENT_API)).take() {
215 match current {
216 WINAPI::NtAllocateVirtualMemory {
217 ProcessHandle,
218 Protect
219 } => {
220 *(((*context).Esp + 0x4) as *mut u32) = ProcessHandle as u32;
221 *(((*context).Esp + 0x18) as *mut u32) = Protect;
222 },
223
224 WINAPI::NtProtectVirtualMemory {
225 ProcessHandle,
226 NewProtect,
227 } => {
228 *(((*context).Esp + 0x4) as *mut u32) = ProcessHandle as u32;
229 *(((*context).Esp + 0x10) as *mut u32) = NewProtect as u32;
230 },
231
232 WINAPI::NtCreateThreadEx {
233 ProcessHandle,
234 ThreadHandle,
235 DesiredAccess,
236 ObjectAttributes
237 } => {
238 *(((*context).Esp + 0x4) as *mut u32) = ThreadHandle as u32;
239 *(((*context).Esp + 0x8) as *mut u32) = DesiredAccess as u32;
240 *(((*context).Esp + 0xC) as *mut u32) = ObjectAttributes as u32;
241 *(((*context).Esp + 0x10) as *mut u32) = ProcessHandle as u32;
242 },
243
244 WINAPI::NtWriteVirtualMemory {
245 ProcessHandle,
246 Buffer,
247 NumberOfBytesToWrite,
248 } => {
249 *(((*context).Esp + 0x4) as *mut u32) = ProcessHandle as u32;
250 *(((*context).Esp + 0xC) as *mut u32) = Buffer as u32;
251 *(((*context).Esp + 0x10) as *mut u32) = NumberOfBytesToWrite as u32;
252 }
253 }
254
255 (*context).Dr0 = 0x00;
256 (*context).Dr6 = 0x00;
257 (*context).Dr7 = set_dr7_bits((*context).Dr7, 0, 1, 0) as u32;
258 }
259
260 return EXCEPTION_CONTINUE_EXECUTION;
261 }
262
263 EXCEPTION_CONTINUE_SEARCH
264}