Fixing `kernel_vmmap` perf issue when symbol file is not added (#3390)

* fix x64 vmmap perf issue

* restore

* fixing tests

* further optimizations

* addressed comments

* handle none value

* addressed comments

pwndbg/aglib/kernel/__init__.py

@@ -116,12 +116,26 @@ def first_kernel_ro_page() -> pwndbg.lib.memory.Page | None:
     if base is None:
         return None
 
-    for mapping in pwndbg.aglib.kernel.paging.get_memory_map_raw():
+    banner = pwndbg.aglib.symbol.lookup_symbol_addr("linux_banner")
+    fallback_mappings = []
+    for mapping in pwndbg.aglib.kernel.vmmap.kernel_vmmap_pages():
         if mapping.vaddr < base:
             continue
+        if banner is not None and banner in mapping:
+            return mapping
+        if not mapping.read or mapping.write or mapping.execute:
+            fallback_mappings.append(mapping)
+            continue
 
         result = next(pwndbg.search.search(b"Linux version", mappings=[mapping]), None)
 
+        if result:
+            return mapping
+    for mapping in fallback_mappings:
+        # this loop handles when the kernel has not finished initialization
+        # and the permission of the first ro page has not been properly set
+        result = next(pwndbg.search.search(b"Linux version", mappings=[mapping]), None)
+
         if result:
             return mapping
 
@@ -137,6 +151,8 @@ def kconfig() -> pwndbg.lib.kernel.kconfig.Kconfig | None:
         config_end = pwndbg.aglib.symbol.lookup_symbol_addr("kernel_config_data_end")
     else:
         mapping = first_kernel_ro_page()
+        if mapping is None:
+            return None
         result = next(pwndbg.search.search(b"IKCFG_ST", mappings=[mapping]), None)
 
         if result is not None:
@@ -162,7 +178,7 @@ def kcmdline() -> str:
 
 
 @pwndbg.lib.cache.cache_until("start")
-def kversion() -> str:
+def kversion() -> str | None:
     try:
         if has_debug_symbols("linux_banner"):
             version_addr = pwndbg.aglib.symbol.lookup_symbol_addr("linux_banner")

pwndbg/aglib/kernel/paging.py

@@ -11,6 +11,7 @@ import pwndbg.aglib.vmmap_custom
 import pwndbg.color.message as M
 import pwndbg.lib.cache
 import pwndbg.lib.memory
+from pwndbg.aglib.kernel.vmmap import kernel_vmmap_pages
 from pwndbg.lib.regs import BitFlags
 
 # don't return None but rather an invalid value for address markers
@@ -18,14 +19,9 @@ from pwndbg.lib.regs import BitFlags
 INVALID_ADDR = 1 << 64
 
 
-@pwndbg.lib.cache.cache_until("stop")
-def get_memory_map_raw() -> Tuple[pwndbg.lib.memory.Page, ...]:
-    return pwndbg.aglib.kernel.vmmap.kernel_vmmap(False)
-
-
 @pwndbg.lib.cache.cache_until("stop")
 def first_kernel_page_start():
-    for page in get_memory_map_raw():
+    for page in kernel_vmmap_pages():
         if page.start and pwndbg.aglib.memory.is_kernel(page.start):
             return page.start
     return INVALID_ADDR
@@ -95,7 +91,7 @@ class ArchPagingInfo:
         raise NotImplementedError()
 
     def kbase_helper(self, address):
-        for mapping in get_memory_map_raw():
+        for mapping in kernel_vmmap_pages():
             # should be page aligned -- either from pt-dump or info mem
 
             # only search in kernel mappings:
@@ -200,7 +196,10 @@ class x86_64PagingInfo(ArchPagingInfo):
         try:
             target = self.physmap.to_bytes(8, byteorder="little")
             mapping = pwndbg.aglib.kernel.first_kernel_ro_page()
-            result = next(pwndbg.search.search(target, mappings=[mapping]), None)
+            result = next(
+                pwndbg.search.search(target, mappings=[mapping], aligned=pwndbg.aglib.arch.ptrsize),
+                None,
+            )
         except Exception as e:
             print(e)
             pass
@@ -215,12 +214,7 @@ class x86_64PagingInfo(ArchPagingInfo):
     def physmap(self):
         result = pwndbg.aglib.kernel.symbol.try_usymbol("page_offset_base")
         if result is None:
-            result = INVALID_ADDR
-            min = 0xFFFF888000000000 if self.paging_level == 4 else 0xFF11000000000000
-            for page in get_memory_map_raw():
-                if page.start and page.start >= min:
-                    result = page.start
-                    break
+            result = first_kernel_page_start()
         return result
 
     @property
@@ -390,22 +384,15 @@ class Aarch64PagingInfo(ArchPagingInfo):
     @property
     @pwndbg.lib.cache.cache_until("stop")
     def module_start(self):
-        # this is only used for marking the end of module_start
-        self.module_end = -1
         res = None
-        for page in get_memory_map_raw():
+        for page in kernel_vmmap_pages()[::-1]:
             if page.start >= self.kbase:
+                continue
+            if page.start < self.vmalloc:
                 break
             if page.execute:
                 res = page.start
-        if res is None:
-            return INVALID_ADDR
-        prev = None
-        for page in get_memory_map_raw():
-            if page.start >= res:
-                if prev is not None and page.start > prev + 0x1000:
-                    break
-                prev = self.module_end = page.end
+                break
         return res
 
     def _PAGE_OFFSET(self, va):  # aka PAGE_START
@@ -582,7 +569,7 @@ class Aarch64PagingInfo(ArchPagingInfo):
             page = pages[i]
             if page.start > self.kbase + self.ksize:
                 continue
-            if self.module_start <= page.start < self.module_end:
+            if self.module_start and self.module_start <= page.start < self.kbase:
                 page.objfile = self.KERNELDRIVER
                 continue
             if page.start < self.kbase:

pwndbg/aglib/kernel/vmmap.py

@@ -34,6 +34,7 @@ class KernelVmmap:
         self.pi = pwndbg.aglib.kernel.arch_paginginfo()
         if self.pi:
             self.sections = self.pi.markers()
+        self.adjust()
 
     def get_name(self, addr: int) -> str:
         if addr is None or self.sections is None:
@@ -396,7 +397,16 @@ Note that the page-tables method will require the QEMU kernel process to be on t
 )
 
 
