From 0948712555683f46092dc26db3923f35edb1bfe2 Mon Sep 17 00:00:00 2001
From: Gulshan Singh <gsingh2011@gmail.com>
Date: Sun, 25 Feb 2024 17:56:56 -0800
Subject: [PATCH] Allow multibit register bitflags (#2029)

---
 pwndbg/color/context.py                     | 30 +++++++++++++++------
 pwndbg/lib/regs.py                          | 24 ++++++++---------
 tests/qemu-tests/tests/user/test_aarch64.py |  2 +-
 3 files changed, 35 insertions(+), 21 deletions(-)

diff --git a/pwndbg/color/context.py b/pwndbg/color/context.py
index e6c08dbe3..d4c0c924c 100644
--- a/pwndbg/color/context.py
+++ b/pwndbg/color/context.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 from pwndbg.color import generateColorFunction
 from pwndbg.color import theme
 from pwndbg.gdblib import config
+from pwndbg.lib.regs import BitFlags
 
 config_prefix_color = theme.add_color_param(
     "code-prefix-color", "none", "color for 'context code' command (prefix marker)"
@@ -87,7 +88,7 @@ def comment(x: object) -> str:
     return generateColorFunction(config.comment_color)(x)
 
 
-def format_flags(value, flags, last=None):
+def format_flags(value, flags: BitFlags, last=None):
     if value is None:
         return "<unavailable>"
 
@@ -97,15 +98,28 @@ def format_flags(value, flags, last=None):
 
     names = []
     for name, bit in flags.items():
-        bit = 1 << bit
-        if value & bit:
-            name = name.upper()
-            name = flag_set(name)
+        # If the size is not specified, assume it's 1
+        if isinstance(bit, int):
+            size = 1
         else:
-            name = name.lower()
-            name = flag_unset(name)
+            assert len(bit) == 2
+            size = bit[1]
+            bit = bit[0]
 
-        if last is not None and value & bit != last & bit:
+        mask = (1 << size) - 1
+        flag_val = (value >> bit) & mask
+
+        # If the bitfield is larger than a single bit, we can't communicate the value
+        # with just the case of the name, so append the actual value
+        if size > 1:
+            name = f"{name}:{flag_val}"
+
+        if flag_val == 0:
+            name = flag_unset(name.lower())
+        else:
+            name = flag_set(name.upper())
+
+        if last is not None and flag_val != (last >> bit) & mask:
             name = flag_changed(name)
         names.append(name)
 
diff --git a/pwndbg/lib/regs.py b/pwndbg/lib/regs.py
index 93b77a9d8..fad3e001b 100644
--- a/pwndbg/lib/regs.py
+++ b/pwndbg/lib/regs.py
@@ -4,11 +4,14 @@ standardized interface to registers like "sp" and "pc".
 """
 from __future__ import annotations
 
-import collections
 from typing import Any
 from typing import Dict
 from typing import Iterator
+from typing import OrderedDict
 from typing import Tuple
+from typing import Union
+
+BitFlags = OrderedDict[str, Union[int, Tuple[int, int]]]
 
 
 class RegisterSet:
@@ -51,7 +54,7 @@ class RegisterSet:
         stack: str = "sp",
         frame: str | None = None,
         retaddr: Tuple[str, ...] = tuple(),
-        flags: Dict[str, Any] = {},
+        flags: Dict[str, BitFlags] = {},
         gpr: Tuple[str, ...] = tuple(),
         misc: Tuple[str, ...] = tuple(),
         args: Tuple[str, ...] = tuple(),
@@ -80,7 +83,7 @@ class RegisterSet:
         yield from self.all
 
 
-arm_cpsr_flags = collections.OrderedDict(
+arm_cpsr_flags = BitFlags(
     [
         ("N", 31),
         ("Z", 30),
@@ -95,11 +98,9 @@ arm_cpsr_flags = collections.OrderedDict(
         ("F", 6),
     ]
 )
-arm_xpsr_flags = collections.OrderedDict(
-    [("N", 31), ("Z", 30), ("C", 29), ("V", 28), ("Q", 27), ("T", 24)]
-)
+arm_xpsr_flags = BitFlags([("N", 31), ("Z", 30), ("C", 29), ("V", 28), ("Q", 27), ("T", 24)])
 
-aarch64_cpsr_flags = collections.OrderedDict(
+aarch64_cpsr_flags = BitFlags(
     [
         ("N", 31),
         ("Z", 30),
@@ -112,8 +113,7 @@ aarch64_cpsr_flags = collections.OrderedDict(
         ("A", 8),
         ("I", 7),
         ("F", 6),
-        # TODO: EL is two bits
-        ("EL", 2),
+        ("EL", (2, 2)),
         ("SP", 0),
     ]
 )
@@ -221,7 +221,7 @@ aarch64 = RegisterSet(
 )
 
 x86flags = {
-    "eflags": collections.OrderedDict(
+    "eflags": BitFlags(
         [("CF", 0), ("PF", 2), ("AF", 4), ("ZF", 6), ("SF", 7), ("IF", 9), ("DF", 10), ("OF", 11)]
     )
 }
@@ -330,7 +330,7 @@ i386 = RegisterSet(
 # r31     Used for local variables or "environment pointers"
 powerpc = RegisterSet(
     retaddr=("lr",),
-    flags={"msr": {}, "xer": {}},
+    flags={"msr": BitFlags(), "xer": BitFlags()},
     gpr=(
         "r0",
         "r1",
@@ -401,7 +401,7 @@ sparc = RegisterSet(
     stack="sp",
     frame="fp",
     retaddr=("i7",),
-    flags={"psr": {}},
+    flags={"psr": BitFlags()},
     gpr=(
         "g1",
         "g2",
diff --git a/tests/qemu-tests/tests/user/test_aarch64.py b/tests/qemu-tests/tests/user/test_aarch64.py
index e5d4911fd..29a334330 100644
--- a/tests/qemu-tests/tests/user/test_aarch64.py
+++ b/tests/qemu-tests/tests/user/test_aarch64.py
@@ -17,7 +17,7 @@ try:
     gdb.execute("auxv", to_string=True)
     assert (
         gdb.execute("cpsr", to_string=True, from_tty=False).strip()
-        == "cpsr 0x60000000 [ n Z C v q pan il d a i f el sp ]"
+        == "cpsr 0x60000000 [ n Z C v q pan il d a i f el:0 sp ]"
     )
     gdb.execute("context", to_string=True)
     gdb.execute("hexdump", to_string=True)