From bbfd1090370345f8906220d2090e0e28b6c1c343 Mon Sep 17 00:00:00 2001
From: kotee4ko <62266948+kotee4ko@users.noreply.github.com>
Date: Wed, 29 Nov 2023 04:10:03 +0200
Subject: [PATCH] [heap] `hi` command -- feature to check if an address belongs
 to a chunk. (#1938)

* add hi -- heap_info

* add default return after first hit

* fix

* ready

* + Add __contains__ method of Chunk class
+ Update verbose output features

* heap_info: new chunk detection feature

* heap_info: new chunk detection feature (lint >_<)

* heap_info: new chunk detection feature (lint >_<) 2

* eap_info: new chunk detection feature (lint >_<) 3

---------

Co-authored-by: Administrator <admin@example.com>
Co-authored-by: Th3C4t <swagcat228@gmail.com>
---
 pwndbg/commands/heap.py | 54 +++++++++++++++++++++++++++++++++++++++++
 pwndbg/heap/ptmalloc.py | 17 +++++++++++++
 2 files changed, 71 insertions(+)

diff --git a/pwndbg/commands/heap.py b/pwndbg/commands/heap.py
index 0ed87da44..244853493 100644
--- a/pwndbg/commands/heap.py
+++ b/pwndbg/commands/heap.py
@@ -183,6 +183,60 @@ def heap(addr=None, verbose=False, simple=False) -> None:
             malloc_chunk(chunk.address, verbose=verbose, simple=simple)
 
 
+parser = argparse.ArgumentParser(
+    formatter_class=argparse.RawTextHelpFormatter,
+    description="""Searches all heaps to find if an address belongs to a chunk. If yes, prints the chunk.""",
+)
+parser.add_argument(
+    "addr",
+    type=int,
+    help="Address of the interest.",
+)
+parser.add_argument(
+    "-v", "--verbose", action="store_true", help="Print all chunk fields, even unused ones."
+)
+parser.add_argument(
+    "-s", "--simple", action="store_true", help="Simply print malloc_chunk struct's contents."
+)
+parser.add_argument(
+    "-f",
+    "--fake",
+    action="store_true",
+    help="Allow fake chunks. If set, displays any memory as a heap chunk (even if its not a real chunk).",
+)
+
+
+@pwndbg.commands.ArgparsedCommand(parser, category=CommandCategory.HEAP)
+@pwndbg.commands.OnlyWhenRunning
+@pwndbg.commands.OnlyWithResolvedHeapSyms
+@pwndbg.commands.OnlyWhenHeapIsInitialized
+def hi(addr, verbose=False, simple=False, fake=False) -> None:
+    try:
+        heap = Heap(addr)
+    except Exception as E:
+        print(f"The provided address {hex(addr)} cannot be interpreted as a heap!\n{E}\n")
+        return
+
+    if fake is False and heap.arena is None:
+        return
+
+    for chunk in heap:
+        if addr in chunk:
+            malloc_chunk(chunk.address, verbose=verbose, simple=simple)
+            if verbose:
+                start = chunk.address + (pwndbg.gdblib.arch.ptrsize if chunk.prev_inuse else 0x00)
+                print(f"Your address: {hex(addr)}")
+                print(f"Head offset: {hex(addr - start)}")
+                if chunk.is_top_chunk is False:
+                    end = (
+                        start
+                        + chunk.real_size
+                        + (pwndbg.gdblib.arch.ptrsize if chunk.prev_inuse is False else 0x00)
+                    )
+                    print(f"Tail offset: {hex(end - addr)}")
+            break
+
+
 parser = argparse.ArgumentParser(
     formatter_class=argparse.RawTextHelpFormatter,
     description="""Print the contents of an arena.
diff --git a/pwndbg/heap/ptmalloc.py b/pwndbg/heap/ptmalloc.py
index ae270f07e..cc7228799 100644
--- a/pwndbg/heap/ptmalloc.py
+++ b/pwndbg/heap/ptmalloc.py
@@ -336,6 +336,23 @@ class Chunk:
         else:
             return None
 
+    def __contains__(self, addr: int) -> bool:
+        """
+        This allow us to avoid extra constructions like 'if strart_addr <= ptr <= end_addr', etc.
+        """
+        size_field_address = self._gdbValue[self.__match_renamed_field("size")].address
+        start_address = size_field_address if self.prev_inuse else self.address
+
+        next = self.next_chunk()
+        # and this is handles chunk's last qword field, depending on prev_inuse bit
+        if next is None:
+            end_address = size_field_address + self.real_size
+        else:
+            next_size_field_address = next._gdbValue[self.__match_renamed_field("size")].address
+            end_address = next_size_field_address if next.prev_inuse else next.address
+
+        return start_address <= addr < end_address
+
 
 class Heap:
     __slots__ = (