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

@@ -460,6 +460,33 @@ class Program:
             return an :class:`Object` or ``None`` if not found.
         """
         ...
+    def add_symbol_finder(
+        self, fn: Callable[[Optional[str], Optional[int], bool], Sequence[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 'one' indicating whether to return only
+        the single best match. When the 'one' flag is True, the callback should
+        return a list containing at most 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. If no
+        result is found, the return should be an empty list.
+
+        Callbacks are called in reverse order of the order they were added
+        (i.e,, the most recently added callback is called first). When the
+        'one' flag is set, the search will short-circuit after the first
+        finder which returns a result, and subsequent finders will not be
+        called. Otherwise, all callbacks will be called, and all results will be
+        returned.
+
+        :param fn: Callable taking name, address, and 'one' flag, and
+            returning a sequence of :class:`Symbol`\\ s.
+        """
+        ...
     def set_core_dump(self, path: Union[Path, int]) -> None:
         """
         Set the program to a core dump.

libdrgn/python/program.c

@@ -477,6 +477,70 @@ 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,
+					    enum drgn_find_symbol_flags flags,
+					    void *data, struct drgn_symbol_result_builder *builder)
+{
+	PyGILState_guard();
+
+	_cleanup_pydecref_ PyObject *name_obj = NULL;
+	if (flags & DRGN_FIND_SYMBOL_NAME) {
+		name_obj = PyUnicode_FromString(name);
+		if (!name_obj)
+			return drgn_error_from_python();
+	} else {
+		name_obj = Py_None;
+		Py_INCREF(name_obj);
+	}
+
+	_cleanup_pydecref_ PyObject *address_obj = NULL;
+	if (flags & DRGN_FIND_SYMBOL_ADDR) {
+		address_obj = PyLong_FromUnsignedLong(addr);
+		if (!address_obj)
+			return drgn_error_from_python();
+	} else {
+		address_obj = Py_None;
+		Py_INCREF(address_obj);
+	}
+
+	_cleanup_pydecref_ PyObject *one_obj = PyBool_FromLong(flags & DRGN_FIND_SYMBOL_ONE);
+
+	_cleanup_pydecref_ PyObject *tmp = PyObject_CallFunction(data, "OOO", name_obj,
+								 address_obj, one_obj);
+	if (!tmp)
+		return drgn_error_from_python();
+
+	_cleanup_pydecref_ PyObject *obj =
+		PySequence_Fast(tmp, "symbol finder must return a sequence");
+	if (!obj)
+		return drgn_error_from_python();
+
+	size_t len = PySequence_Fast_GET_SIZE(obj);
+	if (len > 1 && (flags & DRGN_FIND_SYMBOL_ONE))  {
+		return drgn_error_create(DRGN_ERROR_INVALID_ARGUMENT,
+					 "symbol finder returned multiple elements, but one was requested");
+	}
+
+	for (size_t i = 0; i < len; i++) {
+		PyObject *item = PySequence_Fast_GET_ITEM(obj, i);
+		if (!PyObject_TypeCheck(item, &Symbol_type))
+			return drgn_error_create(DRGN_ERROR_TYPE,
+						 "symbol finder results must be of type Symbol");
+		_cleanup_free_ struct drgn_symbol *sym = malloc(sizeof(*sym));
+		if (!sym)
+			return &drgn_enomem;
+		struct drgn_error *err = drgn_symbol_copy(sym, ((Symbol *)item)->sym);
+		if (err)
+			return err;
+
+		if (!drgn_symbol_result_builder_add(builder, sym))
+			return &drgn_enomem;
+		sym = NULL; // owned by the builder now
+	}
+
+	return NULL;
+}
+
 static PyObject *Program_add_object_finder(Program *self, PyObject *args,
 					   PyObject *kwds)
 {
@@ -506,6 +570,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)
 {
@@ -1157,6 +1249,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,25 @@ 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_lifetime == DRGN_LIFETIME_STATIC) {
+		dst->name = src->name;
+		dst->name_lifetime = DRGN_LIFETIME_STATIC;
+	} else {
+		dst->name = strdup(src->name);
+		if (!dst->name)
+			return &drgn_enomem;
+		dst->name_lifetime = DRGN_LIFETIME_OWNED;
+	}
+	dst->address = src->address;
+	dst->size = src->size;
+	dst->kind = src->kind;
+	dst->binding = src->binding;
+	return NULL;
+}
+
 LIBDRGN_PUBLIC struct drgn_error *
 drgn_symbol_create(const char *name, uint64_t address, uint64_t size,
 		   enum drgn_symbol_binding binding, enum drgn_symbol_kind kind,

libdrgn/symbol.h

@@ -61,4 +61,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 */