From 03997eb58ff311075c7459c1df8bffd53abc158d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Petr=20Bene=C5=A1?= Date: Fri, 13 Dec 2024 07:15:16 +0000 Subject: [PATCH] feat(core): move stack pointer handling into RecipeExecutor --- crates/vmi-utils/src/injector/mod.rs | 3 -- crates/vmi-utils/src/injector/os/windows.rs | 37 ++++---------- crates/vmi-utils/src/injector/recipe.rs | 53 ++++++++++++++++++--- 3 files changed, 55 insertions(+), 38 deletions(-) diff --git a/crates/vmi-utils/src/injector/mod.rs b/crates/vmi-utils/src/injector/mod.rs index 916c823..e7d723d 100644 --- a/crates/vmi-utils/src/injector/mod.rs +++ b/crates/vmi-utils/src/injector/mod.rs @@ -125,9 +125,6 @@ where /// Whether a thread has been successfully hijacked. pub(super) hijacked: bool, - /// Virtual address of the stack in the hijacked thread. - pub(super) sp_va: Option, - /// Virtual address of the instruction pointer in the hijacked thread. pub(super) ip_va: Option, diff --git a/crates/vmi-utils/src/injector/os/windows.rs b/crates/vmi-utils/src/injector/os/windows.rs index 47a24dd..d1572d3 100644 --- a/crates/vmi-utils/src/injector/os/windows.rs +++ b/crates/vmi-utils/src/injector/os/windows.rs @@ -194,7 +194,6 @@ where pid, tid: None, hijacked: false, - sp_va: None, ip_va: None, ip_pa: None, offsets, @@ -414,7 +413,6 @@ where // self.tid = Some(current_tid); - self.sp_va = Some(sp_va); self.ip_va = Some(ip_va); self.ip_pa = Some(ip_pa); @@ -500,37 +498,18 @@ where vmi.monitor_disable(EventMonitor::Register(ControlRegister::Cr3))?; } - let sp = Va(registers.rsp); - if let Some(sp_va) = self.sp_va { - if sp_va > sp { - tracing::trace!( - sp = %sp_va, - current_sp = %sp, - "not the right stack pointer" - ); - - return Ok( - VmiEventResponse::toggle_fast_singlestep().and_set_view(vmi.default_view()) - ); - } - } - // // Execute the next step in the recipe. // - //println!("BEFORE"); - //println!("RIP: 0x{:016X}", registers.rip); - //println!("RSP: 0x{:016X}", registers.rsp); - //let _ = hexdump(vmi, (Va(registers.rsp - 0x200), registers.cr3.into()), 0x400, Representation::U64); - - let new_registers = self.recipe.execute(vmi)?; - self.sp_va = Some(Va(new_registers.rsp)); - - //println!("AFTER"); - //println!("RIP: 0x{:016X}", new_registers.rip); - //println!("RSP: 0x{:016X}", new_registers.rsp); - //let _ = hexdump(vmi, (Va(registers.rsp - 0x200), registers.cr3.into()), 0x400, Representation::U64); + let new_registers = match self.recipe.execute(vmi)? { + Some(registers) => registers, + None => { + return Ok( + VmiEventResponse::toggle_fast_singlestep().and_set_view(vmi.default_view()) + ) + } + }; if self.recipe.done() { // diff --git a/crates/vmi-utils/src/injector/recipe.rs b/crates/vmi-utils/src/injector/recipe.rs index 85ad33c..3296231 100644 --- a/crates/vmi-utils/src/injector/recipe.rs +++ b/crates/vmi-utils/src/injector/recipe.rs @@ -124,6 +124,9 @@ where /// Index of the current step being executed. index: Option, + + /// The previous stack pointer value. + previous_stack_pointer: Option, } impl RecipeExecutor @@ -138,14 +141,25 @@ where cache: ImageSymbolCache::new(), original_registers: None, index: None, + previous_stack_pointer: None, } } /// Executes the next step in the recipe. + /// + /// Returns the new CPU registers after executing the step. + /// If the recipe has finished executing, returns the original registers. pub fn execute( &mut self, vmi: &VmiContext, - ) -> Result<<::Architecture as Architecture>::Registers, VmiError> { + ) -> Result::Architecture as Architecture>::Registers>, VmiError> + { + // If the stack pointer has decreased, we are likely in a recursive call or APC. + // In this case, we should not execute the recipe. + if self.has_stack_pointer_decreased(vmi) { + return Ok(None); + } + let index = match &mut self.index { Some(index) => index, None => { @@ -166,22 +180,25 @@ where cache: &mut self.cache, })?; + // Update the stack pointer after executing the step. + self.previous_stack_pointer = Some(Va(registers.stack_pointer())); + match next { RecipeControlFlow::Continue => { *index += 1; - return Ok(registers); + return Ok(Some(registers)); } RecipeControlFlow::Break => {} RecipeControlFlow::Repeat => { - return Ok(registers); + return Ok(Some(registers)); } RecipeControlFlow::Skip => { *index += 2; - return Ok(registers); + return Ok(Some(registers)); } RecipeControlFlow::Goto(i) => { *index = i; - return Ok(registers); + return Ok(Some(registers)); } } } @@ -193,7 +210,31 @@ where self.index = None; let original_registers = self.original_registers.expect("original_registers"); - Ok(original_registers) + Ok(Some(original_registers)) + } + + /// Returns whether the stack pointer has decreased since the last step. + /// + /// This is used to detect irrelevant reentrancy in the recipe, + /// such as recursive calls, interrupts, or APCs. + fn has_stack_pointer_decreased(&self, vmi: &VmiContext) -> bool { + let previous_stack_pointer = match self.previous_stack_pointer { + Some(previous_stack_pointer) => previous_stack_pointer, + None => return false, + }; + + let current_stack_pointer = Va(vmi.registers().stack_pointer()); + + let result = previous_stack_pointer > current_stack_pointer; + if result { + tracing::trace!( + %previous_stack_pointer, + %current_stack_pointer, + "stack pointer decreased" + ); + } + + result } /// Resets the executor to the initial state.