From 628f1dc154e8896cb0423296821b24885731af84 Mon Sep 17 00:00:00 2001
From: Stephen Brennan <stephen@brennan.io>
Date: Thu, 17 Aug 2023 00:29:18 -0700
Subject: [PATCH] libdrgn: Add program_from_fd

All the other Program.set_xxx() methods have a corresponding
program_from_xxx(), so it makes sense to keep things consistent.

Signed-off-by: Stephen Brennan <stephen@brennan.io>
---
 _drgn.pyi                |  9 +++++++++
 drgn/__init__.py         |  2 ++
 libdrgn/program.c        | 15 +++++++++++++++
 libdrgn/program.h        |  6 ++++++
 libdrgn/python/drgnpy.h  |  1 +
 libdrgn/python/main.c    |  2 ++
 libdrgn/python/program.c | 20 ++++++++++++++++++++
 7 files changed, 55 insertions(+)

diff --git a/_drgn.pyi b/_drgn.pyi
index e18be9244..233f24750 100644
--- a/_drgn.pyi
+++ b/_drgn.pyi
@@ -907,6 +907,15 @@ def program_from_core_dump(path: Path) -> Program:
     """
     ...
 
+def program_from_fd(path: int) -> Program:
+    """
+    Create a :class:`Program` from a core dump file descriptor. The type of
+    program (e.g., userspace or kernel) is determined automatically.
+
+    :param fd: Core dump file descriptor.
+    """
+    ...
+
 def program_from_kernel() -> Program:
     """
     Create a :class:`Program` from the running operating system kernel. This
diff --git a/drgn/__init__.py b/drgn/__init__.py
index 35ff14fdc..cbab45411 100644
--- a/drgn/__init__.py
+++ b/drgn/__init__.py
@@ -82,6 +82,7 @@
     host_platform,
     offsetof,
     program_from_core_dump,
+    program_from_fd,
     program_from_kernel,
     program_from_pid,
     reinterpret,
@@ -134,6 +135,7 @@
     "host_platform",
     "offsetof",
     "program_from_core_dump",
+    "program_from_fd",
     "program_from_kernel",
     "program_from_pid",
     "reinterpret",
diff --git a/libdrgn/program.c b/libdrgn/program.c
index d0e4c1830..9d21f607e 100644
--- a/libdrgn/program.c
+++ b/libdrgn/program.c
@@ -1508,6 +1508,21 @@ struct drgn_error *drgn_program_init_core_dump(struct drgn_program *prog,
 	return err;
 }
 
+struct drgn_error *drgn_program_init_fd(struct drgn_program *prog, int fd)
+{
+	struct drgn_error *err;
+
+	err = drgn_program_set_fd(prog, fd);
+	if (err)
+		return err;
+	err = drgn_program_load_debug_info(prog, NULL, 0, true, true);
+	if (err && err->code == DRGN_ERROR_MISSING_DEBUG_INFO) {
+		drgn_error_destroy(err);
+		err = NULL;
+	}
+	return err;
+}
+
 struct drgn_error *drgn_program_init_kernel(struct drgn_program *prog)
 {
 	struct drgn_error *err;
diff --git a/libdrgn/program.h b/libdrgn/program.h
index 4b8954c51..3c08a67f1 100644
--- a/libdrgn/program.h
+++ b/libdrgn/program.h
@@ -222,6 +222,12 @@ void drgn_program_set_platform(struct drgn_program *prog,
 struct drgn_error *drgn_program_init_core_dump(struct drgn_program *prog,
 					       const char *path);
 
+/**
+ * Implement @ref drgn_program_from_fd() on an initialized @ref
+ * drgn_program.
+ */
+struct drgn_error *drgn_program_init_fd(struct drgn_program *prog, int fd);
+
 /**
  * Implement @ref drgn_program_from_kernel() on an initialized @ref
  * drgn_program.
diff --git a/libdrgn/python/drgnpy.h b/libdrgn/python/drgnpy.h
index 7a65e1d0d..2c6b3ef46 100644
--- a/libdrgn/python/drgnpy.h
+++ b/libdrgn/python/drgnpy.h
@@ -275,6 +275,7 @@ bool Program_hold_reserve(Program *prog, size_t n);
 int Program_type_arg(Program *prog, PyObject *type_obj, bool can_be_none,
 		     struct drgn_qualified_type *ret);
 Program *program_from_core_dump(PyObject *self, PyObject *args, PyObject *kwds);
+Program *program_from_fd(PyObject *self, PyObject *args, PyObject *kwds);
 Program *program_from_kernel(PyObject *self);
 Program *program_from_pid(PyObject *self, PyObject *args, PyObject *kwds);
 
diff --git a/libdrgn/python/main.c b/libdrgn/python/main.c
index a001e000e..c54140d2c 100644
--- a/libdrgn/python/main.c
+++ b/libdrgn/python/main.c
@@ -117,6 +117,8 @@ static PyMethodDef drgn_methods[] = {
 	 METH_VARARGS | METH_KEYWORDS, drgn_container_of_DOC},
 	{"program_from_core_dump", (PyCFunction)program_from_core_dump,
 	 METH_VARARGS | METH_KEYWORDS, drgn_program_from_core_dump_DOC},
+	{"program_from_fd", (PyCFunction)program_from_fd,
+	 METH_VARARGS | METH_KEYWORDS, drgn_program_from_fd_DOC},
 	{"program_from_kernel", (PyCFunction)program_from_kernel,
 	 METH_NOARGS, drgn_program_from_kernel_DOC},
 	{"program_from_pid", (PyCFunction)program_from_pid,
diff --git a/libdrgn/python/program.c b/libdrgn/python/program.c
index adcc94c3d..de5ecc444 100644
--- a/libdrgn/python/program.c
+++ b/libdrgn/python/program.c
@@ -1329,6 +1329,26 @@ Program *program_from_core_dump(PyObject *self, PyObject *args, PyObject *kwds)
 	return_ptr(prog);
 }
 
+Program *program_from_fd(PyObject *self, PyObject *args, PyObject *kwds)
+{
+	static char *keywords[] = {"fd", NULL};
+	struct drgn_error *err;
+	int fd;
+	if (!PyArg_ParseTupleAndKeywords(args, kwds,
+					 "i:program_from_fd", keywords,
+					 &fd))
+		return NULL;
+
+	_cleanup_pydecref_ Program *prog =
+		(Program *)PyObject_CallObject((PyObject *)&Program_type, NULL);
+	if (!prog)
+		return NULL;
+	err = drgn_program_init_fd(&prog->prog, fd);
+	if (err)
+		return set_drgn_error(err);
+	return_ptr(prog);
+}
+
 Program *program_from_kernel(PyObject *self)
 {
 	struct drgn_error *err;