Manually resolve AArch64 conditional branches (#2259)

* Resolve conditionals in AArch64 instruction enhancement

* Manually resolve targets

* lint

* Update comment

pwndbg/gdblib/disasm/aarch64.py

@@ -12,7 +12,45 @@ import pwndbg.gdblib.disasm.arch
 import pwndbg.gdblib.memory
 import pwndbg.gdblib.regs
 from pwndbg.emu.emulator import Emulator
+from pwndbg.gdblib.disasm.instruction import ALL_JUMP_GROUPS
+from pwndbg.gdblib.disasm.instruction import InstructionCondition
 from pwndbg.gdblib.disasm.instruction import PwndbgInstruction
+from pwndbg.gdblib.disasm.instruction import boolean_to_instruction_condition
+
+
+def resolve_condition(condition: int, cpsr: int) -> InstructionCondition:
+    """
+    Given a condition and the NZCV flag bits, determine when the condition is satisfied
+
+    The condition is a Capstone constant
+    """
+
+    n = (cpsr >> 31) & 1
+    z = (cpsr >> 30) & 1
+    c = (cpsr >> 29) & 1
+    v = (cpsr >> 28) & 1
+
+    condition = {
+        ARM64_CC_INVALID: True,  # Capstone uses this code for the 'B' instruction, the unconditional branch
+        ARM64_CC_EQ: z == 1,
+        ARM64_CC_NE: z == 0,
+        ARM64_CC_HS: c == 1,
+        ARM64_CC_LO: c == 0,
+        ARM64_CC_MI: n == 1,
+        ARM64_CC_PL: n == 0,
+        ARM64_CC_VS: v == 1,
+        ARM64_CC_VC: v == 0,
+        ARM64_CC_HI: c == 1 and z == 0,
+        ARM64_CC_LS: not (c == 1 and z == 0),
+        ARM64_CC_GE: n == v,
+        ARM64_CC_LT: n != v,
+        ARM64_CC_GT: z == 0 and n == v,
+        ARM64_CC_LE: not (z == 0 and n == v),
+        ARM64_CC_AL: True,
+        ARM64_CC_NV: True,
+    }.get(condition, False)
+
+    return InstructionCondition.TRUE if condition else InstructionCondition.FALSE
 
 
 class DisassemblyAssistant(pwndbg.gdblib.disasm.arch.DisassemblyAssistant):
@@ -64,6 +102,61 @@ class DisassemblyAssistant(pwndbg.gdblib.disasm.arch.DisassemblyAssistant):
 
             instruction.annotation = f"{left.str} => {super()._telescope_format_list(telescope_addresses, TELESCOPE_DEPTH, emu)}"
 
+    @override
+    def _condition(
+        self, instruction: PwndbgInstruction, emu: Emulator
+    ) -> pwndbg.gdblib.disasm.arch.InstructionCondition:
+        # In ARM64, only branches have the conditional code in the instruction,
+        # as opposed to ARM32 which allows most instructions to be conditional
+        if instruction.id == ARM64_INS_B:
+            flags = super()._read_register_name(instruction, "cpsr", emu)
+            if flags is not None:
+                return resolve_condition(instruction.cs_insn.cc, flags)
+
+        elif instruction.id == ARM64_INS_CBNZ:
+            op_val = instruction.operands[0].before_value
+            return boolean_to_instruction_condition(op_val is not None and op_val != 0)
+
+        elif instruction.id == ARM64_INS_CBZ:
+            op_val = instruction.operands[0].before_value
+            return boolean_to_instruction_condition(op_val is not None and op_val == 0)
+
+        elif instruction.id == ARM64_INS_TBNZ:
+            op_val, bit = (
+                instruction.operands[0].before_value,
+                instruction.operands[1].before_value,
+            )
+
+            if op_val is not None and bit is not None:
+                return boolean_to_instruction_condition(bool((op_val >> bit) & 1))
+
+        elif instruction.id == ARM64_INS_TBZ:
+            op_val, bit = (
+                instruction.operands[0].before_value,
+                instruction.operands[1].before_value,
+            )
+
+            if op_val is not None and bit is not None:
+                return boolean_to_instruction_condition(not ((op_val >> bit) & 1))
+
+        # TODO: Additionally, the "conditional comparisons" and "conditional selects" support conditional execution
+
+        return super()._condition(instruction, emu)
+
+    @override
+    def _resolve_target(self, instruction: PwndbgInstruction, emu: Emulator | None, call=False):
+        if not bool(instruction.groups_set & ALL_JUMP_GROUPS):
+            return None
+
+        if len(instruction.operands) > 0:
+            # For all AArch64 branches, the target is either an immediate or a register and is the last operand
+            return instruction.operands[-1].before_value
+        elif instruction.id == ARM64_INS_RET:
+            # If this is a ret WITHOUT an operand, it means we should read from the LR/x30 register
+            return super()._read_register_name(instruction, "lr", emu)
+
+        return super()._resolve_target(instruction, emu, call)
+
     @override
     def _set_annotation_string(self, instruction: PwndbgInstruction, emu: Emulator) -> None:
         # Dispatch to the correct handler

pwndbg/gdblib/disasm/instruction.py

@@ -81,6 +81,10 @@ class InstructionCondition(Enum):
     UNDETERMINED = 3
 
 
+def boolean_to_instruction_condition(condition: bool) -> InstructionCondition:
+    return InstructionCondition.TRUE if condition else InstructionCondition.FALSE
+
+
 # Only use within the instruction.__repr__ to give a nice output
 CAPSTONE_ARCH_MAPPING_STRING = {
     CS_ARCH_ARM: "arm",
@@ -318,7 +322,7 @@ class PwndbgInstruction:
     def __repr__(self) -> str:
         operands_str = " ".join([repr(op) for op in self.operands])
 
-        return f"""{self.mnemonic} {self.op_str} at {self.address:#x} (size={self.size}) (arch: {CAPSTONE_ARCH_MAPPING_STRING.get(self.cs_insn._cs.arch,None)})
+        info = f"""{self.mnemonic} {self.op_str} at {self.address:#x} (size={self.size}) (arch: {CAPSTONE_ARCH_MAPPING_STRING.get(self.cs_insn._cs.arch,None)})
         ID: {self.id}, {self.cs_insn.insn_name()}
         Raw asm: {'%-06s %s' % (self.mnemonic, self.op_str)}
         New asm: {self.asm_string}
@@ -333,6 +337,12 @@ class PwndbgInstruction:
         Can change PC: {self.can_change_instruction_pointer}
         Syscall: {self.syscall if self.syscall is not None else ""} {self.syscall_name if self.syscall_name is not None else "N/A"}"""
 
+        # Hacky, but this is just for debugging
+        if hasattr(self.cs_insn, "cc"):
+            info += f"\n\tARM condition code: {self.cs_insn.cc}"
+
+        return info
+
 
 class EnhancedOperand:
     def __init__(self, cs_op):