Add bitwise math helper functions (#2278)

* Implement bitwise math rotation operations on numbers of discrete width. Will be used in manually evaluating arm instruction offsets and shifts

* fixes

pwndbg/gdblib/arch.py

@@ -97,6 +97,7 @@ def _get_arch(ptrsize: int):
             elif match.startswith("riscv:"):
                 match = match[6:]
             elif match == "riscv":
+                # If GDB doesn't detect the width, it will just say `riscv`.
                 match = "rv64"
             return match, ptrsize, endian
 

pwndbg/gdblib/disasm/mips.py

@@ -22,6 +22,7 @@ from typing_extensions import override
 
 import pwndbg.gdblib.disasm.arch
 import pwndbg.gdblib.regs
+import pwndbg.lib.disasm.helpers as bit_math
 from pwndbg.emu.emulator import Emulator
 from pwndbg.gdblib.disasm.instruction import FORWARD_JUMP_GROUP
 from pwndbg.gdblib.disasm.instruction import InstructionCondition
@@ -43,24 +44,17 @@ BRANCH_LIKELY_INSTRUCTIONS = {
 }
 
 
-def to_signed(unsigned: int):
-    if pwndbg.gdblib.arch.ptrsize == 8:
-        return unsigned - ((unsigned & 0x80000000_00000000) << 1)
-    else:
-        return unsigned - ((unsigned & 0x80000000) << 1)
-
-
 CONDITION_RESOLVERS: Dict[int, Callable[[List[int]], bool]] = {
     MIPS_INS_BEQZ: lambda ops: ops[0] == 0,
     MIPS_INS_BNEZ: lambda ops: ops[0] != 0,
     MIPS_INS_BEQ: lambda ops: ops[0] == ops[1],
     MIPS_INS_BNE: lambda ops: ops[0] != ops[1],
-    MIPS_INS_BGEZ: lambda ops: to_signed(ops[0]) >= 0,
-    MIPS_INS_BGEZAL: lambda ops: to_signed(ops[0]) >= 0,
-    MIPS_INS_BGTZ: lambda ops: to_signed(ops[0]) > 0,
-    MIPS_INS_BLEZ: lambda ops: to_signed(ops[0]) <= 0,
-    MIPS_INS_BLTZAL: lambda ops: to_signed(ops[0]) < 0,
-    MIPS_INS_BLTZ: lambda ops: to_signed(ops[0]) < 0,
+    MIPS_INS_BGEZ: lambda ops: bit_math.to_signed(ops[0], pwndbg.gdblib.arch.ptrsize * 8) >= 0,
+    MIPS_INS_BGEZAL: lambda ops: bit_math.to_signed(ops[0], pwndbg.gdblib.arch.ptrsize * 8) >= 0,
+    MIPS_INS_BGTZ: lambda ops: bit_math.to_signed(ops[0], pwndbg.gdblib.arch.ptrsize * 8) > 0,
+    MIPS_INS_BLEZ: lambda ops: bit_math.to_signed(ops[0], pwndbg.gdblib.arch.ptrsize * 8) <= 0,
+    MIPS_INS_BLTZAL: lambda ops: bit_math.to_signed(ops[0], pwndbg.gdblib.arch.ptrsize * 8) < 0,
+    MIPS_INS_BLTZ: lambda ops: bit_math.to_signed(ops[0], pwndbg.gdblib.arch.ptrsize * 8) < 0,
 }
 
 

pwndbg/gdblib/disasm/riscv.py

@@ -7,6 +7,7 @@ from typing_extensions import override
 import pwndbg.gdblib.arch
 import pwndbg.gdblib.disasm.arch
 import pwndbg.gdblib.regs
+import pwndbg.lib.disasm.helpers as bit_math
 from pwndbg.emu.emulator import Emulator
 from pwndbg.gdblib.disasm.instruction import InstructionCondition
 from pwndbg.gdblib.disasm.instruction import PwndbgInstruction
@@ -28,14 +29,8 @@ class DisassemblyAssistant(pwndbg.gdblib.disasm.arch.DisassemblyAssistant):
         else:
             src2_unsigned = 0
 
-        if self.architecture == "rv32":
-            src1_signed = src1_unsigned - ((src1_unsigned & 0x80000000) << 1)
-            src2_signed = src2_unsigned - ((src2_unsigned & 0x80000000) << 1)
-        elif self.architecture == "rv64":
-            src1_signed = src1_unsigned - ((src1_unsigned & 0x80000000_00000000) << 1)
-            src2_signed = src2_unsigned - ((src2_unsigned & 0x80000000_00000000) << 1)
-        else:
-            raise NotImplementedError(f"architecture '{self.architecture}' not implemented")
+        src1_signed = bit_math.to_signed(src1_unsigned, pwndbg.gdblib.arch.ptrsize * 8)
+        src2_signed = bit_math.to_signed(src2_unsigned, pwndbg.gdblib.arch.ptrsize * 8)
 
         condition = {
             RISCV_INS_BEQ: src1_signed == src2_signed,

pwndbg/lib/disasm/helpers.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+
+def to_signed(unsigned: int, bit_width: int):
+    """
+    Returns the signed number associated with the two's-complement binary representation of `unsigned`
+    """
+    extract_bit = 1 << (bit_width - 1)
+    return unsigned - ((unsigned & extract_bit) << 1)
+
+
+def logical_shift_left(n: int, shift_amt: int, bit_width: int):
+    return (n << shift_amt) & ((1 << bit_width) - 1)
+
+
+def logical_shift_right(n: int, shift_amt: int, bit_width: int):
+    """
+    `n` is truncated to the width of `bit_width` before the operation takes place.
+    """
+    n = n & ((1 << bit_width) - 1)
+    return n >> shift_amt
+
+
+def rotate_right(n: int, shift_amt: int, bit_width: int):
+    """
+    `n` is truncated to the width of `bit_width` before the operation takes place.
+    """
+    n = n & ((1 << bit_width) - 1)
+    return ((n >> shift_amt) | (n << (bit_width - shift_amt))) & ((1 << bit_width) - 1)
+
+
+def arithmetic_shift_right(n: int, shift_amt: int, bit_width: int):
+    """
+    This returns the value represented by the two's-complement binary representation of the final result.
+    This means the result could be negative (if the top bit of the input is negative)
+
+    `n` is truncated to the width of `bit_width` before the operation takes place.
+    """
+    n = n & ((1 << bit_width) - 1)
+
+    result = logical_shift_right(n, shift_amt, bit_width)
+
+    sign_extension_mask = (1 << (bit_width - shift_amt)) - 1
+    # Replicate the sign bit if it's set
+    if n & (1 << (bit_width - 1)):
+        result |= ~sign_extension_mask
+
+    return result

tests/unit-tests/test_bit_math.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import pwndbg.lib.disasm.helpers as bit_math
+
+# We must import the function under test after all the mocks are imported
+
+
+def test_to_signed():
+    assert bit_math.to_signed(0b0100_0000, 8) == 0b0100_0000
+    assert bit_math.to_signed(0b1000_0000, 8) == -128
+
+    assert bit_math.to_signed(0xFFFFFFFF_FFFFFFFF, 64) == -1
+    assert bit_math.to_signed(0x7FFFFFFF_FFFFFFFF, 64) == 0x7FFFFFFF_FFFFFFFF
+
+    assert bit_math.to_signed(0xFFFF_FFFF, 32) == -1
+    assert bit_math.to_signed(0x8000_0000, 32) == -(2**31)
+
+
+def test_lsl():
+    assert bit_math.logical_shift_left(0b1000_0000, 1, 8) == 0
+    assert bit_math.logical_shift_left(0b0100_0000, 1, 8) == 0b1000_0000
+    assert bit_math.logical_shift_left(0b1111_1111, 1, 8) == 0b1111_1110
+    assert bit_math.logical_shift_left(0b1111_1111, 5, 8) == 0b1110_0000
+
+
+def test_lsr():
+    assert bit_math.logical_shift_right(0b1000_0000, 1, 8) == 0b0100_0000
+    assert bit_math.logical_shift_right(0b0100_0000, 1, 8) == 0b0010_0000
+    assert bit_math.logical_shift_right(0b1111_1111, 1, 8) == 0b0111_1111
+    assert bit_math.logical_shift_right(0b1111_1111, 5, 8) == 0b0000_0111
+    # Should truncate to bit_width before operation
+    assert bit_math.logical_shift_right(0b1_0000_0000, 1, 8) == 0
+
+
+def test_ror():
+    assert bit_math.rotate_right(0b1000_0001, 1, 8) == 0b1100_0000
+    assert bit_math.rotate_right(0b0100_0000, 1, 8) == 0b0010_0000
+    assert bit_math.rotate_right(0b0100_0000, 4, 8) == 0b0000_0100
+    assert bit_math.rotate_right(0b1111_1111, 1, 8) == 0b1111_1111
+    assert bit_math.rotate_right(0b1110_1111, 5, 8) == 0b0111_1111
+
+    # Should truncate to bit_width before operation
+    assert bit_math.rotate_right(0b1_0000_0000, 1, 8) == 0
+    assert bit_math.rotate_right(0b1_0111_1111, 1, 8) == 0b1011_1111
+
+
+def test_asr():
+    # Unsigned numbers should be the same
+    assert bit_math.arithmetic_shift_right(0b0100_0000, 1, 8) == bit_math.logical_shift_right(
+        0b0100_0000, 1, 8
+    )
+    assert bit_math.arithmetic_shift_right(0xFFFF_FF, 1, 32) == bit_math.logical_shift_right(
+        0xFFFF_FF, 1, 32
+    )
+    assert bit_math.arithmetic_shift_right(0xFFFF_FF, 6, 32) == bit_math.logical_shift_right(
+        0xFFFF_FF, 6, 32
+    )
+
+    assert bit_math.arithmetic_shift_right(0b1000_0000, 1, 8) == -64
+    assert bit_math.arithmetic_shift_right(0b1000_0000, 2, 8) == -32
+    assert bit_math.arithmetic_shift_right(0b1000_0000, 7, 8) == -1
+
+    # Should truncate to bit_width before operation
+    assert bit_math.arithmetic_shift_right(0b1_0000_0000, 1, 8) == 0
+    assert bit_math.arithmetic_shift_right(0b1_0111_1111, 7, 8) == 0
+
+    # Unsigned number shifted
+    assert bit_math.arithmetic_shift_right(0x70000000_00000000, 62, 64) == 1
+    assert bit_math.arithmetic_shift_right(0x70000000_00000000, 63, 64) == 0