Merge pull request libdebug#112 from libdebug/unwind-fix

Unwind fix

libdebug/architectures/aarch64/aarch64_stack_unwinder.py

@@ -9,8 +9,10 @@
 from typing import TYPE_CHECKING
 
 from libdebug.architectures.stack_unwinding_manager import StackUnwindingManager
+from libdebug.liblog import liblog
 
 if TYPE_CHECKING:
+    from libdebug.data.memory_map import MemoryMap
     from libdebug.state.thread_context import ThreadContext
 
 
@@ -29,10 +31,18 @@ def unwind(self: Aarch64StackUnwinder, target: ThreadContext) -> list:
         assert hasattr(target.regs, "pc")
 
         frame_pointer = target.regs.x29
-        initial_link_register = target.regs.x30
-        stack_trace = [target.regs.pc, initial_link_register]
 
         vmaps = target._internal_debugger.debugging_interface.maps()
+        initial_link_register = None
+
+        try:
+            initial_link_register = self.get_return_address(target, vmaps)
+        except ValueError:
+            liblog.warning(
+                "Failed to get the return address. Check stack frame registers (e.g., base pointer). The stack trace may be incomplete.",
+            )
+
+        stack_trace = [target.regs.pc, initial_link_register] if initial_link_register else [target.regs.pc]
 
         # Follow the frame chain
         while frame_pointer:
@@ -56,13 +66,19 @@ def unwind(self: Aarch64StackUnwinder, target: ThreadContext) -> list:
 
         return stack_trace
 
-    def get_return_address(self: Aarch64StackUnwinder, target: ThreadContext) -> int:
+    def get_return_address(self: Aarch64StackUnwinder, target: ThreadContext, vmaps: list[MemoryMap]) -> int:
         """Get the return address of the current function.
 
         Args:
             target (ThreadContext): The target ThreadContext.
+            vmaps (list[MemoryMap]): The memory maps of the process.
 
         Returns:
             int: The return address.
         """
-        return target.regs.x30
+        return_address = target.regs.x30
+
+        if not any(vmap.start <= return_address < vmap.end for vmap in vmaps):
+            raise ValueError("Return address not in any valid memory map")
+
+        return return_address

libdebug/architectures/amd64/amd64_stack_unwinder.py

@@ -9,13 +9,13 @@
 from typing import TYPE_CHECKING
 
 from libdebug.architectures.stack_unwinding_manager import StackUnwindingManager
-from libdebug.liblog import logging
+from libdebug.liblog import liblog
 
 if TYPE_CHECKING:
+    from libdebug.data.memory_map import MemoryMap
     from libdebug.state.thread_context import ThreadContext
 
 
-
 class Amd64StackUnwinder(StackUnwindingManager):
     """Class that provides stack unwinding for the x86_64 architecture."""
 
@@ -54,22 +54,26 @@ def unwind(self: Amd64StackUnwinder, target: ThreadContext) -> list:
         # If we are in the prolouge of a function, we need to get the return address from the stack
         # using a slightly more complex method
         try:
-            first_return_address = self.get_return_address(target)
+            first_return_address = self.get_return_address(target, vmaps)
 
-            if first_return_address != stack_trace[1]:
-                stack_trace.insert(1, first_return_address)
+            if len(stack_trace) > 1:
+                if first_return_address != stack_trace[1]:
+                    stack_trace.insert(1, first_return_address)
+            else:
+                stack_trace.append(first_return_address)
         except (OSError, ValueError):
