libdrgn: move _vmemmap() to object finder

Similarly to PAGE_OFFSET, vmemmap makes more sense as part of the Linux
kernel object finder than an internal helper.

While we're here, let's fix the definition for 5-level page tables. This
only matters for kernels with commit 77ef56e4f0fb ("x86: Enable 5-level
paging support via CONFIG_X86_5LEVEL=y") but without eedb92abb9bb
("x86/mm: Make virtual memory layout dynamic for CONFIG_X86_5LEVEL=y")
(namely, v4.14, v4.15, and v4.16); since v4.17, 5-level page table
support enables KASLR.

Author: osandov

drgn/helpers/linux/mm.py

@@ -35,22 +35,13 @@ def pgtable_l5_enabled(prog):
     return _linux_helper_pgtable_l5_enabled(prog)
 
 
-def _vmemmap(prog):
-    try:
-        # KASAN
-        return cast("struct page *", prog["vmemmap_base"])
-    except KeyError:
-        # x86-64
-        return Object(prog, "struct page *", value=0xFFFFEA0000000000)
-
-
 def for_each_page(prog):
     """
     Iterate over all pages in the system.
 
     :return: Iterator of ``struct page *`` objects.
     """
-    vmemmap = _vmemmap(prog)
+    vmemmap = prog["vmemmap"]
     for i in range(prog["max_pfn"]):
         yield vmemmap + i
 
@@ -61,7 +52,7 @@ def page_to_pfn(page):
 
     Get the page frame number (PFN) of a page.
     """
-    return cast("unsigned long", page - _vmemmap(page.prog_))
+    return cast("unsigned long", page - page.prog_["vmemmap"])
 
 
 def pfn_to_page(prog_or_pfn, pfn=None):
@@ -76,7 +67,7 @@ def pfn_to_page(prog_or_pfn, pfn=None):
         pfn = prog_or_pfn
     else:
         prog = prog_or_pfn
-    return _vmemmap(prog) + pfn
+    return prog["vmemmap"] + pfn
 
 
 def virt_to_pfn(prog_or_addr, addr=None):

libdrgn/arch_x86_64.c.in

@@ -344,6 +344,39 @@ out:
 	return err;
 }
 
+static struct drgn_error *
+linux_kernel_get_vmemmap_x86_64(struct drgn_program *prog, uint64_t *ret)
+{
+
+	struct drgn_error *err;
+	struct drgn_object obj;
+
+	/* If KASLR is enabled, vmemmap is vmemmap_base. */
+	drgn_object_init(&obj, prog);
+	err = drgn_program_find_object(prog, "vmemmap_base", NULL,
+				       DRGN_FIND_OBJECT_VARIABLE, &obj);
+	if (!err) {
+		err = drgn_object_read_unsigned(&obj, ret);
+		goto out;
+	}
+	if (err->code == DRGN_ERROR_LOOKUP) {
+		drgn_error_destroy(err);
+		err = NULL;
+	} else {
+		goto out;
+	}
+
+	/* Otherwise, it depends on whether we have 5-level page tables. */
+	if (prog->vmcoreinfo.pgtable_l5_enabled)
+		*ret = UINT64_C(0xffd4000000000000);
+	else
+		*ret = UINT64_C(0xffffea0000000000);
+
+out:
+	drgn_object_deinit(&obj);
+	return err;
+}
+
 const struct drgn_architecture_info arch_info_x86_64 = {
 	ARCHITECTURE_INFO,
 	.default_flags = (DRGN_PLATFORM_IS_64_BIT |
@@ -352,4 +385,5 @@ const struct drgn_architecture_info arch_info_x86_64 = {
 	.num_frame_registers = ARRAY_SIZE(frame_registers_x86_64),
 	.linux_kernel_set_initial_registers = linux_kernel_set_initial_registers_x86_64,
 	.linux_kernel_get_page_offset = linux_kernel_get_page_offset_x86_64,
+	.linux_kernel_get_vmemmap = linux_kernel_get_vmemmap_x86_64,
 };

libdrgn/linux_kernel.c

@@ -318,6 +318,26 @@ struct drgn_error *linux_kernel_object_find(const char *name, size_t name_len,
 						      prog->vmcoreinfo.osrelease,
 						      0, 0,
 						      DRGN_PROGRAM_ENDIAN);
+		} else if (name_len == strlen("vmemmap") &&
+			   memcmp(name, "vmemmap", name_len) == 0) {
+			if (!prog->vmemmap) {
+				if (!prog->has_platform ||
+				    !prog->platform.arch->linux_kernel_get_vmemmap)
+					return &drgn_not_found;
+				err = prog->platform.arch->linux_kernel_get_vmemmap(prog,
+										    &prog->vmemmap);
+				if (err) {
+					prog->vmemmap = 0;
+					return err;
+				}
+			}
+
+			err = drgn_program_find_type(prog, "struct page *",
+						     NULL, &qualified_type);
+			if (err)
+				return err;
+			return drgn_object_set_unsigned(ret, qualified_type,
+							prog->vmemmap, 0);
 		}
 	}
 	return &drgn_not_found;

libdrgn/platform.h

@@ -37,6 +37,8 @@ struct drgn_architecture_info {
 								 const struct drgn_object *);
 	struct drgn_error *(*linux_kernel_get_page_offset)(struct drgn_program *,
 							   uint64_t *);
+	struct drgn_error *(*linux_kernel_get_vmemmap)(struct drgn_program *,
+						       uint64_t *);
 };
 
 static inline const struct drgn_register *

libdrgn/program.h

@@ -71,6 +71,8 @@ struct drgn_program {
 	struct vmcoreinfo vmcoreinfo;
 	/* Cached PAGE_OFFSET. */
 	uint64_t page_offset;
+	/* Cached vmemmap. */
+	uint64_t vmemmap;
 #ifdef WITH_LIBKDUMPFILE
 	kdump_ctx_t *kdump_ctx;
 #endif