libdrgn: make Linux kernel stack unwinding more robust

drgn has a couple of issues unwinding stack traces for kernel core
dumps:

1. It can't unwind the stack for the idle task (PID 0), which commonly
   appears in core dumps.
2. It uses the PID in PRSTATUS, which is racy and can't actually be
   trusted.

The solution for both of these is to look up the PRSTATUS note by CPU
instead of PID.

For the live kernel, drgn refuses to unwind the stack of tasks in the
"R" state. However, the "R" state is running *or runnable*, so in the
latter case, we can still unwind the stack. The solution for this is to
look at on_cpu for the task instead of the state.

Author: osandov

libdrgn/arch_x86_64.c.in

@@ -257,7 +257,9 @@ out:
 
 static struct drgn_error *
 linux_kernel_set_initial_registers_x86_64(Dwfl_Thread *thread,
-					  const struct drgn_object *task_obj)
+					  const struct drgn_object *task_obj,
+					  const void *prstatus,
+					  size_t prstatus_size)
 {
 	struct drgn_error *err;
 	struct drgn_program *prog = task_obj->prog;
@@ -268,8 +270,44 @@ linux_kernel_set_initial_registers_x86_64(Dwfl_Thread *thread,
 
 	drgn_object_init(&sp_obj, prog);
 
-	/*
-	 */
+	if (prstatus) {
+		/*
+		 * If the stack pointer in PRSTATUS is within this task's stack,
+		 * then we can use it. Otherwise, the task either wasn't running
+		 * or was in the middle of context switching. Either way, we
+		 * should use the saved registers instead.
+		 */
+		uint64_t thread_size;
+		uint64_t stack;
+
+		err = linux_kernel_get_thread_size(prog, &thread_size);
+		if (err)
+			goto out;
+		err = drgn_object_member_dereference(&sp_obj, task_obj,
+						     "stack");
+		if (err)
+			goto out;
+		err = drgn_object_read_unsigned(&sp_obj, &stack);
+		if (err)
+			goto out;
+
+		if (prstatus_size < 272) {
+			err = drgn_error_create(DRGN_ERROR_INVALID_ARGUMENT,
+						"registers are truncated");
+			goto out;
+		}
+		memcpy(&sp, (char *)prstatus + 264, sizeof(sp));
+		if (drgn_program_bswap(prog))
+			sp = bswap_64(sp);
+		if (sp > stack && sp <= stack + thread_size) {
+			err = prstatus_set_initial_registers_x86_64(prog,
+								    thread,
+								    prstatus,
+								    prstatus_size);
+			goto out;
+		}
+	}
+
 	err = drgn_object_member_dereference(&sp_obj, task_obj, "thread");
 	if (err)
 		goto out;

libdrgn/linux_kernel.c

@@ -234,8 +234,8 @@ struct drgn_error *read_vmcoreinfo_fallback(struct drgn_memory_reader *reader,
 	return err;
 }
 
-static struct drgn_error *
-linux_kernel_get_thread_size(struct drgn_program *prog, uint64_t *ret)
+struct drgn_error *linux_kernel_get_thread_size(struct drgn_program *prog,
+						uint64_t *ret)
 {
 	struct drgn_error *err;
 	struct drgn_qualified_type thread_union_type;

libdrgn/linux_kernel.h

@@ -24,6 +24,8 @@ struct drgn_error *proc_kallsyms_symbol_addr(const char *name,
 
 struct drgn_error *read_vmcoreinfo_fallback(struct drgn_memory_reader *reader,
 					    struct vmcoreinfo *ret);
+struct drgn_error *linux_kernel_get_thread_size(struct drgn_program *prog,
+						uint64_t *ret);
 
 struct drgn_error *linux_kernel_object_find(const char *name, size_t name_len,
 					    const char *filename,

libdrgn/platform.h

@@ -65,8 +65,35 @@ struct drgn_architecture_info {
 							     Dwfl_Thread *,
 							     const void *,
 							     size_t);
+	/*
+	 * Get a task's registers from the task_struct or PRSTATUS note as
+	 * appropriate.
+	 *
+	 * The given PRSTATUS note is for the CPU that the task is assigned to,
+	 * which may or may not be for the given task. This callback must
+	 * determine that (typically by checking whether the stack pointer in
+	 * PRSTATUS lies within the task's stack).
+	 *
+	 * We find the PRSTATUS note by CPU rather than by PID for two reasons:
+	 *
+	 * 1. The PID is populated by the kernel from "current" (the current
+	 *    task) via a non-maskable interrupt (NMI). During a context switch,
+	 *    the stack pointer and current are not updated atomically, so if
+	 *    the NMI arrives in the middle of a context switch, the stack
+	 *    pointer may not actually be that of current. Therefore, the stack
+	 *    pointer in PRSTATUS may not actually be for the PID in PRSTATUS.
+	 *
+	 *    We go through all of this trouble because blindly trusting the PID
+	 *    could result in a stack trace for the wrong task, which we want to
+	 *    avoid at all costs.
+	 *
+	 * 2. There is an idle task with PID 0 for each CPU, so for an idle task
+	 *    we have no choice but to find the note by CPU.
+	 */
 	struct drgn_error *(*linux_kernel_set_initial_registers)(Dwfl_Thread *,
-								 const struct drgn_object *);
+								 const struct drgn_object *,
+								 const void *prstatus,
+								 size_t prstatus_size);
 	struct drgn_error *(*linux_kernel_get_page_offset)(struct drgn_program *,
 							   uint64_t *);
 	struct drgn_error *(*linux_kernel_get_vmemmap)(struct drgn_program *,

libdrgn/program.c

@@ -28,6 +28,7 @@
 #include "type_index.h"
 #include "vector.h"
 
+DEFINE_VECTOR_FUNCTIONS(drgn_prstatus_vector)
 DEFINE_HASH_TABLE_FUNCTIONS(drgn_prstatus_map, hash_pair_int_type,
 			    hash_table_scalar_eq)
 
@@ -75,15 +76,19 @@ void drgn_program_init(struct drgn_program *prog,
 	drgn_type_index_init(&prog->tindex);
 	drgn_object_index_init(&prog->oindex);
 	prog->core_fd = -1;
-	drgn_prstatus_map_init(&prog->prstatus_cache);
 	if (platform)
 		drgn_program_set_platform(prog, platform);
 }
 
 void drgn_program_deinit(struct drgn_program *prog)
 {
 	free(prog->task_state_chars);
-	drgn_prstatus_map_deinit(&prog->prstatus_cache);
+	if (prog->prstatus_cached) {
+		if (prog->flags & DRGN_PROGRAM_IS_LINUX_KERNEL)
+			drgn_prstatus_vector_deinit(&prog->prstatus_vector);
+		else
+			drgn_prstatus_map_deinit(&prog->prstatus_map);
+	}
 	free(prog->pgtable_it);
 
 	drgn_object_index_deinit(&prog->oindex);
@@ -736,41 +741,64 @@ drgn_program_load_debug_info(struct drgn_program *prog, const char **paths,
 struct drgn_error *drgn_program_cache_prstatus_entry(struct drgn_program *prog,
 						     char *data, size_t size)
 {
-	struct drgn_prstatus_map_entry entry;
-	size_t pr_pid_offset;
-	uint32_t pr_pid;
+	if (prog->flags & DRGN_PROGRAM_IS_LINUX_KERNEL) {
+		struct string *entry;
 
-	pr_pid_offset = drgn_program_is_64_bit(prog) ? 32 : 24;
+		entry = drgn_prstatus_vector_append_entry(&prog->prstatus_vector);
+		if (!entry)
+			return &drgn_enomem;
+		entry->str = data;
+		entry->len = size;
+	} else {
+		struct drgn_prstatus_map_entry entry;
+		size_t pr_pid_offset;
+		uint32_t pr_pid;
 
-	if (size < pr_pid_offset + sizeof(pr_pid))
-		return NULL;
+		pr_pid_offset = drgn_program_is_64_bit(prog) ? 32 : 24;
+		if (size < pr_pid_offset + sizeof(pr_pid))
+			return NULL;
 
-	memcpy(&pr_pid, data + pr_pid_offset, sizeof(pr_pid));
-	if (drgn_program_bswap(prog))
-		pr_pid = bswap_32(pr_pid);
-	if (!pr_pid)
-		return NULL;
+		memcpy(&pr_pid, data + pr_pid_offset, sizeof(pr_pid));
+		if (drgn_program_bswap(prog))
+			pr_pid = bswap_32(pr_pid);
 
-	entry.key = pr_pid;
-	entry.value.str = data;
-	entry.value.len = size;
-	if (drgn_prstatus_map_insert(&prog->prstatus_cache, &entry,
-				     NULL) == -1) {
-		return &drgn_enomem;
+		entry.key = pr_pid;
+		entry.value.str = data;
+		entry.value.len = size;
+		if (drgn_prstatus_map_insert(&prog->prstatus_map, &entry,
+					     NULL) == -1)
+			return &drgn_enomem;
 	}
 	return NULL;
 }
 
 static struct drgn_error *drgn_program_cache_prstatus(struct drgn_program *prog)
 {
+	struct drgn_error *err;
 	size_t phnum, i;
 
+	if (prog->prstatus_cached)
+		return NULL;
+
+	if (prog->flags & DRGN_PROGRAM_IS_LINUX_KERNEL)
+		drgn_prstatus_vector_init(&prog->prstatus_vector);
+	else
+		drgn_prstatus_map_init(&prog->prstatus_map);
+
 #ifdef WITH_LIBKDUMPFILE
-	if (prog->kdump_ctx)
-		return drgn_program_cache_prstatus_kdump(prog);
+	if (prog->kdump_ctx) {
+		err = drgn_program_cache_prstatus_kdump(prog);
+		goto out;
+	}
 #endif
-	if (elf_getphdrnum(prog->core, &phnum) != 0)
-		return drgn_error_libelf();
+	if (!prog->core) {
+		err = NULL;
+		goto out;
+	}
+	if (elf_getphdrnum(prog->core, &phnum) != 0) {
+		err = drgn_error_libelf();
+		goto out;
+	}
 	for (i = 0; i < phnum; i++) {
 		GElf_Phdr phdr_mem, *phdr;
 		Elf_Data *data;
@@ -779,23 +807,26 @@ static struct drgn_error *drgn_program_cache_prstatus(struct drgn_program *prog)
 		size_t name_offset, desc_offset;
 
 		phdr = gelf_getphdr(prog->core, i, &phdr_mem);
-		if (!phdr)
-			return drgn_error_libelf();
+		if (!phdr) {
+			err = drgn_error_libelf();
+			goto out;
+		}
 		if (phdr->p_type != PT_NOTE)
 			continue;
 
 		data = elf_getdata_rawchunk(prog->core, phdr->p_offset,
 					    phdr->p_filesz,
 					    note_header_type(phdr));
-		if (!data)
-			return drgn_error_libelf();
+		if (!data) {
+			err = drgn_error_libelf();
+			goto out;
+		}
 
 		offset = 0;
 		while (offset < data->d_size &&
 		       (offset = gelf_getnote(data, offset, &nhdr, &name_offset,
 					      &desc_offset))) {
 			const char *name;
-			struct drgn_error *err;
 
 			name = (char *)data->d_buf + name_offset;
 			if (strncmp(name, "CORE", nhdr.n_namesz) != 0 ||
@@ -806,26 +837,56 @@ static struct drgn_error *drgn_program_cache_prstatus(struct drgn_program *prog)
 								(char *)data->d_buf + desc_offset,
 								nhdr.n_descsz);
 			if (err)
-				return err;
+				goto out;
 		}
 	}
+
+	err = NULL;
+out:
+	if (err) {
+		if (prog->flags & DRGN_PROGRAM_IS_LINUX_KERNEL)
+			drgn_prstatus_vector_deinit(&prog->prstatus_vector);
+		else
+			drgn_prstatus_map_deinit(&prog->prstatus_map);
+	} else {
+		prog->prstatus_cached = true;
+	}
+	return err;
+}
+
+struct drgn_error *drgn_program_find_prstatus_by_cpu(struct drgn_program *prog,
+						     uint32_t cpu,
+						     struct string *ret)
+{
+	struct drgn_error *err;
+
+	assert(prog->flags & DRGN_PROGRAM_IS_LINUX_KERNEL);
+	err = drgn_program_cache_prstatus(prog);
+	if (err)
+		return err;
+
+	if (cpu < prog->prstatus_vector.size) {
+		*ret = prog->prstatus_vector.data[cpu];
+	} else {
+		ret->str = NULL;
+		ret->len = 0;
+	}
 	return NULL;
 }
 
-struct drgn_error *drgn_program_find_prstatus(struct drgn_program *prog,
-					      uint32_t tid, struct string *ret)
+struct drgn_error *drgn_program_find_prstatus_by_tid(struct drgn_program *prog,
+						     uint32_t tid,
+						     struct string *ret)
 {
 	struct drgn_error *err;
 	struct drgn_prstatus_map_iterator it;
 
-	if (!prog->prstatus_cached) {
-		err = drgn_program_cache_prstatus(prog);
-		if (err)
-			return err;
-		prog->prstatus_cached = true;
-	}
+	assert(!(prog->flags & DRGN_PROGRAM_IS_LINUX_KERNEL));
+	err = drgn_program_cache_prstatus(prog);
+	if (err)
+		return err;
 
-	it = drgn_prstatus_map_search(&prog->prstatus_cache, &tid);
+	it = drgn_prstatus_map_search(&prog->prstatus_map, &tid);
 	if (!it.entry) {
 		ret->str = NULL;
 		ret->len = 0;

libdrgn/program.h

@@ -22,6 +22,7 @@
 #include "object_index.h"
 #include "platform.h"
 #include "type_index.h"
+#include "vector.h"
 
 /**
  * @ingroup Internals
@@ -53,6 +54,7 @@ struct vmcoreinfo {
 	bool pgtable_l5_enabled;
 };
 
+DEFINE_VECTOR_TYPE(drgn_prstatus_vector, struct string)
 DEFINE_HASH_MAP_TYPE(drgn_prstatus_map, uint32_t, struct string)
 
 struct drgn_dwarf_info_cache;
@@ -92,7 +94,16 @@ struct drgn_program {
 	  */
 	pid_t pid;
 	struct drgn_dwarf_info_cache *_dicache;
-	struct drgn_prstatus_map prstatus_cache;
+	union {
+		/*
+		 * For the Linux kernel, PRSTATUS notes indexed by CPU. See @ref
+		 * drgn_architecture_info::linux_kernel_set_initial_registers
+		 * for why we don't use the PID map.
+		 */
+		struct drgn_prstatus_vector prstatus_vector;
+		/* For userspace programs, PRSTATUS notes indexed by PID. */
+		struct drgn_prstatus_map prstatus_map;
+	};
 	/* See @ref drgn_object_stack_trace(). */
 	struct drgn_error *stack_trace_err;
 	/* See @ref drgn_object_stack_trace_next_thread(). */
@@ -172,16 +183,29 @@ static inline bool drgn_program_is_64_bit(struct drgn_program *prog)
 
 struct drgn_error *drgn_program_get_dwfl(struct drgn_program *prog, Dwfl **ret);
 
+/**
+ * Find the @c NT_PRSTATUS note for the given CPU.
+ *
+ * This is only valid for the Linux kernel.
+ *
+ * @param[out] ret Returned note data. If not found, <tt>ret->str</tt> is set to
+ * @c NULL and <tt>ret->len</tt> is set to zero.
+ */
+struct drgn_error *drgn_program_find_prstatus_by_cpu(struct drgn_program *prog,
+						     uint32_t cpu,
+						     struct string *ret);
+
 /**
  * Find the @c NT_PRSTATUS note for the given thread ID.
  *
- * This assumes that <tt>prog->core</tt> is not @c NULL.
+ * This is only valid for userspace programs.
  *
  * @param[out] ret Returned note data. If not found, <tt>ret->str</tt> is set to
  * @c NULL and <tt>ret->len</tt> is set to zero.
  */
-struct drgn_error *drgn_program_find_prstatus(struct drgn_program *prog,
-					      uint32_t tid, struct string *ret);
+struct drgn_error *drgn_program_find_prstatus_by_tid(struct drgn_program *prog,
+						     uint32_t tid,
+						     struct string *ret);
 
 /**
  * Cache the @c NT_PRSTATUS note provided by @p data in @p prog.