-            logging.WARNING(
-                "Failed to get the return address from the stack. Check stack frame registers (e.g., base pointer). The stack trace may be incomplete.",
+            liblog.warning(
+                "Failed to get the return address. Check stack frame registers (e.g., base pointer). The stack trace may be incomplete.",
             )
 
         return stack_trace
 
-    def get_return_address(self: Amd64StackUnwinder, target: ThreadContext) -> int:
+    def get_return_address(self: Amd64StackUnwinder, target: ThreadContext, vmaps: list[MemoryMap]) -> int:
         """Get the return address of the current function.
 
         Args:
             target (ThreadContext): The target ThreadContext.
+            vmaps (list[MemoryMap]): The memory maps of the process.
 
         Returns:
             int: The return address.
@@ -86,7 +90,12 @@ def get_return_address(self: Amd64StackUnwinder, target: ThreadContext) -> int:
         else:
             return_address = target.memory[target.regs.rsp + 8, 8, "absolute"]
 
-        return int.from_bytes(return_address, byteorder="little")
+        return_address = int.from_bytes(return_address, byteorder="little")
+
+        if not any(vmap.start <= return_address < vmap.end for vmap in vmaps):
+            raise ValueError("Return address not in any valid memory map")
+
+        return return_address
 
     def _preamble_state(self: Amd64StackUnwinder, instruction_window: bytes) -> int:
         """Check if the instruction window is a function preamble and if so at what stage.

libdebug/architectures/stack_unwinding_manager.py

@@ -10,6 +10,7 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from libdebug.data.memory_map import MemoryMap
     from libdebug.state.thread_context import ThreadContext
 
 
@@ -21,5 +22,5 @@ def unwind(self: StackUnwindingManager, target: ThreadContext) -> list:
         """Unwind the stack of the target process."""
 
     @abstractmethod
-    def get_return_address(self: StackUnwindingManager, target: ThreadContext) -> int:
+    def get_return_address(self: StackUnwindingManager, target: ThreadContext, vmaps: list[MemoryMap]) -> int:
         """Get the return address of the current function."""

libdebug/ptrace/ptrace_interface.py

@@ -310,7 +310,7 @@ def finish(self: PtraceInterface, thread: ThreadContext, heuristic: str) -> None
             invalidate_process_cache()
         elif heuristic == "backtrace":
             # Breakpoint to return address
-            last_saved_instruction_pointer = thread.current_return_address()
+            last_saved_instruction_pointer = thread.saved_ip
 
             # If a breakpoint already exists at the return address, we don't need to set a new one
             found = False

libdebug/state/thread_context.py

@@ -201,13 +201,14 @@ def print_backtrace(self: ThreadContext) -> None:
             else:
                 print(f"{PrintStyle.RED}{return_address:#x} <{return_address_symbol}> {PrintStyle.RESET}")
 
-    def current_return_address(self: ThreadContext) -> int:
-        """Returns the return address of the current function."""
+    @property
+    def saved_ip(self: ThreadContext) -> int:
+        """The return address of the current function."""
         self._internal_debugger._ensure_process_stopped()
         stack_unwinder = stack_unwinding_provider(self._internal_debugger.arch)
 
         try:
-            return_address = stack_unwinder.get_return_address(self)
+            return_address = stack_unwinder.get_return_address(self, self._internal_debugger.debugging_interface.maps())
         except (OSError, ValueError) as e:
             raise ValueError(
                 "Failed to get the return address. Check stack frame registers (e.g., base pointer).",
@@ -249,8 +250,7 @@ def finish(self: ThreadContext, heuristic: str = "backtrace") -> None:
         self._internal_debugger.finish(self, heuristic=heuristic)
 
     def next(self: ThreadContext) -> None:
-        """Executes the next instruction of the process. If the instruction is a call, the debugger will continue until the called function returns.
-        """
+        """Executes the next instruction of the process. If the instruction is a call, the debugger will continue until the called function returns."""
         self._internal_debugger.next(self)
 
     def si(self: ThreadContext) -> None:
@@ -288,6 +288,5 @@ def fin(self: ThreadContext, heuristic: str = "backtrace") -> None:
         self._internal_debugger.finish(self, heuristic)
 
     def ni(self: ThreadContext) -> None:
-        """Alias for the `next` method. Executes the next instruction of the process. If the instruction is a call, the debugger will continue until the called function returns.
-        """
+        """Alias for the `next` method. Executes the next instruction of the process. If the instruction is a call, the debugger will continue until the called function returns."""
         self._internal_debugger.next(self)

test/aarch64/scripts/finish_test.py

@@ -234,19 +234,19 @@ def test_heuristic_return_address(self):
         # We need to repeat the check for the three stages of the function preamble
 
         # Get current return address
-        curr_srip = stack_unwinder.get_return_address(d)
+        curr_srip = d.saved_ip
         self.assertEqual(curr_srip, RETURN_POINT_FROM_C)
 
         d.step()
 
         # Get current return address
-        curr_srip = stack_unwinder.get_return_address(d)
+        curr_srip = d.saved_ip
         self.assertEqual(curr_srip, RETURN_POINT_FROM_C)
 
         d.step()
 
         # Get current return address
-        curr_srip = stack_unwinder.get_return_address(d)
+        curr_srip = d.saved_ip
         self.assertEqual(curr_srip, RETURN_POINT_FROM_C)
 
         d.kill()

test/amd64/scripts/finish_test.py

@@ -234,19 +234,19 @@ def test_heuristic_return_address(self):
         # We need to repeat the check for the three stages of the function preamble
 
         # Get current return address
-        curr_srip = stack_unwinder.get_return_address(d)
+        curr_srip = d.saved_ip
         self.assertEqual(curr_srip, RETURN_POINT_FROM_C)
 
         d.step()
 
         # Get current return address
-        curr_srip = stack_unwinder.get_return_address(d)
+        curr_srip = d.saved_ip
         self.assertEqual(curr_srip, RETURN_POINT_FROM_C)
 
         d.step()
 
         # Get current return address
-        curr_srip = stack_unwinder.get_return_address(d)
+        curr_srip = d.saved_ip
         self.assertEqual(curr_srip, RETURN_POINT_FROM_C)
 
         d.kill()