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 <[email protected]>

_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

@@ -481,6 +481,89 @@ static struct drgn_error *py_object_find_fn(const char *name, size_t name_len,
 	return drgn_object_copy(ret, &((DrgnObject *)obj)->obj);
 }
 
+static struct drgn_error *py_symbol_find_fn(const char *name, uint64_t addr,
+					    struct drgn_module *module,
+					    enum drgn_find_symbol_flags flags,
+					    void *data, struct drgn_symbol_result_builder *builder)
+{
+	struct drgn_error *err = NULL;
+	PyGILState_STATE gstate;
+	_cleanup_pydecref_ PyObject *name_obj = NULL;
+	_cleanup_pydecref_ PyObject *address_obj = NULL;
+	_cleanup_pydecref_ PyObject *one_obj = NULL;
+	_cleanup_pydecref_ PyObject *obj = NULL;
+
+	gstate = PyGILState_Ensure();
+
+	if (flags & DRGN_FIND_SYM_NAME) {
+		name_obj = PyUnicode_FromString(name);
+		if (!name_obj) {
+			err = drgn_error_from_python();
+			goto out;
+		}
+	} 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;
+		}
+	} else {
+		address_obj = Py_None;
+		Py_INCREF(address_obj);
+	}
+
+	one_obj = PyBool_FromLong(flags & DRGN_FIND_SYM_ONE);
+
+	obj = PyObject_CallFunction(data, "OOO", name_obj, address_obj, one_obj);
+	if (!obj) {
+		err = drgn_error_from_python();
+		goto out;
+	}
+	if (!PyList_Check(obj)) {
+		PyErr_SetString(PyExc_TypeError,
+				"symbol find callback must return list");
+		err = drgn_error_from_python();
+		goto out;
+	}
+
+	size_t len = PyList_GET_SIZE(obj);
+	if (len > 1 && (flags & DRGN_FIND_SYM_ONE))  {
+		PyErr_SetString(PyExc_ValueError,
+				"symbol find callback returned multiple elements, but one was requested");
+		err = drgn_error_from_python();
+		goto out;
+	}
+
+	for (size_t 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;
+		}
+		_cleanup_free_ struct drgn_symbol *sym = malloc(sizeof(*sym));
+		if (!sym) {
+			err = &drgn_enomem;
+			goto out;
+		}
+		err = drgn_symbol_copy(sym, ((Symbol *)item)->sym);
+		if (err)
+			goto out;
+
+		err = drgn_symbol_result_builder_add(builder, no_cleanup_ptr(sym));
+	}
+
+out:
+	PyGILState_Release(gstate);
+	return err;
+}
+
 static PyObject *Program_add_object_finder(Program *self, PyObject *args,
 					   PyObject *kwds)
 {
@@ -510,6 +593,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)
 {
@@ -1125,6 +1236,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

@@ -60,4 +60,7 @@ drgn_symbol_result_builder_single(struct drgn_symbol_result_builder *builder);
 void drgn_symbol_result_builder_array(struct drgn_symbol_result_builder *builder,
 				      struct drgn_symbol ***syms_ret, size_t *count_ret);
 
+struct drgn_error *
+drgn_symbol_copy(struct drgn_symbol *dst, struct drgn_symbol *src);
+
 #endif /* DRGN_SYMBOL_H */