Add drgn.Program.main_thread()

_drgn.pyi

@@ -308,6 +308,15 @@ class Program:
         :raises LookupError: if no thread has the given thread ID
         """
         ...
+    def main_thread(self) -> Thread:
+        """
+        Get the main thread of the program.
+
+        This is only defined for userspace programs.
+
+        :raises ValueError: if the program is the Linux kernel
+        """
+        ...
     def crashed_thread(self) -> Thread:
         """
         Get the thread that caused the program to crash.

docs/user_guide.rst

@@ -203,14 +203,17 @@ Threads
 ^^^^^^^
 
 The :class:`drgn.Thread` class represents a thread.
-:meth:`drgn.Program.threads()`, :meth:`drgn.Program.thread()`, and
-:meth:`drgn.Program.crashed_thread()` can be used to find threads::
+:meth:`drgn.Program.threads()`, :meth:`drgn.Program.thread()`,
+:meth:`drgn.Program.main_thread()`, and :meth:`drgn.Program.crashed_thread()`
+can be used to find threads::
 
     >>> for thread in prog.threads():
     ...     print(thread.tid)
     ...
     39143
     39144
+    >>> print(prog.main_thread().tid)
+    39143
     >>> print(prog.crashed_thread().tid)
     39144
 

libdrgn/drgn.h.in

@@ -2871,6 +2871,16 @@ struct drgn_error *drgn_program_find_thread(struct drgn_program *prog,
 					    uint32_t tid,
 					    struct drgn_thread **ret);
 
+/**
+ * Get the main program thread.
+ *
+ * @param[out] ret Borrowed thread handle. This is valid for the lifetime of @p
+ * prog. This must NOT be destroyed with @ref drgn_thread_destroy().
+ * @return @c NULL on success, non-@c NULL on error.
+ */
+struct drgn_error *drgn_program_main_thread(struct drgn_program *prog,
+					    struct drgn_thread **ret);
+
 /**
  * Get the thread that caused the program to crash.
  *

libdrgn/program.c

@@ -678,6 +678,31 @@ static struct drgn_error *get_prstatus_pid(struct drgn_program *prog, const char
 	return NULL;
 }
 
+static struct drgn_error *get_prpsinfo_pid(struct drgn_program *prog,
+					   const char *data, size_t size,
+					   uint32_t *ret)
+{
+	bool is_64_bit, bswap;
+	struct drgn_error *err = drgn_program_is_64_bit(prog, &is_64_bit);
+	if (err)
+		return err;
+	err = drgn_program_bswap(prog, &bswap);
+	if (err)
+		return err;
+
+	size_t offset = is_64_bit ? 24 : 12;
+	uint32_t pr_pid;
+	if (size < offset + sizeof(pr_pid)) {
+		return drgn_error_create(DRGN_ERROR_OTHER,
+					 "NT_PRPSINFO is truncated");
+	}
+	memcpy(&pr_pid, data + offset, sizeof(pr_pid));
+	if (bswap)
+		pr_pid = bswap_32(pr_pid);
+	*ret = pr_pid;
+	return NULL;
+}
+
 struct drgn_error *drgn_thread_dup_internal(const struct drgn_thread *thread,
 					    struct drgn_thread *ret)
 {
@@ -766,6 +791,8 @@ drgn_program_cache_core_dump_notes(struct drgn_program *prog)
 	size_t phnum, i;
 	bool found_prstatus = false;
 	uint32_t first_prstatus_tid;
+	bool found_prpsinfo = false;
+	uint32_t prpsinfo_pid;
 
 	if (prog->core_dump_notes_cached)
 		return NULL;
@@ -823,42 +850,59 @@ drgn_program_cache_core_dump_notes(struct drgn_program *prog)
 			const char *name;
 
 			name = (char *)data->d_buf + name_offset;
-			if (strncmp(name, "CORE", nhdr.n_namesz) != 0 ||
-			    nhdr.n_type != NT_PRSTATUS)
+			if (strncmp(name, "CORE", nhdr.n_namesz) != 0)
 				continue;
 
-			uint32_t tid;
-			err = drgn_program_cache_prstatus_entry(prog,
-								(char *)data->d_buf + desc_offset,
-								nhdr.n_descsz,
-								&tid);
-			if (err)
-				goto err;
-			/*
-			 * The first PRSTATUS note is the crashed thread. See
-			 * fs/binfmt_elf.c:fill_note_info in the Linux kernel
-			 * and bfd/elf.c:elfcore_grok_prstatus in BFD.
-			 */
-			if (!found_prstatus) {
-				found_prstatus = true;
-				first_prstatus_tid = tid;
+			if (nhdr.n_type == NT_PRPSINFO) {
+				err = get_prpsinfo_pid(prog,
+						       (char *)data->d_buf + desc_offset,
+						       nhdr.n_descsz,
+						       &prpsinfo_pid);
+				if (err)
+					goto err;
+				found_prpsinfo = true;
+			} else if (nhdr.n_type == NT_PRSTATUS) {
+				uint32_t tid;
+				err = drgn_program_cache_prstatus_entry(prog,
+						(char *)data->d_buf + desc_offset,
+						nhdr.n_descsz,
+						&tid);
+				if (err)
+					goto err;
+				/*
+				 * The first PRSTATUS note is the crashed thread. See
+				 * fs/binfmt_elf.c:fill_note_info in the Linux kernel
+				 * and bfd/elf.c:elfcore_grok_prstatus in BFD.
+				 */
+				if (!found_prstatus) {
+					found_prstatus = true;
+					first_prstatus_tid = tid;
+				}
 			}
 		}
 	}
 
 out:
 	prog->core_dump_notes_cached = true;
