mallocng: Show slot statusline when printing meta (#3245)

* add line to dump_meta giving an overview of the allocation state of the group

* print slot state more sensically

* add and fixup status tests

pwndbg/aglib/heap/mallocng.py

@@ -432,6 +432,10 @@ class Slot:
         Raises:
             pwndbg.dbg_mod.Error: When reading meta fails.
         """
+        # Special case (probably) freed chunks:
+        if self.reserved == -1:
+            return 0
+
         # https://elixir.bootlin.com/musl/v1.2.5/source/src/malloc/mallocng/meta.h#L159
         return self.end - self.reserved - self.p
 

pwndbg/commands/mallocng.py

@@ -64,6 +64,16 @@ def get_colored_slot_state(ss: mallocng.SlotState) -> str:
     return C.colorize(ss.value, get_slot_color(ss))
 
 
+def get_colored_slot_state_short(ss: mallocng.SlotState) -> str:
+    match ss:
+        case mallocng.SlotState.ALLOCATED:
+            return C.colorize("U", state_alloc_color)
+        case mallocng.SlotState.FREED:
+            return C.colorize("F", state_freed_color)
+        case mallocng.SlotState.AVAIL:
+            return C.colorize("A", state_avail_color)
+
+
 def dump_group(group: mallocng.Group) -> str:
     try:
         # May fail on corrupt meta.
@@ -99,7 +109,12 @@ def dump_group(group: mallocng.Group) -> str:
     return pp.dump()
 
 
-def dump_meta(meta: mallocng.Meta) -> str:
+def dump_meta(meta: mallocng.Meta, focus_slot: Optional[int] = None) -> str:
+    """
+    Arguments:
+        meta: the meta to dump
+        focus_slot: the index of the slot to highlight in the slot statuses list
+    """
     int_size = str(typeinfo.sint.sizeof * 8)
     avail_binary = "0b" + format(meta.avail_mask, f"0{int_size}b")
     freed_binary = "0b" + format(meta.freed_mask, f"0{int_size}b")
@@ -153,6 +168,26 @@ def dump_meta(meta: mallocng.Meta) -> str:
             print(message.error(f"Could not fetch parent group: {e}"))
         output += C.bold(".\n")
 
+    # Print the slot statuses.
+    slot_statuses = "\nSlot statuses: "
+    for i in range(meta.cnt):
+        this_slot = get_colored_slot_state_short(meta.slotstate_at_index(i))
+
+        if focus_slot is not None and i == focus_slot:
+            this_slot = "[" + this_slot + "]"
+
+        slot_statuses += this_slot
+
+    slot_statuses = C.bold(slot_statuses + "\n")
+    # Explain the notation.
+    slot_statuses += (
+        f"  ({C.bold(get_colored_slot_state_short(mallocng.SlotState.ALLOCATED))}: Inuse (allocated)"
+        f" / {C.bold(get_colored_slot_state_short(mallocng.SlotState.FREED))}: Freed"
+        f" / {C.bold(get_colored_slot_state_short(mallocng.SlotState.AVAIL))}: Available)\n"
+    )
+
+    output += slot_statuses
+
     return output
 
 
@@ -185,7 +220,7 @@ def dump_grouped_slot(gslot: mallocng.GroupedSlot, all: bool) -> str:
 
     if all:
         output += dump_group(gslot.group)
-        output += dump_meta(gslot.meta)
+        output += dump_meta(gslot.meta, gslot.idx)
 
     return output
 
@@ -193,6 +228,9 @@ def dump_grouped_slot(gslot: mallocng.GroupedSlot, all: bool) -> str:
 def dump_slot(
     slot: mallocng.Slot, all: bool, successful_preload: bool, will_dump_gslot: bool
 ) -> str:
+    if successful_preload:
+        assert not will_dump_gslot and "Why?"
+
     pp = PropertyPrinter()
 
     all = all and successful_preload and not will_dump_gslot
@@ -231,6 +269,10 @@ def dump_slot(
                     extra="slot's unused memory / 0x10",
                     alt_value=(slot.slack * mallocng.UNIT),
                 ),
+                Property(
+                    name="state",
+                    value=get_colored_slot_state(slot.meta.slotstate_at_index(slot.idx)),
+                ),
             ]
         )
         pp.end_section()
@@ -279,22 +321,27 @@ def dump_slot(
                 alt_value=cyc_val_alt,
             ),
         )
+    else:
+        # We haven't printed the slot state yet. Will we do it with a grouped slot?
+        if not will_dump_gslot:
+            # Nope, then let's go ahead and guess.
+            inband_group.append(
+                Property(
+                    name="state",
+                    value=get_colored_slot_state(slot.slot_state),
+                    extra="(probably, check the meta)",
+                )
+            )
 
     pp.add(inband_group)
     pp.end_section()
 
     output = pp.dump()
 
-    if not will_dump_gslot:
-        # The grouped_slot will have accurate information on this,
-        # no need for us to guess.
-        output += C.bold(
-            "\nThe slot is (probably) " + get_colored_slot_state(slot.slot_state) + ".\n\n"
-        )
-
     if all:
+        output += "\n"
         output += dump_group(slot.group)
-        output += dump_meta(slot.meta)
+        output += dump_meta(slot.meta, slot.idx)
 
     return output
 

tests/library/dbg/tests/test_mallocng.py

@@ -47,6 +47,7 @@ async def test_mallocng_slot_user(ctrl: Controller, binary: str):
         "  stride:         0x30              distance between adjacent slots",
         """  user size:      0x20              aka "nominal size", `n`""",
         r"  slack:          0x0 \(0x0\)         slot's unused memory \/ 0x10",
+        "  state:          allocated         ",
         "in-band",
         r"  offset:         0x[0-9] \(0x[0-9]{0,1}0\)         distance to first slot start \/ 0x10",
         r"  index:          0x0               index of slot in its group",
@@ -54,10 +55,6 @@ async def test_mallocng_slot_user(ctrl: Controller, binary: str):
         "                                    use ftr reserved",
         "  ftr reserved:   0xc               ",
         r"  cyclic offset:  NA \(not cyclic\)   prevents double free, \(p - start\) / 0x10",
-        "",
-        r"The slot is \(probably\) allocated.",
-        "",
-        "",
     ]
 
     assert len(expected_output) == len(buffer1_out)
@@ -72,12 +69,12 @@ async def test_mallocng_slot_user(ctrl: Controller, binary: str):
     stride_idx = 7
     user_size_idx = 8
     slack_idx = 9
-    offset_idx = 11
-    index_idx = 12
-    hdr_res_idx = 13
-    ftr_res_idx = 15
-    cyclic_idx = 16
-    status_idx = 18
+    state_idx = 10
+    offset_idx = 12
+    index_idx = 13
+    hdr_res_idx = 14
+    ftr_res_idx = 16
+    cyclic_idx = 17
 
     # Check stride
     assert "stride" in buffer2_out[stride_idx] and " 0x30 " in buffer2_out[stride_idx]
@@ -91,6 +88,10 @@ async def test_mallocng_slot_user(ctrl: Controller, binary: str):
     assert "slack" in buffer2_out[slack_idx] and " 0x0 " in buffer2_out[slack_idx]
     assert "slack" in buffer4_out[slack_idx] and " 0x8 (0x80) " in buffer4_out[slack_idx]
 
+    # Check allocation status
+    assert "state" in buffer2_out[state_idx] and " allocated " in buffer2_out[state_idx]
+    assert "state" in buffer4_out[state_idx] and " allocated " in buffer4_out[state_idx]
+
     # Check offset
     assert "offset" in buffer2_out[offset_idx] and " 0x3 (0x30) " in buffer2_out[offset_idx]
     if binary == HEAP_MALLOCNG_STATIC:
@@ -129,10 +130,6 @@ async def test_mallocng_slot_user(ctrl: Controller, binary: str):
             and " NA (not cyclic) " in buffer4_out[cyclic_idx]
         )
 
-    # Check allocation status
-    assert "slot is" in buffer2_out[status_idx] and " allocated." in buffer2_out[status_idx]
-    assert "slot is" in buffer4_out[status_idx] and " allocated." in buffer4_out[status_idx]
-
     # == Check command on free slots ==
     break_at_sym("break_here")
     await ctrl.cont()
@@ -238,6 +235,9 @@ async def test_mallocng_group(ctrl: Controller, binary: str):
         "  maplen:         0x0               ",
         "",
         rf"Group nested in slot of another group \({re_addr}\).",
+        "",
+        "Slot statuses: UUUAAAAAAA",
+        r"  \(U: Inuse \(allocated\) / F: Freed / A: Available\)",
     ]
 
     assert len(expected_out) == len(group1_out)
@@ -246,8 +246,9 @@ async def test_mallocng_group(ctrl: Controller, binary: str):
         assert re.match(expected_out[i], group1_out[i])
 
     # == Check group traversal is done properly.
+    pgline_idx = -4
 
-    assert "another group" in group1_out[-1]
+    assert "another group" in group1_out[pgline_idx]
 
     # We are going to fetch parent groups recursively until
     # we reach the outermost group which is either mmap()-ed in or
@@ -255,18 +256,18 @@ async def test_mallocng_group(ctrl: Controller, binary: str):
     cur_group_out: List[str] = group1_out
     cur_group_addr: int = group_addr
 
-    while "another group" in cur_group_out[-1]:
+    while "another group" in cur_group_out[pgline_idx]:
         cur_group_addr = int(
-            re.search(r"group \((0x[0-9a-fA-F]+)\)", cur_group_out[-1]).group(1), 16
+            re.search(r"group \((0x[0-9a-fA-F]+)\)", cur_group_out[pgline_idx]).group(1), 16
         )
         cur_group_out = color.strip(
             await ctrl.execute_and_capture(f"ng-group {cur_group_addr}")
         ).splitlines()
 
     if binary == HEAP_MALLOCNG_STATIC:
-        assert "mmap()" in cur_group_out[-1]
+        assert "mmap()" in cur_group_out[pgline_idx]
     else:
-        assert "donated by ld" in cur_group_out[-1]
+        assert "donated by ld" in cur_group_out[pgline_idx]
 
 
 @pwndbg_test