python: Add Program.add_symbol_finder()

Expose the Symbol finder API so that Python code can be used to lookup
additional symbols by name or address.

Signed-off-by: Stephen Brennan <stephen.s.brennan@oracle.com>

_drgn.pyi

@@ -453,6 +453,38 @@ class Program:
             return an :class:`Object` or ``None`` if not found.
         """
         ...
+    def add_symbol_finder(
+        self, fn: Callable[[Optional[str], Optional[int], bool], List[Symbol]]
+    ) -> None:
+        """
+        Register a callback for finding symbols in the program.
+
+        The callback should take three arguments: a search name, a search
+        address, and a boolean flag indicating whether to return the first
+        match. When the flag is True, the callback should return just one
+        :class:`Symbol`. When the flag is False, the callback should return a
+        list of all matching :class:`Symbol`\\ s. Both the name and address
+        arguments are optional. If both are provided, then the result(s) should
+        match both. If neither are provided the finder should return all
+        available symbols.
+
+        Callbacks are called in reverse order of the order they were added. When
+        the boolean flag is set, the search will short-circuit after the first
+        finder returns a result. Otherwise, all callbacks will be called, and
+        all results will be returned.
+
+        .. note::
+
+            There is one case where callback order is not respected: drgn's
+            internal symbol resolution. If drgn has debuginfo already loaded for
+            a particular memory address, it will take a shortcut and directly
+            use the built-in ELF symbol finder first. If this fails, it will
+            continue to call each callback in the order described.
+
+        :param fn: Callable taking address, name, and flag, and returning a list
+            of :class:`Symbol`\\ s.
+        """
+        ...
     def set_core_dump(self, path: Path) -> None:
         """
         Set the program to a core dump.

libdrgn/python/program.c

@@ -354,6 +354,137 @@ static struct drgn_error *py_object_find_fn(const char *name, size_t name_len,
 	return err;
 }
 