-	if (found_prstatus &&
-	    !(prog->flags & DRGN_PROGRAM_IS_LINUX_KERNEL)) {
-		/*
-		 * Now that thread_set won't be modified, look up the crashed
-		 * thread entry.
-		 */
-		struct drgn_thread_set_iterator it =
-			drgn_thread_set_search(&prog->thread_set,
-					       &first_prstatus_tid);
-		assert(it.entry);
-		prog->crashed_thread = it.entry;
+	if (!(prog->flags & DRGN_PROGRAM_IS_LINUX_KERNEL)) {
+		if (found_prpsinfo) {
+			struct drgn_thread_set_iterator it =
+				drgn_thread_set_search(&prog->thread_set,
+						       &prpsinfo_pid);
+			/* If the PID isn't found, then this is NULL. */
+			prog->main_thread = it.entry;
+		}
+		if (found_prstatus) {
+			/*
+			 * Now that thread_set won't be modified, look up the crashed
+			 * thread entry.
+			 */
+			struct drgn_thread_set_iterator it =
+				drgn_thread_set_search(&prog->thread_set,
+						       &first_prstatus_tid);
+			assert(it.entry);
+			prog->crashed_thread = it.entry;
+		}
 	}
 	return NULL;
 
@@ -1057,6 +1101,30 @@ drgn_program_kernel_core_dump_cache_crashed_thread(struct drgn_program *prog)
 	return NULL;
 }
 