-def kernel_vmmap(process_pages=True) -> Tuple[pwndbg.lib.memory.Page, ...]:
+@pwndbg.lib.cache.cache_until("stop")
+def kernel_vmmap_pages() -> Tuple[pwndbg.lib.memory.Page, ...]:
+    if kernel_vmmap_mode == "page-tables":
+        return kernel_vmmap_via_page_tables()
+    elif kernel_vmmap_mode == "monitor":
+        return kernel_vmmap_via_monitor_info_mem()
+    return ()
+
+
+def kernel_vmmap() -> Tuple[pwndbg.lib.memory.Page, ...]:
     if not pwndbg.aglib.qemu.is_qemu_kernel():
         return ()
 
@@ -409,23 +419,16 @@ def kernel_vmmap(process_pages=True) -> Tuple[pwndbg.lib.memory.Page, ...]:
     ):
         return ()
 
-    pages = None
-    if kernel_vmmap_mode == "page-tables":
-        pages = kernel_vmmap_via_page_tables()
-    elif kernel_vmmap_mode == "monitor":
-        pages = kernel_vmmap_via_monitor_info_mem()
-    if pages is None:
-        return ()
-    if process_pages:
-        kv = KernelVmmap(pages)
-        kv.adjust()
-        if kernel_vmmap_mode == "monitor" and pwndbg.aglib.arch.name == "x86-64":
-            # TODO: check version here when QEMU displays the x bit for x64
-            for page in pages:
-                if page.objfile == kv.pi.ESPSTACK:
-                    continue
-                entry = pwndbg.aglib.kernel.pagewalk(page.start)[0].entry
-                if entry and entry >> 63 == 0:
-                    page.flags |= 1
+    pages = kernel_vmmap_pages()
+    kv = KernelVmmap(pages)
+    if kernel_vmmap_mode == "monitor" and pwndbg.aglib.arch.name == "x86-64":
+        # TODO: check version here when QEMU displays the x bit for x64
+        # see: https://github.com/pwndbg/pwndbg/pull/3020#issuecomment-2914573242
+        for page in pages:
+            if page.objfile == kv.pi.ESPSTACK:
+                continue
+            entry = pwndbg.aglib.kernel.pagewalk(page.start)[0].entry
+            if entry and entry >> 63 == 0:
+                page.flags |= 1
 
     return tuple(pages)