+static struct drgn_error *py_symbol_find_fn(const char *name, uint64_t addr,
+					    enum drgn_find_symbol_flags flags,
+					    void *data, union drgn_find_symbol_result *ret)
+{
+	struct drgn_error *err = NULL;
+	PyGILState_STATE gstate;
+	PyObject *name_obj;
+	PyObject *address_obj;
+	PyObject *one_obj;
+
+	gstate = PyGILState_Ensure();
+
+	if (flags & DRGN_FIND_SYM_NAME) {
+		name_obj = PyUnicode_FromString(name);
+		if (!name_obj) {
+			err = drgn_error_from_python();
+			goto out_gstate;
+		}
+	} else {
+		name_obj = Py_None;
+		Py_INCREF(name_obj);
+	}
+
+	if (flags & DRGN_FIND_SYM_ADDR) {
+		address_obj = PyLong_FromUnsignedLong(addr);
+		if (!address_obj) {
+			err = drgn_error_from_python();
+			goto out_name_obj;
+		}
+	} else {
+		address_obj = Py_None;
+		Py_INCREF(address_obj);
+	}
+
+	one_obj = PyBool_FromLong(flags & DRGN_FIND_SYM_ONE);
+
+	PyObject *obj = PyObject_CallFunction(data, "OOO", name_obj, address_obj, one_obj);
+	if (!obj) {
+		err = drgn_error_from_python();
+		goto out_one_obj;
+	}
+	if (!PyList_Check(obj)) {
+		PyErr_SetString(PyExc_TypeError,
+				"symbol find callback must return list");
+		err = drgn_error_from_python();
+		goto out_obj;
+	}
+	if (flags & DRGN_FIND_SYM_ONE) {
+		if (PyList_GET_SIZE(obj) == 0) {
+			ret->symbol = NULL;
+			goto out_obj;
+		} else if (PyList_GET_SIZE(obj) == 1) {
+			PyObject *item = PyList_GET_ITEM(obj, 0);
+			if (!PyObject_TypeCheck(item, &Symbol_type)) {
+				PyErr_SetString(PyExc_ValueError,
+						"symbol find callback elements must be type Symbol");
+				err = drgn_error_from_python();
+				goto out_obj;
+			}
+			struct drgn_symbol *sym = malloc(sizeof(*sym));
+			if (!sym) {
+				err = &drgn_enomem;
+				goto out_obj;
+			}
+			err = drgn_symbol_copy(sym, ((Symbol *)item)->sym);
+			if (err)
+				free(sym);
+			ret->symbol = sym;
+		} else {
+			PyErr_SetString(PyExc_ValueError,
+					"symbol find callback returned multiple elements, but one was requested");
+			err = drgn_error_from_python();
+			goto out_obj;
+		}
+	} else {
+		size_t len = PyList_GET_SIZE(obj);
+		size_t i;
+
+		if (len == 0) {
+			ret->symbol_arr = NULL;
+			goto out_obj;
+		}
+
+		struct drgn_symbol **arr = calloc(len, sizeof(struct drgn_symbol *));
+		if (!arr) {
+			err = &drgn_enomem;
+			goto out_obj;
+		}
+		for (i = 0; i < len; i++) {
+			PyObject *item = PyList_GET_ITEM(obj, i);
+			if (!PyObject_TypeCheck(item, &Symbol_type)) {
+				PyErr_SetString(PyExc_ValueError,
+						"symbol find callback elements must be type Symbol");
+				err = drgn_error_from_python();
+				goto out_free_arr;
+			}
+			struct drgn_symbol *sym = malloc(sizeof(*sym));
+			if (!sym) {
+				err = &drgn_enomem;
+				goto out_free_arr;
+			}
+			err = drgn_symbol_copy(sym, ((Symbol *)item)->sym);
+			if (err) {
+				free(sym);
+				goto out_free_arr;
+			}
+			arr[i] = sym;
+		}
+		ret->symbol_arr = arr;
+		ret->count = len;
+		/* On success, we need to skip the destruction of arr, this is
+		 * kind of gross. */
+		goto out_obj;
+	out_free_arr:
+		for (size_t j = 0; j < i; j++)
+			drgn_symbol_destroy(arr[i]);
+		free(arr);
+	}
+
+out_obj:
+	Py_DECREF(obj);
+out_one_obj:
+	Py_DECREF(one_obj);
+	Py_DECREF(address_obj);
+out_name_obj:
+	Py_DECREF(name_obj);
+out_gstate:
+	PyGILState_Release(gstate);
+	return err;
+}
+
 static PyObject *Program_add_object_finder(Program *self, PyObject *args,
 					   PyObject *kwds)
 {
@@ -386,6 +517,34 @@ static PyObject *Program_add_object_finder(Program *self, PyObject *args,
 	Py_RETURN_NONE;
 }
 
+static PyObject *Program_add_symbol_finder(Program *self, PyObject *args,
+					   PyObject *kwds)
+{
+	static char *keywords[] = {"fn", NULL};
+	struct drgn_error *err;
+	PyObject *fn;
+	int ret;
+
+	if (!PyArg_ParseTupleAndKeywords(args, kwds, "O:add_symbol_finder",
+					 keywords, &fn))
+	    return NULL;
+
+	if (!PyCallable_Check(fn)) {
+		PyErr_SetString(PyExc_TypeError, "fn must be callable");
+		return NULL;
+	}
+
+	ret = Program_hold_object(self, fn);
+	if (ret == -1)
+		return NULL;
+
+	err = drgn_program_add_symbol_finder(&self->prog, py_symbol_find_fn,
+					     fn);
+	if (err)
+		return set_drgn_error(err);
+	Py_RETURN_NONE;
+}
+
 static PyObject *Program_set_core_dump(Program *self, PyObject *args,
 				       PyObject *kwds)
 {
@@ -1005,6 +1164,8 @@ static PyMethodDef Program_methods[] = {
 	 METH_VARARGS | METH_KEYWORDS, drgn_Program_add_type_finder_DOC},
 	{"add_object_finder", (PyCFunction)Program_add_object_finder,
 	 METH_VARARGS | METH_KEYWORDS, drgn_Program_add_object_finder_DOC},
+	{"add_symbol_finder", (PyCFunction)Program_add_symbol_finder,
+	 METH_VARARGS | METH_KEYWORDS, drgn_Program_add_symbol_finder_DOC},
 	{"set_core_dump", (PyCFunction)Program_set_core_dump,
 	 METH_VARARGS | METH_KEYWORDS, drgn_Program_set_core_dump_DOC},
 	{"set_kernel", (PyCFunction)Program_set_kernel, METH_NOARGS,

libdrgn/symbol.c

@@ -46,6 +46,24 @@ void drgn_symbol_from_elf(const char *name, uint64_t address,
 		ret->kind = DRGN_SYMBOL_KIND_UNKNOWN;
 }
 
+struct drgn_error *
+drgn_symbol_copy(struct drgn_symbol *dst, struct drgn_symbol *src)
+{
+	if (src->name_owned) {
+		dst->name = strdup(src->name);
+		if (!dst->name)
+			return &drgn_enomem;
+	} else {
+		dst->name = src->name;
+	}
+	dst->name_owned = src->name_owned;
+	dst->address = src->address;
+	dst->size = src->size;
+	dst->kind = src->kind;
+	dst->binding = src->binding;
+	return NULL;
+}
+
 LIBDRGN_PUBLIC const char *drgn_symbol_name(struct drgn_symbol *sym)
 {
 	return sym->name;

libdrgn/symbol.h

@@ -28,6 +28,9 @@ struct drgn_symbol_finder {
 void drgn_symbol_from_elf(const char *name, uint64_t address,
 			  const GElf_Sym *elf_sym, struct drgn_symbol *ret);
 
+struct drgn_error *
+drgn_symbol_copy(struct drgn_symbol *dst, struct drgn_symbol *src);
+
 DEFINE_VECTOR_TYPE(symbolp_vector, struct drgn_symbol *)
 
 #endif /* DRGN_SYMBOL_H */