+LIBDRGN_PUBLIC struct drgn_error *
+drgn_program_main_thread(struct drgn_program *prog, struct drgn_thread **ret)
+{
+	struct drgn_error *err;
+
+	if (prog->flags & DRGN_PROGRAM_IS_LINUX_KERNEL) {
+		return drgn_error_create(DRGN_ERROR_INVALID_ARGUMENT,
+					 "main thread is not defined for the Linux kernel");
+	}
+	if (prog->flags & DRGN_PROGRAM_IS_LIVE) {
+		return drgn_error_create(DRGN_ERROR_INVALID_ARGUMENT,
+					 "finding threads is not yet supported for live processes");
+	}
+	err = drgn_program_cache_core_dump_notes(prog);
+	if (err)
+		return err;
+	if (!prog->main_thread) {
+		return drgn_error_create(DRGN_ERROR_OTHER,
+					 "main thread not found");
+	}
+	*ret = prog->main_thread;
+	return NULL;
+}
+
 LIBDRGN_PUBLIC struct drgn_error *
 drgn_program_crashed_thread(struct drgn_program *prog, struct drgn_thread **ret)
 {

libdrgn/program.h

@@ -154,6 +154,7 @@ struct drgn_program {
 		/* For userspace programs, threads indexed by PID. */
 		struct drgn_thread_set thread_set;
 	};
+	struct drgn_thread *main_thread;
 	struct drgn_thread *crashed_thread;
 	bool core_dump_notes_cached;
 	bool prefer_orc_unwinder;

libdrgn/python/program.c

@@ -870,6 +870,16 @@ static PyObject *Program_thread(Program *self, PyObject *args, PyObject *kwds)
 	return ret;
 }
 
+static PyObject *Program_main_thread(Program *self)
+{
+	struct drgn_error *err;
+	struct drgn_thread *thread;
+	err = drgn_program_main_thread(&self->prog, &thread);
+	if (err)
+		return set_drgn_error(err);
+	return Thread_wrap(thread);
+}
+
 static PyObject *Program_crashed_thread(Program *self)
 {
 	struct drgn_error *err;
@@ -1028,6 +1038,8 @@ static PyMethodDef Program_methods[] = {
 	 drgn_Program_threads_DOC},
 	{"thread", (PyCFunction)Program_thread,
 	 METH_VARARGS | METH_KEYWORDS, drgn_Program_thread_DOC},
+	{"main_thread", (PyCFunction)Program_main_thread, METH_NOARGS,
+	 drgn_Program_main_thread_DOC},
 	{"crashed_thread", (PyCFunction)Program_crashed_thread, METH_NOARGS,
 	 drgn_Program_crashed_thread_DOC},
 	{"void_type", (PyCFunction)Program_void_type,

tests/helpers/linux/test_threads.py

@@ -37,6 +37,13 @@ def test_thread(self):
         self.assertEqual(thread.tid, pid)
         self.assertEqual(thread.object, find_task(self.prog, pid))
 
+    def test_main_thread(self):
+        self.assertRaisesRegex(
+            ValueError,
+            "main thread is not defined for the Linux kernel",
+            self.prog.main_thread,
+        )
+
     def test_crashed_thread(self):
         self.assertRaisesRegex(
             ValueError,

tests/linux_kernel/vmcore/test_vmcore.py

@@ -24,6 +24,13 @@ def test_thread(self):
         crashed_thread_tid = self.prog.crashed_thread().tid
         self.assertEqual(self.prog.thread(crashed_thread_tid).tid, crashed_thread_tid)
 
+    def test_main_thread(self):
+        self.assertRaisesRegex(
+            ValueError,
+            "main thread is not defined for the Linux kernel",
+            self.prog.main_thread,
+        )
+
     def test_crashed_thread(self):
         crashed_thread = self.prog.crashed_thread()
         self.assertGreater(crashed_thread.tid, 0)

tests/test_thread.py

@@ -28,6 +28,7 @@ class TestCoreDump(TestCase):
         2265425,
     )
 
+    MAIN_TID = 2265413
     CRASHED_TID = 2265419
 
     @classmethod
@@ -61,5 +62,8 @@ def test_thread(self):
             self.assertEqual(self.prog.thread(tid).tid, tid)
         self.assertRaises(LookupError, self.prog.thread, 99)
 
+    def test_main_thread(self):
+        self.assertEqual(self.prog.main_thread().tid, self.MAIN_TID)
+
     def test_crashed_thread(self):
         self.assertEqual(self.prog.crashed_thread().tid, self.CRASHED_TID)