From dbc95bc7d1a9fe5c6474c6f824f42738f230831b Mon Sep 17 00:00:00 2001 From: Stephen Brennan Date: Fri, 1 Mar 2024 16:46:53 -0800 Subject: [PATCH] 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 --- _drgn.pyi | 27 ++++++++++++ libdrgn/python/program.c | 94 ++++++++++++++++++++++++++++++++++++++++ libdrgn/symbol.c | 19 ++++++++ libdrgn/symbol.h | 3 ++ 4 files changed, 143 insertions(+) diff --git a/_drgn.pyi b/_drgn.pyi index e2bf50c10..ca9e4da87 100644 --- a/_drgn.pyi +++ b/_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. diff --git a/libdrgn/python/program.c b/libdrgn/python/program.c index 5992ab781..c2708b207 100644 --- a/libdrgn/python/program.c +++ b/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, diff --git a/libdrgn/symbol.c b/libdrgn/symbol.c index 53ba15545..581fdc265 100644 --- a/libdrgn/symbol.c +++ b/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, diff --git a/libdrgn/symbol.h b/libdrgn/symbol.h index 1c75e64af..b2e880af4 100644 --- a/libdrgn/symbol.h +++ b/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 */