Adds `break-if-taken` and `break-if-not-taken` (#1799)

This commit adds the `break-if-taken` and `break-if-not-taken` commands,
which attach breakpoints to branch instructions that will stop the
inferior if said branch is taken or is not taken, respectively. It adds
an extra class, `pwndbg.gdblib.bpoint.Breakpoint`, which clears caches
before calling `stop()`, allowing for the use of register values inside
that function in breakpoint classes that derive from it. Additionally,
checking of whether the conditions for a branch to be taken have been
fulfilled is done through `DisassemblyAssistant.condition()`.

pwndbg/commands/__init__.py

@@ -613,6 +613,7 @@ def load_commands() -> None:
     import pwndbg.commands.asm
     import pwndbg.commands.attachp
     import pwndbg.commands.auxv
+    import pwndbg.commands.branch
     import pwndbg.commands.canary
     import pwndbg.commands.checksec
     import pwndbg.commands.comments

pwndbg/commands/branch.py

@@ -0,0 +1,111 @@
+import argparse
+
+import gdb
+from capstone import CS_GRP_JUMP
+
+import pwndbg.color.message as message
+import pwndbg.commands
+import pwndbg.disasm
+import pwndbg.gdblib.bpoint
+import pwndbg.gdblib.next
+
+
+class BreakOnConditionalBranch(pwndbg.gdblib.bpoint.Breakpoint):
+    """
+    A breakpoint that only stops the inferior if a given branch is taken or not taken.
+    """
+
+    def __init__(self, instruction, taken):
+        super().__init__("*%#x" % instruction.address, type=gdb.BP_BREAKPOINT, internal=False)
+        self.instruction = instruction
+        self.taken = taken
+
+    def should_stop(self):
+        # Use the assistant to figure out which if all the conditions this
+        # branch requires in order to be taken have been met.
+        assistant = pwndbg.disasm.arch.DisassemblyAssistant.for_current_arch()
+        condition_met = assistant.condition(self.instruction)
+        if condition_met is None:
+            # This branch is unconditional.
+            condition_met = 1
+        condition_met = condition_met != 0
+
+        return condition_met == self.taken
+
+
+parser = argparse.ArgumentParser(description="Breaks on a branch if it is taken.")
+parser.add_argument(
+    "branch",
+    type=str,
+    help="The branch instruction to break on.",
+)
+
+
+@pwndbg.commands.ArgparsedCommand(parser, command_name="break-if-taken")
+@pwndbg.commands.OnlyWhenRunning
+def break_if_taken(branch) -> None:
+    install_breakpoint(branch, taken=True)
+
+
+parser = argparse.ArgumentParser(description="Breaks on a branch if it is not taken.")
+parser.add_argument(
+    "branch",
+    type=str,
+    help="The branch instruction to break on.",
+)
+
+
+@pwndbg.commands.ArgparsedCommand(parser, command_name="break-if-not-taken")
+@pwndbg.commands.OnlyWhenRunning
+def break_if_not_taken(branch) -> None:
+    install_breakpoint(branch, taken=False)
+
+
+def install_breakpoint(branch, taken):
+    # Do our best to interpret branch as an address locspec. Untimately, though,
+    # we're limited in what we can do from inside Python in that front.
+    #
+    # https://sourceware.org/gdb/onlinedocs/gdb/Address-Locations.html#Address-Locations
+    address = None
+    try:
+        # Try to interpret branch as an address literal
+        address = int(branch, 0)
+    except ValueError:
+        # That didn't work. Defer to GDB and see if it can make something out of
+        # the address value we were given.
+        try:
+            value = gdb.parse_and_eval(branch)
+            if value.address is None:
+                print(message.warn(f"Value {branch} has no address, trying its value"))
+                address = int(value)
+            else:
+                address = int(value.address)
+        except gdb.error as e:
+            # No such luck. Report to the user and quit.
+            print(message.error(f"Could not resolve branch location {branch}: {e}"))
+            return
+
+    # We should've picked something by now, or errored out.
+    instruction = pwndbg.disasm.one(address)
+    if instruction is None:
+        print(message.error(f"Could not decode instruction at address {address:#x}"))
+        return
+    if CS_GRP_JUMP not in instruction.groups:
+        print(
+            message.error(
+                f"Instruction '{instruction.mnemonic} {instruction.op_str}' at address {address:#x} is not a branch"
+            )
+        )
+        return
+
+    # Not all architectures have assistants we can use for conditionals.
+    if pwndbg.disasm.arch.DisassemblyAssistant.for_current_arch() is None:
+        print(
+            message.error(
+                "The current architecture is not supported for breaking on conditional branches"
+            )
+        )
+        return
+
+    # Install the breakpoint.
+    BreakOnConditionalBranch(instruction, taken)

pwndbg/disasm/arch.py

@@ -36,6 +36,10 @@ class DisassemblyAssistant:
             CS_OP_MEM: self.memory_sz,
         }
 
