Store Instruction Annotations for RISCV, Arm, AArch64, and MIPS (#2363)

* merge store annotations for all arches

* x86 fix

* lint

pwndbg/gdblib/disasm/aarch64.py

@@ -49,6 +49,7 @@ AARCH64_SINGLE_LOAD_INSTRUCTIONS: Dict[int, int | None] = {
     ARM64_INS_LDAR: None,
 }
 
+# None indicates that the write size depends on the source register
 AARCH64_SINGLE_STORE_INSTRUCTIONS: Dict[int, int | None] = {
     ARM64_INS_STRB: 1,
     ARM64_INS_STURB: 1,
@@ -177,6 +178,15 @@ class DisassemblyAssistant(pwndbg.gdblib.disasm.arch.DisassemblyAssistant):
                 instruction.operands[0].str,
                 instruction.operands[1].str,
             )
+        elif instruction.id in AARCH64_SINGLE_STORE_INSTRUCTIONS:
+            self._common_store_annotator(
+                instruction,
+                emu,
+                instruction.operands[1].before_value,
+                instruction.operands[0].before_value,
+                AARCH64_SINGLE_STORE_INSTRUCTIONS[instruction.id],
+                instruction.operands[1].str,
+            )
         else:
             self.annotation_handlers.get(instruction.id, lambda *a: None)(instruction, emu)
 

pwndbg/gdblib/disasm/arch.py

@@ -908,5 +908,48 @@ class DisassemblyAssistant:
             if telescope_print is not None:
                 instruction.annotation += f" => {telescope_print}"
 
+    def _common_store_annotator(
+        self,
+        instruction: PwndbgInstruction,
+        emu: Emulator,
+        address: int | None,
+        value: int | None,
+        write_size: int | None,
+        address_str: str,
+    ) -> None:
+        """
+        This function annotates store functions - moving data from a register to memory.
+
+        The `value` is truncated to match the `write_size`, if `write_size` is not None.
+
+        The annotation will indicate if the instruction will segfault.
+
+        `write_size`: number of bytes of `value` that will be written
+        """
+
+        if address is None:
+            return
+
+        if not pwndbg.gdblib.memory.peek(address):
+            instruction.annotation = MessageColor.error(
+                f"<Cannot dereference [{MemoryColor.get(address)}]>"
+            )
+        elif value is not None:
+            # To make this annotation work with emulation disabled,
+            # we telescope the value that is going to be placed in the memory operand
+            TELESCOPE_DEPTH = max(0, int(pwndbg.config.disasm_telescope_depth))
+
+            if write_size is not None:
+                value &= (1 << (write_size * 8)) - 1
+
+            telescope_addresses = self._telescope(
+                value,
+                TELESCOPE_DEPTH,
+                instruction,
+                emu,
+            )
+
+            instruction.annotation = f"{address_str} => {self._telescope_format_list(telescope_addresses, TELESCOPE_DEPTH, emu)}"
+
 
 generic_assistant = DisassemblyAssistant(None)

pwndbg/gdblib/disasm/arm.py

@@ -85,6 +85,15 @@ class DisassemblyAssistant(pwndbg.gdblib.disasm.arch.DisassemblyAssistant):
                 instruction.operands[0].str,
                 instruction.operands[1].str,
             )
+        elif instruction.id in ARM_SINGLE_STORE_INSTRUCTIONS:
+            self._common_store_annotator(
+                instruction,
+                emu,
+                instruction.operands[1].before_value,
+                instruction.operands[0].before_value,
+                ARM_SINGLE_STORE_INSTRUCTIONS[instruction.id],
+                instruction.operands[1].str,
+            )
         else:
             self.annotation_handlers.get(instruction.id, lambda *a: None)(instruction, emu)
 

pwndbg/gdblib/disasm/mips.py

@@ -113,6 +113,13 @@ MIPS_LOAD_INSTRUCTIONS = {
     MIPS_INS_LDPC: 8,
 }
 
+MIPS_STORE_INSTRUCTIONS = {
+    MIPS_INS_SB: 1,
+    MIPS_INS_SH: 2,
+    MIPS_INS_SW: 4,
+    MIPS_INS_SD: 8,
+}
+
 
 # This class enhances 32-bit, 64-bit, and micro MIPS
 class DisassemblyAssistant(pwndbg.gdblib.disasm.arch.DisassemblyAssistant):
@@ -134,6 +141,15 @@ class DisassemblyAssistant(pwndbg.gdblib.disasm.arch.DisassemblyAssistant):
                 instruction.operands[0].str,
                 instruction.operands[1].str,
             )
