Add control flow redirect mechanism for routine hooks

quarkslab · Nov 11, 2024 · 20e0bca · 20e0bca
1 parent 07fd1a5
commit 20e0bca
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 7 deletions.
diff --git a/tritondse/process_state.py b/tritondse/process_state.py
@@ -126,6 +126,9 @@ def __init__(self, endianness: Endian = Endian.LITTLE, time_inc_coefficient: flo
         # The memory mapping of the program ({vaddr_s : vaddr_e})
         self.__program_segments_mapping = {}
 
+        # Address to return to from a routine hook.
+        self.rtn_redirect_addr = None
+
     @property
     def threads(self) -> List[ThreadContext]:
         """

diff --git a/tritondse/routines.py b/tritondse/routines.py
@@ -159,7 +159,9 @@ def rtn_libc_start_main(se: 'SymbolicExecutor', pstate: 'ProcessState'):
     # WARNING: Dirty trick to make sure to jump on main after
     # the emulation of that stub
     if pstate.architecture == Architecture.AARCH64:
-        pstate.cpu.x30 = main
+        pass
+    elif pstate.architecture == Architecture.ARM32:
+        pass
     elif pstate.architecture in [Architecture.X86_64, Architecture.X86]:
         # Push the return value to jump into the main() function
         pstate.push_stack_value(main)
@@ -224,6 +226,10 @@ def rtn_libc_start_main(se: 'SymbolicExecutor', pstate: 'ProcessState'):
         pstate.write_argument_value(1 + 1, b_argv)
     else:
         pstate.write_argument_value(1, b_argv)
+
+    # Return to the main function.
+    pstate.rtn_redirect_addr = main
+
     return None
 
 

diff --git a/tritondse/symbolic_executor.py b/tritondse/symbolic_executor.py
@@ -25,6 +25,7 @@
 from tritondse.thread_context import ThreadContext
 from tritondse.exception import AbortExecutionException, SkipInstructionException, StopExplorationException, ProbeException
 from tritondse.memory import MemoryAccessViolation, Perm
+from tritondse.thread_context import ThreadState
 import tritondse.logging
 
 logger = tritondse.logging.get("executor")
@@ -492,14 +493,24 @@ def _routines_handler(self, instruction: Instruction):
                 self.pstate.current_thread.count = self.config.thread_scheduling+1
                 return
 
-            if self.pstate.architecture == Architecture.AARCH64:
+            # If an return address was set, return to that address.
+            if self.pstate.rtn_redirect_addr != None:
                 # Get the return address
-                ret_addr = self.pstate.read_register('x30')
-            elif self.pstate.architecture in [Architecture.X86, Architecture.X86_64]:
-                # Get the return address and restore RSP (simulate RET)
-                ret_addr = self.pstate.pop_stack_value()
+                ret_addr = self.pstate.rtn_redirect_addr
+                # Clear rtn_redirect_addr.
+                self.pstate.rtn_redirect_addr = None
             else:
-                raise Exception("Architecture not supported")
+                if self.pstate.architecture == Architecture.AARCH64:
+                    # Get the return address.
+                    ret_addr = self.pstate.read_register('x30')
+                elif self.pstate.architecture == Architecture.ARM32:
+                    # Get the return address.
+                    ret_addr = self.pstate.read_register('r14')
+                elif self.pstate.architecture in [Architecture.X86, Architecture.X86_64]:
+                    # Get the return address and restore RSP (simulate RET).
+                    ret_addr = self.pstate.pop_stack_value()
+                else:
+                    raise Exception("Architecture not supported")
 
             # Hijack RIP to skip the call
             self.pstate.cpu.program_counter = ret_addr