feat(core): move stack pointer handling into RecipeExecutor

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<Va>,
-
     /// Virtual address of the instruction pointer in the hijacked thread.
     pub(super) ip_va: Option<Va>,
 

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() {
             //

crates/vmi-utils/src/injector/recipe.rs

@@ -124,6 +124,9 @@ where
 
     /// Index of the current step being executed.
     index: Option<usize>,
+
+    /// The previous stack pointer value.
+    previous_stack_pointer: Option<Va>,
 }
 
 impl<Driver, Os, T> RecipeExecutor<Driver, Os, T>
@@ -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<Driver, Os>,
-    ) -> Result<<<Driver as VmiDriver>::Architecture as Architecture>::Registers, VmiError> {
+    ) -> Result<Option<<<Driver as VmiDriver>::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<Driver, Os>) -> 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.