+        elif instruction.id in MIPS_STORE_INSTRUCTIONS:
+            self._common_store_annotator(
+                instruction,
+                emu,
+                instruction.operands[1].before_value,
+                instruction.operands[0].before_value,
+                MIPS_STORE_INSTRUCTIONS[instruction.id],
+                instruction.operands[1].str,
+            )
         elif instruction.id in MIPS_SIMPLE_DESTINATION_INSTRUCTIONS:
             self._common_generic_register_destination(instruction, emu)
 

pwndbg/gdblib/disasm/riscv.py

@@ -42,7 +42,9 @@ RISCV_STORE_INSTRUCTIONS = {
 # TODO: remove this when updating to Capstone 6
 RISCV_COMPRESSED_STORE_INSTRUCTIONS = {
     RISCV_INS_C_SW: 4,
+    RISCV_INS_C_SWSP: 4,
     RISCV_INS_C_SD: 8,
+    RISCV_INS_C_SDSP: 8,
 }
 
 
@@ -87,6 +89,31 @@ class DisassemblyAssistant(pwndbg.gdblib.disasm.arch.DisassemblyAssistant):
                     dest_str,
                 )
 
+        if instruction.id in RISCV_STORE_INSTRUCTIONS:
+            self._common_store_annotator(
+                instruction,
+                emu,
+                instruction.operands[1].before_value,
+                instruction.operands[0].before_value,
+                RISCV_STORE_INSTRUCTIONS[instruction.id],
+                instruction.operands[1].str,
+            )
+        elif instruction.id in RISCV_COMPRESSED_STORE_INSTRUCTIONS:
+            # TODO: remove this branch when updating to Capstone 6
+            address = self._resolve_compressed_target_addr(instruction, emu)
+
+            if address is not None:
+                dest_str = f"[{MemoryColor.get_address_or_symbol(address)}]"
+
+                self._common_store_annotator(
+                    instruction,
+                    emu,
+                    address,
+                    instruction.operands[0].before_value,
+                    RISCV_COMPRESSED_STORE_INSTRUCTIONS[instruction.id],
+                    dest_str,
+                )
+
         return super()._set_annotation_string(instruction, emu)
 
     def _resolve_compressed_target_addr(

pwndbg/gdblib/disasm/x86.py

@@ -91,11 +91,20 @@ class DisassemblyAssistant(pwndbg.gdblib.disasm.arch.DisassemblyAssistant):
                 left.str,
                 right.str,
             )
-        # Handle other cases of MOV
-        elif right.before_value is not None:
+        elif left.type == CS_OP_MEM:
+            # Store operation, MOV [MEM], REG|IMM
+            self._common_store_annotator(
+                instruction,
+                emu,
+                instruction.operands[0].before_value,
+                instruction.operands[1].before_value,
+                right.cs_op.size,
+                instruction.operands[0].str,
+            )
+        elif left.type == CS_OP_REG and right.before_value is not None:
+            # MOV REG, REG|IMM
             TELESCOPE_DEPTH = max(0, int(pwndbg.config.disasm_telescope_depth))
 
-            # +1 to ensure we telescope enough to read at least one address for the last "elif" below
             telescope_addresses = super()._telescope(
                 right.before_value,
                 TELESCOPE_DEPTH + 1,
@@ -106,21 +115,7 @@ class DisassemblyAssistant(pwndbg.gdblib.disasm.arch.DisassemblyAssistant):
             if not telescope_addresses:
                 return
 
-            # MOV [MEM], REG or IMM
-            if (
-                left.type == CS_OP_MEM and left.before_value is not None
-            ):  # right.type must then be either CS_OP_REG or CS_OP_IMM. Cannot MOV mem to mem
-                # If the memory isn't mapped, we will segfault
-                if not pwndbg.gdblib.memory.peek(left.before_value):
-                    instruction.annotation = MessageColor.error(
-                        f"<Cannot dereference [{MemoryColor.get(left.before_value)}]>"
-                    )
-                else:
-                    instruction.annotation = f"{left.str} => {super()._telescope_format_list(telescope_addresses, TELESCOPE_DEPTH, emu)}"
-
-            # MOV REG, REG or IMM
-            elif left.type == CS_OP_REG and right.type in (CS_OP_REG, CS_OP_IMM):
-                instruction.annotation = f"{left.str} => {super()._telescope_format_list(telescope_addresses, TELESCOPE_DEPTH, emu)}"
+            instruction.annotation = f"{left.str} => {super()._telescope_format_list(telescope_addresses, TELESCOPE_DEPTH, emu)}"
 
     def handle_vmovaps(self, instruction: PwndbgInstruction, emu: Emulator) -> None:
         # If the source or destination is in memory, it must be aligned to: