Add some basic heap heuristics supports for riscv and powerpc (#1649)

* Add some basic heap heuristics support for riscv and powerpc

Use the relocation section to find the main_arena address if possible

* Refactor the code we used to get the field offset

pwndbg/gdblib/elf.py

@@ -17,6 +17,8 @@ from typing import Tuple
 import gdb
 from elftools.elf.constants import SH_FLAGS
 from elftools.elf.elffile import ELFFile
+from elftools.elf.relocation import Relocation
+from elftools.elf.relocation import RelocationSection
 
 import pwndbg.auxv
 import pwndbg.gdblib.abi
@@ -188,6 +190,23 @@ def dump_section_by_name(
         return (section["sh_addr"], section["sh_size"], section.data()) if section else None
 
 
+def dump_relocations_by_section_name(
+    filepath: str, section_name: str, try_local_path: bool = False
+) -> Optional[Tuple[Relocation, ...]]:
+    """
+    Dump the relocation entries of a section from an ELF file, return a generator of Relocation objects.
+    """
+    # TODO: We should have some cache mechanism or something at `pndbg.gdblib.file.get_file()` in the future to avoid downloading the same file multiple times when we are debugging a remote process
+    local_path = pwndbg.gdblib.file.get_file(filepath, try_local_path=try_local_path)
+
+    with open(local_path, "rb") as f:
+        elffile = ELFFile(f)
+        section = elffile.get_section_by_name(section_name)
+        if section is None or not isinstance(section, RelocationSection):
+            return None
+        return tuple(section.iter_relocations())
+
+
 @pwndbg.gdblib.proc.OnlyWhenRunning
 @pwndbg.lib.memoize.reset_on_start
 def exe():

pwndbg/gdblib/proc.py

@@ -13,6 +13,7 @@ from typing import Optional
 from typing import Tuple
 
 import gdb
+from elftools.elf.relocation import Relocation
 
 import pwndbg.gdblib.qemu
 import pwndbg.lib.memoize
@@ -113,6 +114,18 @@ class module(ModuleType):
         """
         return pwndbg.gdblib.elf.dump_section_by_name(self.exe, ".data", try_local_path=True)
 
+    @pwndbg.lib.memoize.reset_on_start
+    @pwndbg.lib.memoize.reset_on_objfile
+    def dump_relocations_by_section_name(
+        self, section_name: str
+    ) -> Optional[Tuple[Relocation, ...]]:
+        """
+        Dump relocations of a section by section name of current process's ELF file
+        """
+        return pwndbg.gdblib.elf.dump_relocations_by_section_name(
+            self.exe, section_name, try_local_path=True
+        )
+
     @pwndbg.lib.memoize.reset_on_start
     @pwndbg.lib.memoize.reset_on_objfile
     def get_data_section_address(self) -> int:

pwndbg/glibc.py

@@ -9,6 +9,7 @@ from typing import Optional
 from typing import Tuple
 
 import gdb
+from elftools.elf.relocation import Relocation
 
 import pwndbg.gdblib.config
 import pwndbg.gdblib.elf
@@ -104,6 +105,8 @@ def get_libc_filename_from_info_sharedlibrary() -> Optional[str]:
 
 
 @pwndbg.gdblib.proc.OnlyWhenRunning
+@pwndbg.lib.memoize.reset_on_start
+@pwndbg.lib.memoize.reset_on_objfile
 def dump_elf_data_section() -> Optional[Tuple[int, int, bytes]]:
     """
     Dump .data section of libc ELF file
@@ -115,6 +118,22 @@ def dump_elf_data_section() -> Optional[Tuple[int, int, bytes]]:
     return pwndbg.gdblib.elf.dump_section_by_name(libc_filename, ".data", try_local_path=True)
 
 
+@pwndbg.gdblib.proc.OnlyWhenRunning
+@pwndbg.lib.memoize.reset_on_start
+@pwndbg.lib.memoize.reset_on_objfile
+def dump_relocations_by_section_name(section_name: str) -> Optional[Tuple[Relocation, ...]]:
+    """
+    Dump relocations of a section by section name of libc ELF file
+    """
+    libc_filename = get_libc_filename_from_info_sharedlibrary()
+    if not libc_filename:
+        # libc not loaded yet, or it's static linked
+        return None
+    return pwndbg.gdblib.elf.dump_relocations_by_section_name(
+        libc_filename, section_name, try_local_path=True
+    )
+
+
 @pwndbg.gdblib.proc.OnlyWhenRunning
 @pwndbg.lib.memoize.reset_on_start
 @pwndbg.lib.memoize.reset_on_objfile

pwndbg/heap/ptmalloc.py

@@ -1514,31 +1514,99 @@ class HeuristicHeap(GlibcMemoryAllocator):
 
         if not self._main_arena_addr:
             if self.is_statically_linked():
-                section = pwndbg.gdblib.proc.dump_elf_data_section()
-                section_address = pwndbg.gdblib.proc.get_data_section_address()
+                data_section = pwndbg.gdblib.proc.dump_elf_data_section()
+                data_section_address = pwndbg.gdblib.proc.get_data_section_address()
             else:
-                section = pwndbg.glibc.dump_elf_data_section()
-                section_address = pwndbg.glibc.get_data_section_address()
-            if section and section_address:
-                data_section_offset, size, data = section
-
-                # try to find the default main_arena struct in the .data section
-                for i in range(size - self.malloc_state.sizeof):
-                    # https://github.com/bminor/glibc/blob/glibc-2.37/malloc/malloc.c#L1902-L1907
-                    # static struct malloc_state main_arena =
-                    # {
-                    #   .mutex = _LIBC_LOCK_INITIALIZER,
-                    #   .next = &main_arena,
-                    #   .attached_threads = 1
-                    # };
-                    expected = self.malloc_state._c_struct()
-                    expected.next = data_section_offset + i
-                    expected.attached_threads = 1
-                    expected = bytes(expected)
-                    if expected == data[i : i + len(expected)]:
-                        self._main_arena_addr = section_address + i
+                data_section = pwndbg.glibc.dump_elf_data_section()
+                data_section_address = pwndbg.glibc.get_data_section_address()
+            if data_section and data_section_address:
+                data_section_offset, size, data_section_data = data_section
+                # Try to find the default main_arena struct in the .data section
+                # https://github.com/bminor/glibc/blob/glibc-2.37/malloc/malloc.c#L1902-L1907
+                # static struct malloc_state main_arena =
+                # {
+                #   .mutex = _LIBC_LOCK_INITIALIZER,
+                #   .next = &main_arena,
+                #   .attached_threads = 1
+                # };
+                expected = self.malloc_state._c_struct()
+                expected.attached_threads = 1
+                next_field_offset = self.malloc_state.get_field_offset("next")
+                malloc_state_size = self.malloc_state.sizeof
+
+                # Since RELR relocations might also have .rela.dyn section, we check it first
+                for section_name in (".relr.dyn", ".rela.dyn", ".rel.dyn"):
+                    if self._main_arena_addr:
+                        # If we have found the main_arena, we can stop searching
                         break
 
+                    if self.is_statically_linked():
+                        relocations = pwndbg.gdblib.proc.dump_relocations_by_section_name(
+                            section_name
+                        )
+                    else:
+                        relocations = pwndbg.glibc.dump_relocations_by_section_name(section_name)
+                    if not relocations:
+                        continue
+
+                    for relocation in relocations:
+                        r_offset = relocation.entry.r_offset
+
+                        # We only care about the relocation in .data section
+                        if r_offset - next_field_offset < data_section_offset:
+                            continue
+                        elif r_offset - next_field_offset >= data_section_offset + size:
+                            break
+
+                        # To find addend:
+                        # .relr.dyn and .rel.dyn need to read the data from r_offset
+                        # .rela.dyn has the addend in the entry
+                        if section_name != ".rela.dyn":
+                            addend = int.from_bytes(
+                                data_section_data[
+                                    r_offset
+                                    - data_section_offset : r_offset
+                                    - data_section_offset
+                                    + pwndbg.gdblib.arch.ptrsize
+                                ],
+                                pwndbg.gdblib.arch.endian,
+                            )
+                        else:
+                            addend = relocation.entry.r_addend
+
+                        # If addend is the offset of main_arena, then r_offset should be the offset of main_arena.next
+                        if r_offset - next_field_offset == addend:
+                            # Check if we can construct the default main_arena struct we expect
+                            tmp = data_section_data[
+                                addend
+                                - data_section_offset : addend
+                                - data_section_offset
+                                + malloc_state_size
+                            ]
+                            # Note: Although RELA relocations have r_addend, some compiler will still put the addend in the location of r_offset, so we still need to check both cases
+                            found = False
+                            expected.next = addend
+                            found |= bytes(expected) == tmp
+                            if not found:
+                                expected.next = 0
+                                found |= bytes(expected) == tmp
+                            if found:
+                                # This might be a false positive, but it is very unlikely, so should be fine :)
+                                self._main_arena_addr = (
+                                    data_section_address + addend - data_section_offset
+                                )
+                                break
+
+                # If we are still not able to find the main_arena, probably we are debugging a binary with statically linked libc and no PIE enabled
+                if not self._main_arena_addr:
+                    # Try to find the default main_arena struct in the .data section
+                    for i in range(0, size - self.malloc_state.sizeof, pwndbg.gdblib.arch.ptrsize):
+                        expected.next = data_section_offset + i
+                        if bytes(expected) == data_section_data[i : i + malloc_state_size]:
+                            # This also might be a false positive, but it is very unlikely too, so should also be fine :)
+                            self._main_arena_addr = data_section_address + i
+                            break
+
         if pwndbg.gdblib.memory.is_readable_address(self._main_arena_addr):
             self._main_arena = Arena(self._main_arena_addr)
             return self._main_arena

pwndbg/heap/structs.py

@@ -213,6 +213,13 @@ class CStruct2GDB:
         """
         return self.address + getattr(self._c_struct, field).offset
 
+    @classmethod
+    def get_field_offset(cls, field: str) -> int:
+        """
+        Returns the offset of the specified field.
+        """
+        return getattr(cls._c_struct, field).offset
+
     def items(self) -> tuple:
         """
         Returns a tuple of (field name, field value) pairs.