From 2b9beef7afc99db312667ac9bf90f866596287b2 Mon Sep 17 00:00:00 2001
From: Lennard Hofmann <lennard.hofmann@web.de>
Date: Wed, 22 May 2024 14:10:52 +0200
Subject: [PATCH] Small fixes for RISC-V (#2172)

* Fix i386-32 syscall name printing

pwndbg-git from AUR shows hexadecimal constants in masm syntax
(e.g. 80h) for some reason (as if the option CS_OPT_SYNTAX_MASM was set).
This commit makes syscall name printing work regardless of hex syntax.

* riscv: Fix AssertionError on "jalr ra, ra, 0x252"

When the PC was on this instruction, the pwndbg context would not be
printed due to this AssertionError.

* riscv: Fix AssertionError on "c.jalr a5"

According to the specification, "C.JALR expands to jalr x1, 0(rs1)".
---
 pwndbg/arguments.py                         | 16 +++++-----------
 pwndbg/disasm/riscv.py                      |  8 ++++----
 tests/qemu-tests/tests/user/test_riscv64.py | 15 ++++++++++++---
 3 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/pwndbg/arguments.py b/pwndbg/arguments.py
index 0b4006eba70..ee9e90d51ac 100644
--- a/pwndbg/arguments.py
+++ b/pwndbg/arguments.py
@@ -15,6 +15,7 @@
 import pwndbg.chain
 import pwndbg.constants
 import pwndbg.disasm
+import pwndbg.disasm.arch
 import pwndbg.gdblib.arch
 import pwndbg.gdblib.file
 import pwndbg.gdblib.memory
@@ -71,7 +72,7 @@ def get_syscall_name(instruction: PwndbgInstruction) -> str | None:
     if syscall_register in ("eax", "rax"):
         mnemonic = instruction.mnemonic
 
-        is_32bit = mnemonic == "int" and instruction.op_str == "0x80"
+        is_32bit = mnemonic == "int" and instruction.operands[0].before_value == 0x80
         if not (mnemonic == "syscall" or is_32bit):
             return None
 
@@ -87,7 +88,7 @@ def get_syscall_name(instruction: PwndbgInstruction) -> str | None:
 def get(instruction: PwndbgInstruction) -> List[Tuple[pwndbg.lib.functions.Argument, int]]:
     """
     Returns an array containing the arguments to the current function,
-    if $pc is a 'call' or 'bl' type instruction.
+    if $pc is a 'call', 'bl', or 'jalr' type instruction.
 
     Otherwise, returns None.
     """
@@ -105,19 +106,12 @@ def get(instruction: PwndbgInstruction) -> List[Tuple[pwndbg.lib.functions.Argum
         except KeyError:
             return []
 
-        # Not sure of any OS which allows multiple operands on
-        # a call instruction.
-        assert len(instruction.operands) == 1
-
-        target = instruction.operands[0].before_value
+        assistant = pwndbg.disasm.arch.DisassemblyAssistant.for_current_arch()
+        target = assistant.resolve_target(instruction, None, call=True)
 
         if not target:
             return []
 
-        if pwndbg.gdblib.arch.name in ["rv32", "rv64"]:
-            target += instruction.address
-            target &= pwndbg.gdblib.arch.ptrmask
-
         name = pwndbg.gdblib.symbol.get(target)
         if not name:
             return []
diff --git a/pwndbg/disasm/riscv.py b/pwndbg/disasm/riscv.py
index 7500a6cdabf..c8d0bdcb157 100644
--- a/pwndbg/disasm/riscv.py
+++ b/pwndbg/disasm/riscv.py
@@ -95,10 +95,10 @@ def resolve_target(self, instruction: PwndbgInstruction, emu: Emulator | None, c
 
         # Determine the target address of the indirect jump
         if instruction.id in [RISCV_INS_JALR, RISCV_INS_C_JALR]:
-            target = (
-                instruction.op_find(CS_OP_REG, 1).before_value
-                + instruction.op_find(CS_OP_IMM, 1).imm
-            ) & ptrmask
+            target = instruction.op_find(CS_OP_REG, 1).before_value
+            if instruction.id == RISCV_INS_JALR:
+                target += instruction.op_find(CS_OP_IMM, 1).imm
+            target &= ptrmask
             # Clear the lowest bit without knowing the register width
             return target ^ (target & 1)
 
diff --git a/tests/qemu-tests/tests/user/test_riscv64.py b/tests/qemu-tests/tests/user/test_riscv64.py
index 0bcf9e52d88..48961fbd18d 100644
--- a/tests/qemu-tests/tests/user/test_riscv64.py
+++ b/tests/qemu-tests/tests/user/test_riscv64.py
@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import re
 import sys
 import traceback
 
@@ -13,11 +12,21 @@
     assert pwndbg.gdblib.symbol.address("main") == 0x4000000668
     gdb.execute("continue")
 
-    gdb.execute("nextcall", to_string=True)
+    gdb.execute("stepuntilasm jalr")
 
     # verify call argument are enriched
     assembly = gdb.execute("nearpc", to_string=True)
-    assert re.search(r"s.*'Not enough args'", assembly), assembly
+    assert "'Not enough args'" in assembly
+
+    gdb.execute("stepuntilasm c.jalr")
+
+    # verify jump target is correct
+    assembly = gdb.execute("nearpc 0", to_string=True)
+    target = assembly.splitlines()[0].split()[-1]
+    gdb.execute("stepi")
+    assembly = gdb.execute("nearpc 0", to_string=True)
+    assert assembly.split()[2] == target, (assembly.split()[2], target)
+
 except AssertionError:
     traceback.print_exc(file=sys.stdout)
     sys.exit(1)