Basic support for arm64

README.md

@@ -7,7 +7,7 @@ Gives complete trace of a function execution including all sub-calls.
 * Attach to a process, for example to `SimpleTraceTarget` contained in this repository: ```$ lldb -p `pgrep SimpleTrace` ```
 * Import the trace script: ```(lldb) command script import $DIR/trace.py```
 * Stop the process execution at the root function you're interested in tracing, e.g. using breakpoints.
-* Trace: ```(lldb) trace```
+* Trace: ```(lldb) function-trace```
 * Output: 
 ```
 a + 0x1c ==> b
@@ -45,7 +45,7 @@ a + 0x1c ==> b
  where `==>` denotes call to a function, `===` jmp to a different symbol, `<==` is a return, `Syscall ID` is where a syscall of a given `ID` is executed.
 
 ### Help
- `(lldb) trace -h` gives a list of options `trace` accepts, currently:
+ `(lldb) function-trace -h` gives a list of options `function-trace` accepts, currently:
  ```
 Options:
   -h, --help            show this help message and exit
@@ -57,17 +57,17 @@ Options:
   ```
 
 ### More advanced usage
- * `trace` can trace any symbol in the current execution backtrace. Select the desired frame with `(lldb) frame select FRAME_ID` before executing `trace`
- * Long-running `trace` doesn't produce any output by default, until the command execution is finished (seems to be by design). `trace` makes it possible to produce incremental output by either outputing to a specified file `(lldb) trace -f FILE` or by outputting directly to stdout (this may not always work properly since it's against the LLDB's policy) `(lldb) trace -s`
- * If you're not interested in calling external symbols (external to the binary module in which the root symbol is defined) use `(lldb) trace -m`. Every time extrenal symbol is detected appropriate message (such as `Not instrumenting since module X isn't the same as Y`) will be printed.
+ * `function-trace` can trace any symbol in the current execution backtrace. Select the desired frame with `(lldb) frame select FRAME_ID` before executing `function-trace`
+ * Long-running `function-trace` doesn't produce any output by default, until the command execution is finished (seems to be by design). `function-trace` makes it possible to produce incremental output by either outputing to a specified file `(lldb) function-trace -f FILE` or by outputting directly to stdout (this may not always work properly since it's against the LLDB's policy) `(lldb) function-trace -s`
+ * If you're not interested in calling external symbols (external to the binary module in which the root symbol is defined) use `(lldb) function-trace -m`. Every time extrenal symbol is detected appropriate message (such as `Not instrumenting since module X isn't the same as Y`) will be printed.
 
 ### Debugging
-* Verbose output is available with `(lldb) trace -v`
+* Verbose output is available with `(lldb) function-trace -v`
 * If tracing fails, some of the breakpoints established for tracing may be left over. They are usually easy to distinguish from other breakpoinst since they are per thread (`Options: enabled tid: X`). All breakpoints can be deleted with `(lldb) breakpoint delete`
-* `Ctrl-C` in LLDB session should stop the tracing, in some circumstances however (if the program managed to break-out from the `trace`'s control), tracing may have to be stopped by sending the target process a singnal (e.g. `$ pkill X`).
+* `Ctrl-C` in LLDB session should stop the tracing, in some circumstances however (if the program managed to break-out from the `function-trace`'s control), tracing may have to be stopped by sending the target process a singnal (e.g. `$ pkill X`).
 
 ### TODOs
-* `(lldb) trace -m` should give an ability to trace over `objc_sendMsg`
+* `(lldb) function-trace -m` should give an ability to trace over `objc_sendMsg`
 * LLDB config for autoimporting
 * wrapper script for running from breakpoint command (`breakpoint command add -F trace.trace_wrapper ID`)
 

trace.py

@@ -79,13 +79,50 @@ def __init__(self, target, thread, frame):
         self.jmp_breakpoints = {}
         self.syscall_breakpoints = {}
         self.subsequent_instruction = {}
+        self.arch = self.target.triple.split("-")[0]
 
     def update_frame(self, frame):
         self.frame = frame
 
     def is_frame_valid(self):
         return self.frame.IsValid()
 
+    ptrauth_suffixes = ["", "aa", "aaz", "ab", "abz"]
+
+    @staticmethod
+    def check_ptrauth_instruction(base, mnemonic):
+        for suffix in __class__.ptrauth_suffixes:
+            if mnemonic == base + suffix:
+                return True
+        return False
+
+    def is_call_instruction(self, mnemonic):
+        if self.arch == "arm64":
+            return mnemonic == "bl" or __class__.check_ptrauth_instruction("blr", mnemonic)
+        elif self.arch == "x86_64":
+            return mnemonic.startswith('call')
+        else:
+            print(f"Unhandled arch {arch}")
+            return False
+
+    def is_branch_instruction(self, mnemonic):
+        if self.arch == "arm64":
+            return mnemonic == "b" or __class__.check_ptrauth_instruction("br", mnemonic)
+        elif self.arch == "x86_64":
+            return mnemonic.startswith('jmp')
+        else:
+            print(f"Unhandled arch {arch}")
+            return False
+
+    def is_syscall_instruction(self, mnemonic):
+        if self.arch == "arm64":
+            return mnemonic.startswith('svc')
+        elif self.arch == "x86_64":
+            return mnemonic.startswith('syscall')
+        else:
+            print(f"Unhandled arch {arch}")
+            return False
+
     def instrument_calls_syscalls_and_jmps(self):
         # TODO: symbols vs functions
         symbol = self.frame.GetSymbol()
@@ -108,13 +145,13 @@ def instrument_calls_syscalls_and_jmps(self):
                 self.subsequent_instruction[previous_breakpoint_address] = address
                 previous_breakpoint_address = 0
             mnemonic = i.GetMnemonic(self.target)
-            if mnemonic is not None and mnemonic.startswith('call'):
+            if mnemonic is not None and self.is_call_instruction(mnemonic):
                 log_v('Putting call breakpoint at 0x%lx' % address)
                 breakpoint = self.target.BreakpointCreateByAddress(address)
                 breakpoint.SetThreadID(self.thread.GetThreadID())
                 self.call_breakpoints[address] = breakpoint
                 previous_breakpoint_address = address
-            if mnemonic is not None and mnemonic.startswith('jmp'):
+            if mnemonic is not None and self.is_branch_instruction(mnemonic):
                 try:
                     jmp_destination = int(i.GetOperands(self.target), 16)
                 except:
@@ -125,7 +162,7 @@ def instrument_calls_syscalls_and_jmps(self):
                     breakpoint = self.target.BreakpointCreateByAddress(address)
                     breakpoint.SetThreadID(self.thread.GetThreadID())
                     self.jmp_breakpoints[address] = breakpoint
-            if mnemonic is not None and mnemonic.startswith('syscall'):
+            if mnemonic is not None and self.is_syscall_instruction(mnemonic):
                 log_v('Putting syscall breakpoint at 0x%lx' % address)
                 breakpoint = self.target.BreakpointCreateByAddress(address)
                 breakpoint.SetThreadID(self.thread.GetThreadID())
@@ -220,7 +257,7 @@ def __init__(self, result):
         self.exited = False
 
     def get_prog_name(self):
-        return "trace"
+        return "function-trace"
 
     def exit(self, status=0, msg=None):
         if msg is not None:
@@ -422,10 +459,18 @@ def trace(debugger: lldb.SBDebugger, command: str, result: lldb.SBCommandReturnO
         elif instrumented_frame.is_stopped_on_syscall(frame):
             rax = -1;
             register_sets = frame.GetRegisters()
+            arch = target.triple.split("-")[0]
+            if arch == "arm64":
+                syscall_register = "x0"
+            elif arch == "x86_64":
+                syscall_register = "rax"
+            else:
+                print(f"Unhandled arch {arch}")
+                syscall_register = ""
             for register_set in register_sets:
                 if register_set.GetName() == "General Purpose Registers":
                     for register in register_set:
-                        if register.GetName() == "rax":
+                        if register.GetName() == syscall_register:
                             rax = register.GetValue()
             log('{} Syscall {}'.format(spacer * depth, rax))
             thread.StepInstruction(False)
@@ -463,5 +508,5 @@ def trace(debugger: lldb.SBDebugger, command: str, result: lldb.SBCommandReturnO
 
 # And the initialization code to add your commands
 def __lldb_init_module(debugger, internal_dict):
-    debugger.HandleCommand('command script add -f trace.trace trace')
-    print('The "trace" python command has been installed and is ready for use.')
+    debugger.HandleCommand('command script add -f trace.trace function-trace')
+    print('The "function-trace" python command has been installed and is ready for use.')