Fetch chunk metadata (pwndbg#2152)

* Add resolve_renamed_struct_field()

* Add ChunkField enum

* Raise ValueError when field name not found

* Add fetch_chunk_metadata()

* Use None for include_only_fields argument

* Use None as default value for include/exclude filters in fetch_struct_as_dictionary

* Resolve C408

* Remove unused typing.Set import

* Correct use of Set in type hints

Thanks to @gsingh93 for figuring out what I'd done wrong here: capitalized "Set" must be used for compatibility with Python 3.8

* Correct use of Set in type hints

* Import Set from Typing in ptmalloc.py

pwndbg/gdblib/memory.py

@@ -372,8 +372,8 @@ def update_min_addr() -> None:
 def fetch_struct_as_dictionary(
     struct_name: str,
     struct_address: int,
-    include_only_fields: Set[str] = set(),
-    exclude_fields: Set[str] = set(),
+    include_only_fields: Set[str] | None = None,
+    exclude_fields: Set[str] | None = None,
 ) -> GdbDict:
     struct_type = gdb.lookup_type("struct " + struct_name)
     fetched_struct = poi(struct_type, struct_address)
@@ -383,12 +383,15 @@ def fetch_struct_as_dictionary(
 
 def pack_struct_into_dictionary(
     fetched_struct: gdb.Value,
-    include_only_fields: Set[str] = set(),
-    exclude_fields: Set[str] = set(),
+    include_only_fields: Set[str] | None = None,
+    exclude_fields: Set[str] | None = None,
 ) -> GdbDict:
     struct_as_dictionary = {}
 
-    if len(include_only_fields) != 0:
+    if exclude_fields is None:
+        exclude_fields = set()
+
+    if include_only_fields is not None:
         for field_name in include_only_fields:
             key = field_name
             value = convert_gdb_value_to_python_value(fetched_struct[field_name])
@@ -419,3 +422,13 @@ def convert_gdb_value_to_python_value(gdb_value: gdb.Value) -> int | GdbDict:
         return pack_struct_into_dictionary(gdb_value)
 
     raise NotImplementedError
+
+
+def resolve_renamed_struct_field(struct_name: str, possible_field_names: Set[str]) -> str:
+    struct_type = gdb.lookup_type("struct " + struct_name)
+
+    for field_name in possible_field_names:
+        if gdb.types.has_field(struct_type, field_name):
+            return field_name
+
+    raise ValueError(f"Field name did not match any of {possible_field_names}.")

pwndbg/heap/ptmalloc.py

@@ -19,6 +19,7 @@
 from typing import Generic
 from typing import List
 from typing import OrderedDict as OrderedDictType
+from typing import Set
 from typing import Tuple
 from typing import Type
 from typing import TypeVar
@@ -166,6 +167,64 @@ def heap_for_ptr(ptr: int) -> int:
     return ptr & ~(HEAP_MAX_SIZE - 1)
 
 
+class ChunkField(Enum):
+    PREV_SIZE = 1
+    SIZE = 2
+    FD = 3
+    BK = 4
+    FD_NEXTSIZE = 5
+    BK_NEXTSIZE = 6
+
+
+def fetch_chunk_metadata(address: int, include_only_fields: Set[ChunkField] | None = None):
+    prev_size_field_name = pwndbg.gdblib.memory.resolve_renamed_struct_field(
+        "malloc_chunk", {"prev_size", "mchunk_prev_size"}
+    )
+    size_field_name = pwndbg.gdblib.memory.resolve_renamed_struct_field(
+        "malloc_chunk", {"size", "mchunk_size"}
+    )
+
+    if include_only_fields is None:
+        fetched_struct = pwndbg.gdblib.memory.fetch_struct_as_dictionary("malloc_chunk", address)
+    else:
+        requested_fields: Set[str] = set()
+
+        for field in include_only_fields:
+            if field is ChunkField.PREV_SIZE:
+                requested_fields.add(prev_size_field_name)
+            elif field is ChunkField.SIZE:
+                requested_fields.add(size_field_name)
+            elif field is ChunkField.FD:
+                requested_fields.add("fd")
+            elif field is ChunkField.BK:
+                requested_fields.add("bk")
+            elif field is ChunkField.FD_NEXTSIZE:
+                requested_fields.add("fd_nextsize")
+            elif field is ChunkField.BK_NEXTSIZE:
+                requested_fields.add("bk_nextsize")
+
+        fetched_struct = pwndbg.gdblib.memory.fetch_struct_as_dictionary(
+            "malloc_chunk", address, include_only_fields=requested_fields
+        )
+
+    normalized_struct = {}
+    for field in fetched_struct:
+        if field == prev_size_field_name:
+            normalized_struct[ChunkField.PREV_SIZE] = fetched_struct[prev_size_field_name]
+        elif field == size_field_name:
+            normalized_struct[ChunkField.SIZE] = fetched_struct[size_field_name]
+        elif field == "fd":
+            normalized_struct[ChunkField.FD] = fetched_struct["fd"]
+        elif field == "bk":
+            normalized_struct[ChunkField.BK] = fetched_struct["bk"]
+        elif field == "fd_nextsize":
+            normalized_struct[ChunkField.FD_NEXTSIZE] = fetched_struct["fd_nextsize"]
+        elif field == "bk_nextsize":
+            normalized_struct[ChunkField.BK_NEXTSIZE] = fetched_struct["bk_nextsize"]
+
+    return normalized_struct
+
+
 class Chunk:
     __slots__ = (
         "_gdbValue",