[PATCH makedumpfile v2 8/9] Add userstack extension

Some kernel crashes are triggered by userspace. It can be helpful to
debug the user processes to understand why this occurred, but including
all userspace memory in is usually not feasible. The minimal useful
information is a stack trace, which can be pieced together from the top
of the userspace stack, and the unwind info from the executable and
DSOs. The drgn contrib/pstack.py script can be used to create such a
stack trace, assuming that the stack memory pages are available. Add
an extension which can include those memory pages into a vmcore without
including *all* of userspace memory.

To implement this extension, add helpers which walk the VMA tree for
both rbtree and maple tree implementations. To avoid issues with buggy
kernel data, use cycle detection when walking linked lists. Iterate each
task, identify the stack pointer, and find the corresponding VMA for the
region. Build a sorted list of anon_vmas along with the desired start
and end offset. When filtering pages, for any anon_vma, query this list
and include the desired pages.

Signed-off-by: Stephen Brennan <[email protected]>
---
 extensions/Makefile     |   4 +-
 extensions/list.h       | 106 ++++++++++++
 extensions/userstack.c  | 351 ++++++++++++++++++++++++++++++++++++++++
 extensions/vma_mtree.c  | 140 ++++++++++++++++
 extensions/vma_mtree.h  |   7 +
 extensions/vma_rbtree.c |  56 +++++++
 extensions/vma_rbtree.h |  12 ++
 7 files changed, 675 insertions(+), 1 deletion(-)
 create mode 100644 extensions/list.h
 create mode 100644 extensions/userstack.c
 create mode 100644 extensions/vma_mtree.c
 create mode 100644 extensions/vma_mtree.h
 create mode 100644 extensions/vma_rbtree.c
 create mode 100644 extensions/vma_rbtree.h

diff --git a/extensions/Makefile b/extensions/Makefile
index 25e0a07..e8c41d8 100644
--- a/extensions/Makefile
+++ b/extensions/Makefile
@@ -1,5 +1,5 @@
 CC ?= gcc
-CONTRIB_SO := sample.so erase_sample.so
+CONTRIB_SO := sample.so erase_sample.so userstack.so
 
 all: $(CONTRIB_SO)
 
@@ -8,6 +8,8 @@ CFLAGS += -fPIC -shared -Wl,-T,../makedumpfile.ld
 $(CONTRIB_SO): %.so: %.c
 	$(CC) $(CFLAGS) -o $@ $^
 
+userstack.so: vma_rbtree.c vma_mtree.c
+
 clean:
 	rm -f $(CONTRIB_SO)
 
