Fix stepsyscall + add a test (#2884)

pwndbg/aglib/next.py

@@ -13,23 +13,28 @@ import capstone
 import pwndbg.aglib.disasm
 import pwndbg.aglib.proc
 import pwndbg.aglib.regs
+from pwndbg.aglib.disasm.instruction import PwndbgInstruction
 from pwndbg.color import message
 from pwndbg.dbg import BreakpointLocation
 
 interrupts = {capstone.CS_GRP_INT}
 
 
-def next_int(address=None):
+def next_int(address=None, honor_current_branch=False):
     """
     If there is a syscall in the current basic black,
     return the instruction of the one closest to $PC.
 
-    Otherwise, return None.
+    If honor_current_branch is True, then if the address is already a branch, return None.
+
+    If no interrupt exists or a jump is in the way, return None.
     """
     if address is None:
         ins = pwndbg.aglib.disasm.one(pwndbg.aglib.regs.pc)
         if not ins:
             return None
+        if honor_current_branch and ins.jump_like:
+            return None
         address = ins.next
 
     ins = pwndbg.aglib.disasm.one(address)
@@ -43,11 +48,18 @@ def next_int(address=None):
     return None
 
 
-def next_branch(address=None):
+def next_branch(address=None, including_current=False) -> PwndbgInstruction | None:
+    """
+    Return the next branch instruction that the process will encounter with repeated usage of the "nexti" command.
+
+    If including_current == True, then if the instruction at the address is already a branch, return it.
+    """
     if address is None:
         ins = pwndbg.aglib.disasm.one(pwndbg.aglib.regs.pc)
         if not ins:
             return None
+        if including_current and ins.jump_like:
+            return ins
         address = ins.next
 
     ins = pwndbg.aglib.disasm.one(address)
@@ -101,24 +113,42 @@ def next_matching_until_branch(address=None, mnemonic=None, op_str=None):
     return None
 
 
-async def break_next_branch(ec: pwndbg.dbg_mod.ExecutionController, address=None):
-    ins = next_branch(address)
+async def break_next_branch(
+    ec: pwndbg.dbg_mod.ExecutionController, address=None, including_current=False
+):
+    """
+    If including_current == True, do not step in case we are currently on a branch
+    """
+    ins = next_branch(address, including_current=including_current)
 
     proc = pwndbg.dbg.selected_inferior()
     if ins:
-        with proc.break_at(BreakpointLocation(ins.address), internal=True) as bp:
-            await ec.cont(bp)
+        # If the branch we found was not at the current program counter, we should step to it.
+        # Otherwise, return the current instruction.
+        if ins.address != pwndbg.aglib.regs.pc:
+            with proc.break_at(BreakpointLocation(ins.address), internal=True) as bp:
+                await ec.cont(bp)
         return ins
 
 
-async def break_next_interrupt(ec: pwndbg.dbg_mod.ExecutionController, address=None):
-    ins = next_int(address)
+async def break_next_interrupt(
+    ec: pwndbg.dbg_mod.ExecutionController, address=None, honor_current_branch=False
+) -> PwndbgInstruction | None:
+    """
+    Break at the next interrupt if there is one in the current basic block
+    and no jumps are between the current instruction and the interrupt.
+
+    If no such interrupt exists or a jump is in the way, return None.
+    """
+    ins = next_int(address, honor_current_branch=honor_current_branch)
     proc = pwndbg.dbg.selected_inferior()
     if ins:
         with proc.break_at(BreakpointLocation(ins.address), internal=True) as bp:
             await ec.cont(bp)
         return ins
 
+    return None
+
 
 async def break_next_call(ec: pwndbg.dbg_mod.ExecutionController, symbol_regex=None):
     symbol_regex = re.compile(symbol_regex) if symbol_regex else None
@@ -179,7 +209,6 @@ async def break_on_next_matching_instruction(
         ins = next_matching_until_branch(mnemonic=mnemonic, op_str=op_str)
         if ins is not None:
             if ins.address != pwndbg.aglib.regs.pc:
-                print("Found instruction")
                 # Only set breakpoints at a different PC location, otherwise we
                 # will continue until we hit a breakpoint that's not related to
                 # this opeeration, or the program halts.
@@ -192,8 +221,7 @@ async def break_on_next_matching_instruction(
                 pass
         else:
             # Move to the next branch instruction.
-            print("Moving to next branch")
-            nb = next_branch(pwndbg.aglib.regs.pc)
+            nb = next_branch(pwndbg.aglib.regs.pc, including_current=True)
             if nb is not None:
                 if nb.address != pwndbg.aglib.regs.pc:
                     # Stop right at the next branch instruction.

pwndbg/commands/next.py

@@ -161,8 +161,8 @@ async def _stepsyscall(ec: pwndbg.dbg_mod.ExecutionController):
 
     while (
         pwndbg.aglib.proc.alive
-        and not (await pwndbg.aglib.next.break_next_interrupt(ec))
-        and (await pwndbg.aglib.next.break_next_branch(ec))
+        and not (await pwndbg.aglib.next.break_next_interrupt(ec, honor_current_branch=True))
+        and (await pwndbg.aglib.next.break_next_branch(ec, including_current=True))
     ):
         # Here we are e.g. on a CALL instruction (temporarily breakpointed by `break_next_branch`)
         # We need to step so that we take this branch instead of ignoring it

tests/gdb-tests/tests/binaries/stepsyscall_x64.asm

@@ -0,0 +1,33 @@
+section .text
+    global _start
+
+
+write:
+    mov rax, 1
+    mov rdi, 1
+    mov rsi, msg
+    mov rdx, len
+syscall_write_label:
+    syscall
+    ret
+
+_start:
+    nop
+    jmp label1
+    nop
+
+label1:
+    call write
+
+exit:
+    mov rax, 60
+    mov rdi, 0
+
+syscall_exit_label:
+    syscall
+
+
+
+section .data
+    msg db 'hello world', 0xA
+    len equ $ - msg

tests/gdb-tests/tests/test_command_stepsyscall.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+import gdb
+
+import pwndbg.aglib.regs
+import tests
+
+STEPSYSCALL_X64_BINARY = tests.binaries.get("stepsyscall_x64.out")
+
+
+def test_command_stepsyscall(start_binary):
+    start_binary(STEPSYSCALL_X64_BINARY)
+
+    # Test that the logic correctly handles multiple consecutive jumps
+    gdb.execute("stepsyscall")
+    address = int(gdb.parse_and_eval("&syscall_write_label"))
+    assert pwndbg.aglib.regs.pc == address
+
+    gdb.execute("stepsyscall")
+    address = int(gdb.parse_and_eval("&syscall_exit_label"))
+    assert pwndbg.aglib.regs.pc == address
+
+
+def test_command_nextsyscall(start_binary):
+    start_binary(STEPSYSCALL_X64_BINARY)
+
+    gdb.execute("nextsyscall")
+    address = int(gdb.parse_and_eval("&syscall_exit_label"))
+    assert pwndbg.aglib.regs.pc == address