+    @staticmethod
+    def for_current_arch():
+        return DisassemblyAssistant.assistants.get(pwndbg.gdblib.arch.current, None)
+
     @staticmethod
     def enhance(instruction) -> None:
         enhancer = DisassemblyAssistant.assistants.get(

pwndbg/gdblib/__init__.py

@@ -20,6 +20,7 @@ def load_gdblib() -> None:
     import pwndbg.gdblib.abi
     import pwndbg.gdblib.android
     import pwndbg.gdblib.argv
+    import pwndbg.gdblib.bpoint
     import pwndbg.gdblib.ctypes
     import pwndbg.gdblib.elf
     import pwndbg.gdblib.events

pwndbg/gdblib/bpoint.py

@@ -0,0 +1,26 @@
+import gdb
+
+import pwndbg.lib.cache
+
+
+class Breakpoint(gdb.Breakpoint):
+    """
+    Breakpoint class, similar to gdb.Breakpoint, but clears the caches
+    associated with the stop event before determining whether it should stop the
+    inferior or not.
+
+    Unlike gdb.Breakpoint, users of this class should override `should_stop()`,
+    instead of `stop()`, as the latter is used to do cache invalidation.
+    """
+
+    def stop(self):
+        # Clear the cache for the stop event.
+        pwndbg.gdblib.regs.__getattr__.cache.clear()
+        return self.should_stop()
+
+    def should_stop(self):
+        """
+        This function is called whenever this breakpoint is hit in the code and
+        its return value determines whether the inferior will be stopped.
+        """
+        return True

tests/gdb-tests/tests/binaries/conditional_branch_breakpoints_x64.asm

@@ -0,0 +1,41 @@
+section .text
+    global _start
+    global break_here
+    global break_here0
+    global break_here1
+    global branch0
+    global branch1
+    global branch2
+    global branch3
+
+_start:
+break_here:
+    mov rax, 0
+    cmp rax, 0
+branch0:
+    ; Break on branch taken. Branch will be taken. (test for PC=branch0)
+    jz branch0_done
+branch0_done:
+    nop
+branch1:
+    ; Break on branch taken. Branch will not be taken. (test for PC=break_here0)
+    jnz branch1_done
+branch1_done:
+break_here0:
+    mov rax, 10
+    cmp rax, 0
+branch2:
+    ; Break on branch not taken. Branch will be taken. (test for PC=break_here1)
+    jne branch2_done
+branch2_done:
+break_here1:
+    nop
+branch3:
+    ; Break on branch not taken. Branch will not be taken. (test for PC=branch3)
+    je branch3_done
+branch3_done:
+exit:
+    ; Call sys_exit(0) on Linux.
+    mov rax, 60
+    mov rsi, 0
+    syscall

tests/gdb-tests/tests/test_command_branch.py

@@ -0,0 +1,36 @@
+import gdb
+import pytest
+
+import pwndbg.gdblib
+import tests
+
+CONDBR_X64_BINARY = tests.binaries.get("conditional_branch_breakpoints_x64.out")
+
+
+@pytest.mark.parametrize("binary", [CONDBR_X64_BINARY], ids=["x86-64"])
+def test_command_break_if_x64(start_binary, binary):
+    """
+    Tests the chain for a non-nested linked list
+    """
+
+    start_binary(binary)
+    gdb.execute("break break_here")
+    gdb.execute("run")
+
+    gdb.execute("break break_here0")
+    gdb.execute("break break_here1")
+    gdb.execute("break-if-taken branch0")
+    gdb.execute("break-if-taken branch1")
+    gdb.execute("break-if-not-taken branch2")
+    gdb.execute("break-if-not-taken branch3")
+
+    continue_and_test_pc("branch0")
+    continue_and_test_pc("break_here0")
+    continue_and_test_pc("break_here1")
+    continue_and_test_pc("branch3")
+
+
+def continue_and_test_pc(stop_label):
+    gdb.execute("continue")
+    address = int(gdb.parse_and_eval(f"&{stop_label}"))
+    assert pwndbg.gdblib.regs.pc == address