diff --git a/extensions/list.h b/extensions/list.h
new file mode 100644
index 0000000..75c63b0
--- /dev/null
+++ b/extensions/list.h
@@ -0,0 +1,106 @@
+#ifndef MAKEDUMPFILE_EXT_LIST_H
+#define MAKEDUMPFILE_EXT_LIST_H
+
+#include "../makedumpfile.h"
+#include "../btf_info.h"
+#include "../detect_cycle.h"
+
+#define LIST_READERR  0UL
+#define LIST_CYCLE    1UL
+#define LIST_INVAL    2UL
+#define LIST_ALLOCERR 3UL
+#define LIST_END      4UL
+
+#define LIST_ERR(pos) ((pos) <= LIST_ALLOCERR)
+#define LIST_STOP(pos) ((pos) <= LIST_END)
+
+DECLARE_MOD_STRUCT_MEMBER(vmlinux, list_head, next);
+DECLARE_MOD_STRUCT_MEMBER(vmlinux, list_head, prev);
+
+static void *list_next(void *prev, void *data)
+{
+	unsigned long head = (unsigned long)data;
+	unsigned long next;
+	unsigned long nextprev;
+
+	if (!readmem(VADDR, (unsigned long)prev, &next, sizeof(next)))
+		return (void *)LIST_READERR;
+	if (!readmem(VADDR, next + (GET_MOD_STRUCT_MEMBER_MOFF(vmlinux, list_head, prev) / 8),
+		     &nextprev, sizeof(nextprev)))
+		return (void *)LIST_READERR;
+	if ((unsigned long)prev != nextprev)
+		return (void *)LIST_INVAL;
+	if (next == head)
+		return (void *)LIST_END;
+	return (void *)next;
+}
+
+static inline struct detect_cycle *list_iterator_start(unsigned long head)
+{
+	return dc_init((void *)head, (void *)head, list_next);
+}
+
+static inline unsigned long list_iterator_next(struct detect_cycle *dc, unsigned long offset)
+{
+	void *nextp;
+	int is_cycle;
+	if (!dc)
+		return LIST_ALLOCERR;
+	is_cycle = dc_next(dc, &nextp);
+	if (is_cycle) {
+		free(dc);
+		return LIST_CYCLE;
+	}
+	unsigned long next = (unsigned long) nextp;
+	if (!LIST_STOP(next))
+		next -= offset;
+	else
+		free(dc);
+	return next;
+}
+
+static inline int list_iterator_errmsg(unsigned long pos, const char *prefix) {
+	switch (pos) {
+	case LIST_READERR:
+		ERRMSG("%s: error reading next pointer\n", prefix);
+		break;
+	case LIST_CYCLE:
+		ERRMSG("%s: detected cycle\n", prefix);
+		break;
+	case LIST_INVAL:
+		ERRMSG("%s: corrupt list (invalid prev pointer)\n", prefix);
+		break;
+	case LIST_ALLOCERR:
+		ERRMSG("%s: allocation error\n", prefix);
+		break;
+	default:
+		ERRMSG("%s: BUG: list_iterator_errmsg() with no error\n", prefix);
+		break;
+	}
+	return FALSE;
+}
+
+/*
+ * Iterate over each object in a linked list.
+ *
+ * pos: name of an unsigned long variable which is set to the address of
+ *    each object
+ * head: address of the list_head anchoring the list
+ * offset: offset of the list_head within each object
+ *
+ * Unlike the standard kernel list_for_each_entry() implementation, we perform
+ * list validation. That is, we ensure that (a) each prev pointer points back to
+ * the correct head, and (b) there are no cycles. The loop terminates for errors
+ * or success, so the "pos" variable does double-duty as a status variable. At
+ * the end of the loop, users must use LIST_ERR() to check whether an iteration
+ * error occurred, and if so, they are encouraged to use list_iterator_errmsg to
+ * report an error and return FALSE:
+ *
+ *    if (LIST_ERR(pos))
+ *        return list_iterator_errmsg(pos, "descriptive prefix");
+ */
+#define list_for_each_entry(pos, head, offset) \
+	for (struct detect_cycle *__dc = list_iterator_start(head); \
+	     !LIST_STOP(pos = list_iterator_next(__dc, offset));)
+
+#endif // MAKEDUMPFILE_EXT_LIST_H
diff --git a/extensions/userstack.c b/extensions/userstack.c
new file mode 100644
index 0000000..a26656a
--- /dev/null
+++ b/extensions/userstack.c
@@ -0,0 +1,351 @@
+/*
+ * userstack.c: An extension for preserving userspace stack memory pages
+ *
+ * It can be useful to know what userspace tasks were doing at the time of a
+ * crash, but including all userspace memory is usually too much: usually a
+ * simple stack trace would do the trick. This extension preserves the topmost
+ * userspace stack pages for each thread in each process, making it possible
+ * to create a stack trace with a tool such as contrib/pstack.py in drgn.
+ */
+#include <assert.h>
+#include <stdbool.h>
+
+#include "../extension.h"
+#include "../makedumpfile.h"
+#include "../btf_info.h"
+#include "../kallsyms.h"
+#include "vma_mtree.h"
+#include "vma_rbtree.h"
+#include "list.h"
+
+/* Required struct fields */
+INIT_MOD_STRUCT_MEMBER(vmlinux, task_struct, tasks);
+INIT_MOD_STRUCT_MEMBER(vmlinux, task_struct, signal);
+INIT_MOD_STRUCT_MEMBER(vmlinux, task_struct, thread_node);
+INIT_MOD_STRUCT_MEMBER(vmlinux, task_struct, stack);
+INIT_MOD_STRUCT_MEMBER(vmlinux, task_struct, mm);
+INIT_MOD_STRUCT_MEMBER(vmlinux, vm_area_struct, anon_vma);
+INIT_MOD_STRUCT_MEMBER(vmlinux, vm_area_struct, vm_pgoff);
+INIT_MOD_STRUCT_MEMBER(vmlinux, page, index);
+INIT_MOD_STRUCT_MEMBER(vmlinux, signal_struct, thread_head);
+INIT_MOD_STRUCT_MEMBER(vmlinux, list_head, next);
+INIT_MOD_STRUCT_MEMBER(vmlinux, list_head, prev);
+INIT_MOD_STRUCT_MEMBER(vmlinux, pt_regs, sp);
+INIT_MOD_STRUCT(vmlinux, pt_regs);
+
+/* Optional struct fields */
+INIT_OPT_MOD_STRUCT_MEMBER(vmlinux, thread_union, stack);
+INIT_OPT_MOD_STRUCT_MEMBER(vmlinux, mm_struct, mm_mt);
+INIT_OPT_MOD_STRUCT_MEMBER(vmlinux, mm_struct, mm_rb);
+
+/* Required symbols */
+INIT_MOD_SYM(vmlinux, init_task);
+
+/* Optional symbols */
+INIT_OPT_MOD_SYM(vmlinux, fred_rsp0);
+INIT_OPT_MOD_SYM(vmlinux, __start_init_stack);
+INIT_OPT_MOD_SYM(vmlinux, __end_init_stack);
+INIT_OPT_MOD_SYM(vmlinux, __start_init_task);
+INIT_OPT_MOD_SYM(vmlinux, __end_init_task);
+
+unsigned long THREAD_SIZE;
+bool ready;
+
+#define MEMBER_OFF(S, M) \
+	(GET_MOD_STRUCT_MEMBER_MOFF(vmlinux, S, M) / 8)
+
+
+static int for_each_task(bool (*task_fn)(unsigned long))
+{
+	unsigned long curr_proc;
+
+	// NOTE: this explicitly skips "init_task" because it treats it as the
+	// head of the list. This is fine: init_task is a kernel thread, so
+	// never has a stack to retain.
+	list_for_each_entry(curr_proc,
+			    GET_MOD_SYM(vmlinux, init_task) + MEMBER_OFF(task_struct, tasks),
+			    MEMBER_OFF(task_struct, tasks)) {
+		unsigned long mm;
+		if (!readmem(VADDR, curr_proc + MEMBER_OFF(task_struct, mm),
+			     &mm, sizeof(mm))) {
+			ERRMSG("error: failed to read task.mm\n");
+		}
+		if (!mm)
+			continue;
+
+		unsigned long signal;
+		if (!readmem(VADDR, curr_proc + MEMBER_OFF(task_struct, signal),
+			     &signal, sizeof(signal))) {
+			ERRMSG("error: failed to read task.signal\n");
+			break;
+		}
+
+		unsigned long curr_thread;
+		list_for_each_entry(curr_thread,
+				    signal + MEMBER_OFF(signal_struct, thread_head),
+				    MEMBER_OFF(task_struct, thread_node)) {
+			if (!task_fn(curr_thread))
+				return FALSE;
+		}
+		if (LIST_ERR(curr_thread))
+			return list_iterator_errmsg(curr_thread, "iterating thread list");
+	}
+	if (LIST_ERR(curr_proc))
+		return list_iterator_errmsg(curr_proc, "iterating task list");
+
+	return TRUE;
+}
+
+static unsigned long task_sp(unsigned long taskp)
+{
+	// The stack pointer is stored on entry to the kernel at the top of the
+	// kernel stack. If the task in on-cpu, the stack pointer will be in the
+	// PRSTATUS, but the stale value is very likely to be useful enough.
+	unsigned long user_sp_loc;
+	if (!readmem(VADDR, taskp + MEMBER_OFF(task_struct, stack),
+		     &user_sp_loc, sizeof(user_sp_loc)))
+		return 0;
+
+	user_sp_loc += THREAD_SIZE;
+	user_sp_loc -= GET_MOD_STRUCT_SSIZE(vmlinux, pt_regs);
+	if (MOD_SYM_EXIST(vmlinux, fred_rsp0))
+		user_sp_loc -= 16;
+	user_sp_loc += MEMBER_OFF(pt_regs, sp);
+
+	unsigned long sp;
+	if (!readmem(VADDR, user_sp_loc, &sp, sizeof(sp)))
+		return 0;
+	return sp;
+}
+
+struct task_stack {
+	unsigned long anon_vma;
+	unsigned long index_start;
+	unsigned long index_end;
+};
+
+static struct task_stack *stacks;
+static size_t stacks_count;
+static size_t stacks_alloc;
+
+// Avoid using too much memory when processing an especially large vmcore, or in
+// the case of a bug that causes us to create too many entries. 1M threads
+// requires 24 MiB of memory to track. While it's not the maximum amount we
+// could see, by a long shot, it's enough where most common workloads won't hit
+// it, and any more than this will make it far more likely that the dump will
+// hit an OOM issue.
+#define MAX_TASK_STACKS (1LU << 20) /* 1M * 24 bytes = 24 MiB */
+
+static bool append_task_stack(struct task_stack *newstack)
+{
+	if (stacks_count >= MAX_TASK_STACKS) {
+		ERRMSG("userstack error: hit maximum stack count %lu, aborting collection\n", MAX_TASK_STACKS);
+		goto fail;
+	}
+	if (stacks_count == stacks_alloc) {
+		if (stacks_alloc)
+			stacks_alloc *= 2;
+		else
+			stacks_alloc = 512;
+		struct task_stack *newarr = realloc(stacks, stacks_alloc * sizeof(stacks[0]));
+		if (!newarr) {
+			ERRMSG("userstack: allocation error for stack tracking (size %lu)\n", stacks_alloc);
+			goto fail;
+		}
+		stacks = newarr;
+	}
+	stacks[stacks_count++] = *newstack;
+	return TRUE;
+
+fail:
+	free(stacks);
+	stacks = NULL;
+	stacks_count = stacks_alloc = 0;
+	return FALSE;
+}
+
+static bool record_task_stack(unsigned long taskp)
+{
+	unsigned long task_mm;
+	if (!readmem(VADDR, taskp + MEMBER_OFF(task_struct, mm), &task_mm, sizeof(task_mm)))
+		/* Propagate failure to read a value from the task_struct */
+		return FALSE;
+	if (!task_mm)
+		/* NULL task_mm is expected, continue */
+		return TRUE;
+
+	unsigned long sp = task_sp(taskp);
+	if (!sp)
+		/* Propagate error reading SP */
+		return FALSE;
+
+	int ret;
+	unsigned long vma = 0;
+	if (MOD_STRUCT_MEMBER_EXIST(vmlinux, mm_struct, mm_mt))
+		ret = find_vma_mtree(task_mm + MEMBER_OFF(mm_struct, mm_mt), sp, &vma);
+	else if (MOD_STRUCT_MEMBER_EXIST(vmlinux, mm_struct, mm_rb))
+		ret = find_vma_rbtree(task_mm + MEMBER_OFF(mm_struct, mm_rb), sp, &vma);
+	else
+		assert(FALSE); /* should be impossible */
+	if (!ret)
+		/* propagate error from find_vma_xxx() */
+		return FALSE;
+	else if (!vma)
+		/* No VMA found for the stack. This is unexpected, but gracefully
+		 * handle the condition and continue. */
+		return TRUE;
+
+	unsigned long vm_start, vm_end, anon_vma, vm_pgoff;
+	if (!readmem(VADDR, vma + MEMBER_OFF(vm_area_struct, anon_vma), &anon_vma, sizeof(anon_vma)))
+		return FALSE;
+
+	if (!anon_vma)
+		/* Not an anonymous VMA. This is unexpected but valid. Move on to
+		 * the next task. */
+		return TRUE;
+
+	if (!readmem(VADDR, vma + MEMBER_OFF(vm_area_struct, vm_start), &vm_start, sizeof(vm_start)) ||
+	    !readmem(VADDR, vma + MEMBER_OFF(vm_area_struct, vm_end), &vm_end, sizeof(vm_end)) ||
+	    !readmem(VADDR, vma + MEMBER_OFF(vm_area_struct, vm_pgoff), &vm_pgoff, sizeof(vm_pgoff)))
+		return FALSE;
+
+	/* Construct a range of indices we would like to retain. This is the
+	 * range of stack pages starting with the stack pointer, and continuing
+	 * to the top of the stack vma, or until a limit of 128 pages per task
+	 * is reached. */
+	unsigned long pgoff_start = (sp - vm_start) >> PAGESHIFT();
+	pgoff_start += vm_pgoff;
+	unsigned long pgoff_end = (vm_end - vm_start) >> PAGESHIFT();
+	pgoff_end += vm_pgoff;
+	if (pgoff_start + 128 < pgoff_end)
+		pgoff_end = pgoff_start + 128;
+
+	struct task_stack stack = {anon_vma | 1, pgoff_start, pgoff_end};
+	if (!append_task_stack(&stack))
+		return FALSE;
+
+	return TRUE;
+}
+
+static int stack_compar(const void *lhs, const void *rhs)
+{
+	const struct task_stack *lhss = lhs, *rhss = rhs;
+	if (lhss->anon_vma < rhss->anon_vma)
+		return -1;
+	else if (lhss->anon_vma > rhss->anon_vma)
+		return 1;
+	else
+		return 0;
+}
+
+void extension_init(void)
+{
+	if (MOD_STRUCT_MEMBER_EXIST(vmlinux, mm_struct, mm_mt)) {
+		if (!vma_mtree_init())
+			return;
+	} else if (MOD_STRUCT_MEMBER_EXIST(vmlinux, mm_struct, mm_rb)) {
+		if (!vma_rbtree_init())
+			return;
+	} else {
+		ERRMSG("error: Neither mtree nor rbtree available for VMA walking\n");
+		return;
+	}
+	if (!MOD_STRUCT_EXIST(vmlinux, pt_regs)) {
+		ERRMSG("error: missing pt_regs incfo\n");
+		return;
+	}
+	if (!MOD_SYM_EXIST(vmlinux, init_task)) {
+		ERRMSG("error: missing init_task symbol\n");
+		return;
+	}
+
+	// Determine THREAD_SIZE, which is necessary to find the offset of the
+	// userspace stack pointer register from the kernel thread stack.
+	//
+	// - Prior to v4.16, 0500871f21b23 ("Construct init thread stack in the
+	//   linker script rather than by union"), it was found in thread_union.
+	// - Between v4.16 and v6.10, 8f69cba096b5c ("x86: Rename
+	//   __{start,end}_init_task to __{start,end}_init_stack"), the stack
+	//   size can be inferred by the __{start,end}_init_task symbols.
+	// - Since v6.10, the size is inferred by __{start,end}_init_stack.
+	if (MOD_STRUCT_MEMBER_EXIST(vmlinux, thread_union, stack)) {
+		THREAD_SIZE = GET_MOD_STRUCT_MEMBER_MSIZE(vmlinux, thread_union, stack);
+	} else if (MOD_SYM_EXIST(vmlinux, __start_init_stack) &&
+		   MOD_SYM_EXIST(vmlinux, __end_init_stack) &&
+		   GET_MOD_SYM(vmlinux, __end_init_stack) > GET_MOD_SYM(vmlinux, __start_init_stack)) {
+		THREAD_SIZE = GET_MOD_SYM(vmlinux, __end_init_stack) - GET_MOD_SYM(vmlinux, __start_init_stack);
+	} else if (MOD_SYM_EXIST(vmlinux, __start_init_task) &&
+		   MOD_SYM_EXIST(vmlinux, __end_init_task) &&
+		   GET_MOD_SYM(vmlinux, __end_init_task) > GET_MOD_SYM(vmlinux, __start_init_task)) {
+		THREAD_SIZE = GET_MOD_SYM(vmlinux, __end_init_task) - GET_MOD_SYM(vmlinux, __start_init_task);
+	} else {
+		ERRMSG("Could not determine THREAD_SIZE: neither __start_init_stack "
+		       "nor __start_init_task found in kallsyms, nor is thread_union "
+		       "found in BTF.\n");
+		return;
+	}
+
+	if (!for_each_task(&record_task_stack)) {
+		free(stacks);
+		stacks = NULL;
+		stacks_alloc = stacks_count = 0;
+		ERRMSG("Could not gather all task stack VMAs, userstack disabled\n");
+		return;
+	}
+	struct task_stack *tmp = realloc(stacks, stacks_count * sizeof(*tmp));
+	if (tmp) {
+		stacks = tmp;
+		stacks_alloc = stacks_count;
+	}
+	qsort(stacks, stacks_count, sizeof(*stacks), &stack_compar);
+	ready = TRUE;
+}
+
+static int count_retained;
+static int count_checked;
+static int count_cached;
+int extension_callback(unsigned long pfn, const void *pcache, const struct pginfo *i)
+{
+	unsigned long index;
+	static struct {
+		unsigned long mapping;
+		struct task_stack *result;
+	} cache;
+
+	if (!ready || !isAnon(i))
+		return PG_UNDECID;
+
+	index = ULONG(pcache + MEMBER_OFF(page, index));
+	if (!(i->mapping & 1))
+		return PG_UNDECID;
+
+	if (i->mapping != cache.mapping) {
+		count_checked++;
+		struct task_stack search = {i->mapping, 0, 0};
+		struct task_stack *result = bsearch(&search, stacks, stacks_count,
+						sizeof(search), &stack_compar);
+		if (!result)
+			return PG_UNDECID;
+
+		cache.mapping = i->mapping;
+		cache.result = result;
+	} else {
+		count_cached++;
+	}
+
+	if (cache.result->index_start <= index && index < cache.result->index_end) {
+		count_retained++;
+		return PG_INCLUDE;
+	} else {
+		return PG_UNDECID;
+	}
+}
+
+__attribute__((destructor))
+static void userstack_exit(void) {
+	if (count_retained || count_checked || count_cached || stacks_count) {
+		REPORT_MSG("Extension userstack:\n");
+		REPORT_MSG("  PFNs retained: %d searched: %d, cached: %d\n", count_retained, count_checked, count_cached);
+		REPORT_MSG("  Recorded %zu stack anon_vmas\n", stacks_count);
+	}
+}
diff --git a/extensions/vma_mtree.c b/extensions/vma_mtree.c
new file mode 100644
index 0000000..ac8912e
--- /dev/null
+++ b/extensions/vma_mtree.c
@@ -0,0 +1,140 @@
+#include <stdbool.h>
+#include "../btf_info.h"
+#include "../kallsyms.h"
+#include "../makedumpfile.h"
+
+INIT_OPT_MOD_STRUCT(vmlinux, maple_tree);
+INIT_OPT_MOD_STRUCT(vmlinux, maple_node);
+INIT_OPT_MOD_STRUCT_MEMBER(vmlinux, maple_tree, ma_root);
+INIT_OPT_MOD_STRUCT_MEMBER(vmlinux, maple_arange_64, pivot);
+INIT_OPT_MOD_STRUCT_MEMBER(vmlinux, maple_arange_64, slot);
+INIT_OPT_MOD_STRUCT_MEMBER(vmlinux, maple_arange_64, meta);
+INIT_OPT_MOD_STRUCT_MEMBER(vmlinux, maple_range_64, pivot);
+INIT_OPT_MOD_STRUCT_MEMBER(vmlinux, maple_range_64, slot);
+INIT_OPT_MOD_STRUCT_MEMBER(vmlinux, maple_range_64, meta);
+INIT_OPT_MOD_STRUCT_MEMBER(vmlinux, maple_metadata, end);
+
+#define MEMBER_OFF(S, M) \
+	GET_MOD_STRUCT_MEMBER_MOFF(vmlinux, S, M) / 8
+#define KERN_STRUCT_MEMBER_EXIST(S, M) \
+	MOD_STRUCT_MEMBER_EXIST(vmlinux, S, M)
+#define GET_KERN_SYM(SYM) GET_MOD_SYM(vmlinux, SYM)
+#define KERN_SYM_EXIST(SYM) MOD_SYM_EXIST(vmlinux, SYM)
+#define GET_KERN_STRUCT_SSIZE(S) \
+	GET_MOD_STRUCT_SSIZE(vmlinux, S)
+#define KERN_STRUCT_EXIST(SYM) MOD_STRUCT_EXIST(vmlinux, SYM)
+
+#define MAPLE_NODE_MASK			255UL
+#define MAPLE_NODE_TYPE_MASK		0x0F
+#define MAPLE_NODE_TYPE_SHIFT		0x03
+#define XA_ZERO_ENTRY			xa_mk_internal(257)
+
+static unsigned long xa_mk_internal(unsigned long v)
+{
+	return (v << 2) | 2;
+}
+
+static bool xa_is_internal(unsigned long entry)
+{
+	return (entry & 3) == 2;
+}
+
+static bool xa_is_node(unsigned long entry)
+{
+	return xa_is_internal(entry) && entry > 4096;
+}
+
+bool vma_mtree_init(void)
+{
+	if (!KERN_STRUCT_EXIST(maple_tree) ||
+	    !KERN_STRUCT_EXIST(maple_node) ||
+	    !KERN_STRUCT_MEMBER_EXIST(maple_tree, ma_root) ||
+	    !KERN_STRUCT_MEMBER_EXIST(maple_arange_64, pivot) ||
+	    !KERN_STRUCT_MEMBER_EXIST(maple_arange_64, slot) ||
+	    !KERN_STRUCT_MEMBER_EXIST(maple_arange_64, meta) ||
+	    !KERN_STRUCT_MEMBER_EXIST(maple_range_64, pivot) ||
+	    !KERN_STRUCT_MEMBER_EXIST(maple_range_64, slot) ||
+	    !KERN_STRUCT_MEMBER_EXIST(maple_range_64, meta) ||
+	    !KERN_STRUCT_MEMBER_EXIST(maple_metadata, end)) {
+		ERRMSG("Missing required maple tree syms/types\n");
+		return false;
+	}
+
+	return true;
+}
+
+int find_vma_mtree(unsigned long mt, unsigned long index, unsigned long *ret)
+{
+	unsigned long long entry;
+
+	if (!readmem(VADDR, mt + MEMBER_OFF(maple_tree, ma_root), &entry, sizeof(entry)))
+		return FALSE;
+
+	if (!xa_is_node(entry)) {
+		if (index == 0)
+			*ret = entry;
+		else
+			*ret = 0;
+		return TRUE;
+	}
+	unsigned long long max = ULONGLONG_MAX;
+	void *node = malloc(GET_KERN_STRUCT_SSIZE(maple_node));
+	if (!node) {
+		ERRMSG("failed to allocate memory\n");
+		return FALSE;
+	}
+
+	for (;;) {
+		if (!readmem(VADDR, entry & ~MAPLE_NODE_MASK, node, GET_KERN_STRUCT_SSIZE(maple_node))) {
+			ERRMSG("failed to read maple node: %llx\n", entry & ~MAPLE_NODE_MASK);
+			free(node);
+			return FALSE;
+		}
+
+		int node_type = (entry >> MAPLE_NODE_TYPE_SHIFT) & MAPLE_NODE_TYPE_MASK;
+		unsigned long long *pivot, *slot;
+		uint8_t end;
+		if (node_type == 3) {
+			pivot = node + MEMBER_OFF(maple_arange_64, pivot);
+			slot = node + MEMBER_OFF(maple_arange_64, slot);
+			end = ((uint8_t *)node)[MEMBER_OFF(maple_arange_64, meta) + MEMBER_OFF(maple_metadata, end)];
+		} else if (node_type == 1 || node_type == 2) {
+			pivot = node + MEMBER_OFF(maple_range_64, pivot);
+			slot = node + MEMBER_OFF(maple_range_64, slot);
+			unsigned long long p = *(slot - 1);
+			if (!p)
+				end = ((uint8_t *)node)[MEMBER_OFF(maple_range_64, meta) + MEMBER_OFF(maple_metadata, end)];
+			else {
+				end = slot - pivot;
+				if (p == max)
+					end--;
+			}
+		} else {
+			ERRMSG("unrecognized maple node type: %d\n", node_type);
+			free(node);
+			return FALSE;
+		}
+		int offset = 0;
+		for (offset = 0; offset < end; offset++) {
+			if (pivot[offset] >= index) {
+				max = pivot[offset];
+				break;
+			}
+		}
+		if (&pivot[offset] >= slot)
+			offset = end;
+
+		entry = slot[offset];
+		if (node_type == 1) {
+			// leaf:
+			free(node);
+			if (entry == XA_ZERO_ENTRY)
+				*ret = 0;
+			else
+				*ret = entry;
+			return TRUE;
+		}
+	}
+	*ret = 0;
+	return TRUE;
+}
diff --git a/extensions/vma_mtree.h b/extensions/vma_mtree.h
new file mode 100644
index 0000000..69d9448
--- /dev/null
+++ b/extensions/vma_mtree.h
@@ -0,0 +1,7 @@
+#ifndef _MAPLE_TREE_H
+#define _MAPLE_TREE_H
+#include <stdbool.h>
+bool vma_mtree_init(void);
+int find_vma_mtree(unsigned long mt, unsigned long address, unsigned long *ret);
+#endif /* _MAPLE_TREE_H */
+
diff --git a/extensions/vma_rbtree.c b/extensions/vma_rbtree.c
new file mode 100644
index 0000000..c947fbc
--- /dev/null
+++ b/extensions/vma_rbtree.c
@@ -0,0 +1,56 @@
+#include "../makedumpfile.h"
+#include "../btf_info.h"
+#include "vma_rbtree.h"
+
+INIT_MOD_STRUCT_MEMBER(vmlinux, vm_area_struct, vm_start);
+INIT_MOD_STRUCT_MEMBER(vmlinux, vm_area_struct, vm_end);
+
+INIT_OPT_MOD_STRUCT_MEMBER(vmlinux, vm_area_struct, vm_rb);
+INIT_OPT_MOD_STRUCT_MEMBER(vmlinux, rb_root, rb_node);
+INIT_OPT_MOD_STRUCT_MEMBER(vmlinux, rb_node, rb_left);
+INIT_OPT_MOD_STRUCT_MEMBER(vmlinux, rb_node, rb_right);
+
+#define MEMBER_OFF(S, M) \
+	(GET_MOD_STRUCT_MEMBER_MOFF(vmlinux, S, M) / 8)
+
+#define MEMBER_EXIST(S, M) \
+	(MOD_STRUCT_MEMBER_EXIST(vmlinux, S, M))
+
+int find_vma_rbtree(unsigned long rb_root, unsigned long address, unsigned long *ret)
+{
+	unsigned long node, vma, vm_start, vm_end, rb_left, rb_right;
+	if (!readmem(VADDR, rb_root, &node, sizeof(node)))
+		return FALSE;
+
+	*ret = 0;
+	while (node > MEMBER_OFF(vm_area_struct, vm_rb)) {
+		vma = node - MEMBER_OFF(vm_area_struct, vm_rb);
+		if (!readmem(VADDR, vma + MEMBER_OFF(vm_area_struct, vm_start), &vm_start, sizeof(vm_start)) ||
+		    !readmem(VADDR, vma + MEMBER_OFF(vm_area_struct, vm_end), &vm_end, sizeof(vm_end)) ||
+		    !readmem(VADDR, node + MEMBER_OFF(rb_node, rb_left), &rb_left, sizeof(rb_left)) ||
+		    !readmem(VADDR, node + MEMBER_OFF(rb_node, rb_right), &rb_right, sizeof(rb_right)))
+			return FALSE;
+
+		if (address < vm_start) {
+			node = rb_left;
+		} else if (address >= vm_end) {
+			node = rb_right;
+		} else {
+			*ret = vma;
+			break;
+		}
+	}
+	return TRUE;
+}
+
+bool vma_rbtree_init(void)
+{
+	if (!MEMBER_EXIST(vm_area_struct, vm_rb) ||
+	    !MEMBER_EXIST(rb_root, rb_node) ||
+	    !MEMBER_EXIST(rb_node, rb_left) ||
+	    !MEMBER_EXIST(rb_node, rb_right)) {
+		ERRMSG("error: missing required vm_area_struct & rbtree definitions");
+		return false;
+	}
+	return true;
+}
diff --git a/extensions/vma_rbtree.h b/extensions/vma_rbtree.h
new file mode 100644
index 0000000..a8b8d4f
--- /dev/null
+++ b/extensions/vma_rbtree.h
@@ -0,0 +1,12 @@
+#ifndef RBTREE_H_
+#define RBTREE_H_
+#include <stdbool.h>
+
+#include "../btf_info.h"
+
+DECLARE_MOD_STRUCT_MEMBER(vmlinux, vm_area_struct, vm_start);
+DECLARE_MOD_STRUCT_MEMBER(vmlinux, vm_area_struct, vm_end);
+
+int find_vma_rbtree(unsigned long rb_root, unsigned long address, unsigned long *ret);
+bool vma_rbtree_init(void);
+#endif // RBTREE_H_
-- 
2.52.0