diff --git a/_drgn.pyi b/_drgn.pyi index 487e11b2..86e356b8 100644 --- a/_drgn.pyi +++ b/_drgn.pyi @@ -2219,3 +2219,41 @@ def _linux_helper_kaslr_offset(prog: Program) -> int: def _linux_helper_pgtable_l5_enabled(prog: Program) -> bool: """Return whether 5-level paging is enabled.""" ... + +def _linux_helper_radix_tree_for_each(root: Object) -> Iterator[Tuple[int, Object]]: + """ + Iterate over all of the entries in a radix tree. + + :param root: ``struct radix_tree_root *`` + :return: Iterator of (index, ``void *``) tuples. + """ + ... + +def _linux_helper_idr_for_each(idr: Object) -> Iterator[Tuple[int, Object]]: + """ + Iterate over all of the entries in an IDR. + + :param idr: ``struct idr *`` + :return: Iterator of (index, ``void *``) tuples. + """ + ... + +def _linux_helper_for_each_pid(prog_or_ns: Union[Program, Object]) -> Iterator[Object]: + """ + Iterate over all PIDs in a namespace. + + :param prog_or_ns: ``struct pid_namespace *`` to iterate over, or + :class:`Program` to iterate over initial PID namespace. + :return: Iterator of ``struct pid *`` objects. + """ + ... + +def _linux_helper_for_each_task(prog_or_ns: Union[Program, Object]) -> Iterator[Object]: + """ + Iterate over all of the tasks visible in a namespace. + + :param prog_or_ns: ``struct pid_namespace *`` to iterate over, or + :class:`Program` to iterate over initial PID namespace. + :return: Iterator of ``struct task_struct *`` objects. + """ + ... diff --git a/drgn/helpers/linux/idr.py b/drgn/helpers/linux/idr.py index c1d1ba2d..7535fbb6 100644 --- a/drgn/helpers/linux/idr.py +++ b/drgn/helpers/linux/idr.py @@ -11,28 +11,12 @@ an ID to a pointer. This currently only supports Linux v4.11+; before this, IDRs were not based on radix trees. """ -from typing import Iterator, Tuple - -from _drgn import _linux_helper_idr_find as idr_find -from drgn import Object -from drgn.helpers.linux.radixtree import radix_tree_for_each +from _drgn import ( + _linux_helper_idr_find as idr_find, + _linux_helper_idr_for_each as idr_for_each, +) __all__ = ( "idr_find", "idr_for_each", ) - - -def idr_for_each(idr: Object) -> Iterator[Tuple[int, Object]]: - """ - Iterate over all of the entries in an IDR. - - :param idr: ``struct idr *`` - :return: Iterator of (index, ``void *``) tuples. - """ - try: - base = idr.idr_base.value_() - except AttributeError: - base = 0 - for index, entry in radix_tree_for_each(idr.idr_rt.address_of_()): - yield index + base, entry diff --git a/drgn/helpers/linux/pid.py b/drgn/helpers/linux/pid.py index ca14ad35..152d15b2 100644 --- a/drgn/helpers/linux/pid.py +++ b/drgn/helpers/linux/pid.py @@ -9,16 +9,13 @@ The ``drgn.helpers.linux.pid`` module provides helpers for looking up process IDs and processes. """ -from typing import Iterator, Union - from _drgn import ( _linux_helper_find_pid as find_pid, _linux_helper_find_task as find_task, + _linux_helper_for_each_pid as for_each_pid, + _linux_helper_for_each_task as for_each_task, _linux_helper_pid_task as pid_task, ) -from drgn import Object, Program, cast, container_of -from drgn.helpers.linux.idr import idr_for_each -from drgn.helpers.linux.list import hlist_for_each_entry __all__ = ( "find_pid", @@ -27,49 +24,3 @@ __all__ = ( "for_each_task", "pid_task", ) - - -def for_each_pid(prog_or_ns: Union[Program, Object]) -> Iterator[Object]: - """ - Iterate over all PIDs in a namespace. - - :param prog_or_ns: ``struct pid_namespace *`` to iterate over, or - :class:`Program` to iterate over initial PID namespace. - :return: Iterator of ``struct pid *`` objects. - """ - if isinstance(prog_or_ns, Program): - prog = prog_or_ns - ns = prog_or_ns["init_pid_ns"].address_of_() - else: - prog = prog_or_ns.prog_ - ns = prog_or_ns - if hasattr(ns, "idr"): - for nr, entry in idr_for_each(ns.idr): - yield cast("struct pid *", entry) - else: - pid_hash = prog["pid_hash"] - for i in range(1 << prog["pidhash_shift"].value_()): - for upid in hlist_for_each_entry( - "struct upid", pid_hash[i].address_of_(), "pid_chain" - ): - if upid.ns == ns: - yield container_of(upid, "struct pid", f"numbers[{int(ns.level)}]") - - -def for_each_task(prog_or_ns: Union[Program, Object]) -> Iterator[Object]: - """ - Iterate over all of the tasks visible in a namespace. - - :param prog_or_ns: ``struct pid_namespace *`` to iterate over, or - :class:`Program` to iterate over initial PID namespace. - :return: Iterator of ``struct task_struct *`` objects. - """ - if isinstance(prog_or_ns, Program): - prog = prog_or_ns - else: - prog = prog_or_ns.prog_ - PIDTYPE_PID = prog["PIDTYPE_PID"].value_() - for pid in for_each_pid(prog_or_ns): - task = pid_task(pid, PIDTYPE_PID) - if task: - yield task diff --git a/drgn/helpers/linux/radixtree.py b/drgn/helpers/linux/radixtree.py index 0339b0a0..090835fc 100644 --- a/drgn/helpers/linux/radixtree.py +++ b/drgn/helpers/linux/radixtree.py @@ -9,54 +9,12 @@ The ``drgn.helpers.linux.radixtree`` module provides helpers for working with radix trees from :linux:`include/linux/radix-tree.h`. """ -from typing import Iterator, Tuple - -from _drgn import _linux_helper_radix_tree_lookup as radix_tree_lookup -from drgn import Object, cast +from _drgn import ( + _linux_helper_radix_tree_for_each as radix_tree_for_each, + _linux_helper_radix_tree_lookup as radix_tree_lookup, +) __all__ = ( "radix_tree_for_each", "radix_tree_lookup", ) - -_RADIX_TREE_ENTRY_MASK = 3 - - -def _is_internal_node(node: Object, internal_node: int) -> bool: - return (node.value_() & _RADIX_TREE_ENTRY_MASK) == internal_node - - -def _entry_to_node(node: Object, internal_node: int) -> Object: - return Object(node.prog_, node.type_, value=node.value_() & ~internal_node) - - -def _radix_tree_root_node(root: Object) -> Tuple[Object, int]: - try: - node = root.xa_head - except AttributeError: - return root.rnode.read_(), 1 - else: - return cast("struct xa_node *", node).read_(), 2 - - -def radix_tree_for_each(root: Object) -> Iterator[Tuple[int, Object]]: - """ - Iterate over all of the entries in a radix tree. - - :param root: ``struct radix_tree_root *`` - :return: Iterator of (index, ``void *``) tuples. - """ - node, RADIX_TREE_INTERNAL_NODE = _radix_tree_root_node(root) - - def aux(node: Object, index: int) -> Iterator[Tuple[int, Object]]: - if _is_internal_node(node, RADIX_TREE_INTERNAL_NODE): - parent = _entry_to_node(node, RADIX_TREE_INTERNAL_NODE) - for i, slot in enumerate(parent.slots): - yield from aux( - cast(parent.type_, slot).read_(), - index + (i << parent.shift.value_()), - ) - elif node: - yield index, cast("void *", node) - - yield from aux(node, 0) diff --git a/libdrgn/helpers.h b/libdrgn/helpers.h index 95d3a4eb..72b8a6e2 100644 --- a/libdrgn/helpers.h +++ b/libdrgn/helpers.h @@ -15,6 +15,8 @@ #include #include +#include "drgn.h" +#include "vector.h" struct drgn_object; struct drgn_program; @@ -43,4 +45,100 @@ struct drgn_error *linux_helper_find_task(struct drgn_object *res, const struct drgn_object *ns, uint64_t pid); +/** + * Iterator convention: + * + * For all of the iterators defined below, the convention for each of the + * `*_next` functions is that upon returning, `*ret` will point to space + * allocated inside of `iter`. The caller is free to do what they wish with + * this return value, but should note that it will be overwritten the next time + * the `*_next` function is called. + */ + +DEFINE_VECTOR_TYPE(linux_helper_radix_tree_iter_frame_vector, + struct linux_helper_radix_tree_iter_frame) + +struct linux_helper_radix_tree_iter_entry { + uint64_t index; + struct drgn_object node; +}; + +struct linux_helper_radix_tree_iter { + bool started; + struct drgn_object root; + // Current value to be yielded + struct linux_helper_radix_tree_iter_entry entry; + // We need this for later initialization of `drgn_object`s + struct drgn_program *prog; + // Frames to keep track of generator state + struct linux_helper_radix_tree_iter_frame_vector frames; + // One-time setup values that are persistent + uint64_t RADIX_TREE_INTERNAL_NODE; + uint64_t RADIX_TREE_MAP_MASK; + struct drgn_qualified_type node_type; +}; + +struct drgn_error *linux_helper_radix_tree_iter_init(struct linux_helper_radix_tree_iter *iter, + const struct drgn_object *root); + +void linux_helper_radix_tree_iter_deinit(struct linux_helper_radix_tree_iter *iter); + +struct drgn_error *linux_helper_radix_tree_iter_next(struct linux_helper_radix_tree_iter *iter, + struct linux_helper_radix_tree_iter_entry **ret); + +struct linux_helper_idr_iter { + struct linux_helper_radix_tree_iter iter; + uint64_t base; +}; + +struct drgn_error *linux_helper_idr_iter_init(struct linux_helper_idr_iter *iter, + const struct drgn_object *idr); + +void linux_helper_idr_iter_deinit(struct linux_helper_idr_iter *iter); + +struct drgn_error *linux_helper_idr_iter_next(struct linux_helper_idr_iter *iter, + struct linux_helper_radix_tree_iter_entry **ret); + +struct linux_helper_pid_iter { + bool has_idr; + struct drgn_qualified_type pid_type; + union { + // if has_idr + struct linux_helper_idr_iter iter; + // else + struct { + struct drgn_qualified_type upid_type; + struct drgn_object pid_hash; + struct drgn_object pos; // a `struct hlist_node*` + struct drgn_object ns; + struct drgn_object entry; // Current value of the iterator + size_t index; // Current loop index + char member_specifier[sizeof("numbers[]") + 20]; + // 20 = maximum length of a uint64_t as a string + // Space for the null terminator is included as part of the sizeof on the string literal + }; + }; +}; + +struct drgn_error *linux_helper_pid_iter_init(struct linux_helper_pid_iter *iter, + const struct drgn_object *ns); + +void linux_helper_pid_iter_deinit(struct linux_helper_pid_iter *iter); + +struct drgn_error *linux_helper_pid_iter_next(struct linux_helper_pid_iter *iter, + struct drgn_object **ret); + +struct linux_helper_task_iter { + struct linux_helper_pid_iter iter; + uint64_t PIDTYPE_PID; +}; + +struct drgn_error *linux_helper_task_iter_init(struct linux_helper_task_iter *iter, + const struct drgn_object *ns); + +void linux_helper_task_iter_deinit(struct linux_helper_task_iter *iter); + +struct drgn_error *linux_helper_task_iter_next(struct linux_helper_task_iter *iter, + struct drgn_object **ret); + #endif /* DRGN_HELPERS_H */ diff --git a/libdrgn/linux_kernel_helpers.c b/libdrgn/linux_kernel_helpers.c index 54bdc282..3a2736e9 100644 --- a/libdrgn/linux_kernel_helpers.c +++ b/libdrgn/linux_kernel_helpers.c @@ -6,10 +6,13 @@ #include #include "drgn.h" +#include "helpers.h" #include "minmax.h" #include "platform.h" #include "program.h" +static const uint64_t RADIX_TREE_ENTRY_MASK = 3; + struct drgn_error *linux_helper_read_vm(struct drgn_program *prog, uint64_t pgtable, uint64_t virt_addr, void *buf, size_t count) @@ -99,12 +102,67 @@ struct drgn_error *linux_helper_read_vm(struct drgn_program *prog, return err; } +static struct drgn_error * +radix_tree_init(struct drgn_program *prog, const struct drgn_object *root, + uint64_t *RADIX_TREE_INTERNAL_NODE_ret, + uint64_t *RADIX_TREE_MAP_MASK_ret, + struct drgn_qualified_type *node_type_ret, + struct drgn_object *node_ret) +{ + struct drgn_error *err = + drgn_object_member_dereference(node_ret, root, "xa_head"); + /* node = root->xa_head */ + if (!err) { + err = drgn_program_find_type(prog, "struct xa_node *", NULL, + node_type_ret); + if (err) + return err; + *RADIX_TREE_INTERNAL_NODE_ret = 2; + } else if (err->code == DRGN_ERROR_LOOKUP) { + drgn_error_destroy(err); + /* node = (void *)root.rnode */ + err = drgn_object_member_dereference(node_ret, root, "rnode"); + if (err) + return err; + err = drgn_program_find_type(prog, "void *", NULL, + node_type_ret); + if (err) + return err; + err = drgn_object_cast(node_ret, *node_type_ret, node_ret); + if (err) + return err; + err = drgn_program_find_type(prog, "struct radix_tree_node *", + NULL, node_type_ret); + if (err) + return err; + *RADIX_TREE_INTERNAL_NODE_ret = 1; + } else { + return err; + } + + struct drgn_type_member *member; + uint64_t member_bit_offset; + err = drgn_type_find_member(drgn_type_type(node_type_ret->type).type, + "slots", &member, &member_bit_offset); + if (err) + return err; + struct drgn_qualified_type member_type; + err = drgn_member_type(member, &member_type, NULL); + if (err) + return err; + if (drgn_type_kind(member_type.type) != DRGN_TYPE_ARRAY) + return drgn_error_create( + DRGN_ERROR_TYPE, + "struct radix_tree_node slots member is not an array"); + *RADIX_TREE_MAP_MASK_ret = drgn_type_length(member_type.type) - 1; + return NULL; +} + struct drgn_error * linux_helper_radix_tree_lookup(struct drgn_object *res, const struct drgn_object *root, uint64_t index) { struct drgn_error *err; - static const uint64_t RADIX_TREE_ENTRY_MASK = 3; uint64_t RADIX_TREE_INTERNAL_NODE; uint64_t RADIX_TREE_MAP_MASK; struct drgn_object node, tmp; @@ -112,55 +170,11 @@ linux_helper_radix_tree_lookup(struct drgn_object *res, drgn_object_init(&node, drgn_object_program(res)); drgn_object_init(&tmp, drgn_object_program(res)); - - /* node = root->xa_head */ - err = drgn_object_member_dereference(&node, root, "xa_head"); - if (!err) { - err = drgn_program_find_type(drgn_object_program(res), - "struct xa_node *", NULL, - &node_type); - if (err) - goto out; - RADIX_TREE_INTERNAL_NODE = 2; - } else if (err->code == DRGN_ERROR_LOOKUP) { - drgn_error_destroy(err); - /* node = (void *)root.rnode */ - err = drgn_object_member_dereference(&node, root, "rnode"); - if (err) - goto out; - err = drgn_program_find_type(drgn_object_program(res), "void *", - NULL, &node_type); - if (err) - goto out; - err = drgn_object_cast(&node, node_type, &node); - if (err) - goto out; - err = drgn_program_find_type(drgn_object_program(res), - "struct radix_tree_node *", NULL, - &node_type); - if (err) - goto out; - RADIX_TREE_INTERNAL_NODE = 1; - } else { - goto out; - } - - struct drgn_type_member *member; - uint64_t member_bit_offset; - err = drgn_type_find_member(drgn_type_type(node_type.type).type, - "slots", &member, &member_bit_offset); + err = radix_tree_init(drgn_object_program(root), root, + &RADIX_TREE_INTERNAL_NODE, &RADIX_TREE_MAP_MASK, + &node_type, &node); if (err) goto out; - struct drgn_qualified_type member_type; - err = drgn_member_type(member, &member_type, NULL); - if (err) - goto out; - if (drgn_type_kind(member_type.type) != DRGN_TYPE_ARRAY) { - err = drgn_error_create(DRGN_ERROR_TYPE, - "struct radix_tree_node slots member is not an array"); - goto out; - } - RADIX_TREE_MAP_MASK = drgn_type_length(member_type.type) - 1; for (;;) { uint64_t value; @@ -243,6 +257,36 @@ out: return err; } +static struct drgn_error *pid_hash_init(struct drgn_program *prog, + const struct drgn_object *ns, + struct drgn_qualified_type *upid_type_ret, + uint64_t *pidhash_length_ret, uint64_t *ns_level_ret) +{ + struct drgn_error *err; + struct drgn_object ns_level, pidhash_shift; + drgn_object_init(&ns_level, prog); + drgn_object_init(&pidhash_shift, prog); + err = drgn_program_find_type(prog, "struct upid", NULL, upid_type_ret); + if (err) + goto out; + err = drgn_program_find_object(prog, "pidhash_shift", NULL, DRGN_FIND_OBJECT_ANY, + &pidhash_shift); + if (err) + goto out; + err = drgn_object_read_unsigned(&pidhash_shift, pidhash_length_ret); + if (err) + goto out; + // *pidhash_length_ret = 1 << pidhash_shift + *pidhash_length_ret = *pidhash_length_ret >= 64 ? 0 : UINT64_C(1) << *pidhash_length_ret; + err = drgn_object_member_dereference(&ns_level, ns, "level"); + if (err) + goto out; + err = drgn_object_read_unsigned(&ns_level, ns_level_ret); +out: + drgn_object_deinit(&ns_level); + return err; +} + /* * Before Linux kernel commit 95846ecf9dac ("pid: replace pid bitmap * implementation with IDR API") (in v4.15), (struct pid_namespace).idr does not @@ -257,15 +301,27 @@ find_pid_in_pid_hash(struct drgn_object *res, const struct drgn_object *ns, { struct drgn_error *err; + struct drgn_object node, tmp; + drgn_object_init(&node, drgn_object_program(res)); + drgn_object_init(&tmp, drgn_object_program(res)); + + err = drgn_object_read(&tmp, ns); + if (err) + goto out; + struct drgn_qualified_type upid_type; + uint64_t i, ns_level; + err = pid_hash_init(drgn_object_program(res), &tmp, &upid_type, &i, + &ns_level); + if (err) + goto out; struct drgn_qualified_type pidp_type; - err = drgn_program_find_type(drgn_object_program(res), "struct pid *", - NULL, &pidp_type); + err = drgn_program_find_type(drgn_object_program(res), "struct pid *", NULL, + &pidp_type); if (err) return err; - struct drgn_qualified_type upid_type; - err = drgn_program_find_type(drgn_object_program(res), "struct upid", - NULL, &upid_type); + uint64_t ns_addr; + err = drgn_object_read_unsigned(&tmp, &ns_addr); if (err) return err; @@ -298,40 +354,6 @@ find_pid_in_pid_hash(struct drgn_object *res, const struct drgn_object *ns, if (err) return err; - struct drgn_object node, tmp; - drgn_object_init(&node, drgn_object_program(res)); - drgn_object_init(&tmp, drgn_object_program(res)); - - err = drgn_object_read(&tmp, ns); - if (err) - goto out; - uint64_t ns_addr; - err = drgn_object_read_unsigned(&tmp, &ns_addr); - if (err) - goto out; - union drgn_value ns_level; - err = drgn_object_member_dereference(&tmp, &tmp, "level"); - if (err) - goto out; - err = drgn_object_read_integer(&tmp, &ns_level); - if (err) - goto out; - - /* i = 1 << pidhash_shift */ - err = drgn_program_find_object(drgn_object_program(res), - "pidhash_shift", NULL, - DRGN_FIND_OBJECT_ANY, &tmp); - if (err) - goto out; - union drgn_value pidhash_shift; - err = drgn_object_read_integer(&tmp, &pidhash_shift); - if (err) - goto out; - uint64_t i; - if (pidhash_shift.uvalue >= 64) - i = 0; - else - i = UINT64_C(1) << pidhash_shift.uvalue; while (i--) { /* for (node = pid_hash[i].first; node; node = node->next) */ err = drgn_object_subscript(&node, pid_hash, i); @@ -382,7 +404,7 @@ find_pid_in_pid_hash(struct drgn_object *res, const struct drgn_object *ns, goto next; sprintf(member, "numbers[%" PRIu64 "].pid_chain", - ns_level.uvalue); + ns_level); err = drgn_object_container_of(res, &node, drgn_type_type(pidp_type.type), member); @@ -533,3 +555,381 @@ out: drgn_object_deinit(&pid_obj); return err; } + +struct linux_helper_radix_tree_iter_frame { + struct drgn_object slots; + uint64_t index; + uint64_t shift; + uint64_t next_slot; +}; + +DEFINE_VECTOR_FUNCTIONS(linux_helper_radix_tree_iter_frame_vector) + +struct drgn_error *linux_helper_radix_tree_iter_init(struct linux_helper_radix_tree_iter *iter, + const struct drgn_object *root) +{ + struct drgn_program *prog = drgn_object_program(root); + iter->started = false; + drgn_object_init(&iter->root, prog); + drgn_object_init(&iter->entry.node, prog); + iter->entry.index = 0; + iter->prog = prog; + + struct drgn_error *err = + radix_tree_init(prog, root, &iter->RADIX_TREE_INTERNAL_NODE, + &iter->RADIX_TREE_MAP_MASK, &iter->node_type, &iter->root); + + if (err) { + drgn_object_deinit(&iter->root); + drgn_object_deinit(&iter->entry.node); + return err; + } + + linux_helper_radix_tree_iter_frame_vector_init(&iter->frames); + return NULL; +} + +void linux_helper_radix_tree_iter_deinit(struct linux_helper_radix_tree_iter *iter) +{ + drgn_object_deinit(&iter->root); + drgn_object_deinit(&iter->entry.node); + while (iter->frames.size) { + drgn_object_deinit( + &linux_helper_radix_tree_iter_frame_vector_pop(&iter->frames)->slots); + } + linux_helper_radix_tree_iter_frame_vector_deinit(&iter->frames); +} + +static struct drgn_error *radix_tree_iter_handle_node(struct linux_helper_radix_tree_iter *iter, + struct drgn_object *_node, uint64_t index, + bool *entry_populated_ret) +{ + struct drgn_object *node = &iter->entry.node; + struct drgn_error *err; + uint64_t value; + + err = drgn_object_read(node, _node); + if (err) + return err; + err = drgn_object_read_unsigned(node, &value); + if (err) + return err; + if ((value & RADIX_TREE_ENTRY_MASK) != iter->RADIX_TREE_INTERNAL_NODE) { + // Base-case, node is NOT internal + if (value) { + *entry_populated_ret = true; + iter->entry.index = index; + } + return NULL; + } + + *entry_populated_ret = false; + + // We are dealing with an internal node, and must iterate over its slots + + err = drgn_object_set_unsigned(node, iter->node_type, + value & ~iter->RADIX_TREE_INTERNAL_NODE, 0); + if (err) + return err; + struct linux_helper_radix_tree_iter_frame *frame = + linux_helper_radix_tree_iter_frame_vector_append_entry(&iter->frames); + if (!frame) + return &drgn_enomem; + frame->index = index; + frame->next_slot = 0; + drgn_object_init(&frame->slots, iter->prog); + // We temporarily use `frame->slots` to hold `shift` in order to avoid + // using another `struct drgn_object`. + err = drgn_object_member_dereference(&frame->slots, node, "shift"); + if (err) + goto err_frame; + err = drgn_object_read_unsigned(&frame->slots, &frame->shift); + if (err) + goto err_frame; + // Now `frame->slots` is actually used for `slots`. + err = drgn_object_member_dereference(&frame->slots, node, "slots"); + if (err) + goto err_frame; + return NULL; + +err_frame: + drgn_object_deinit(&frame->slots); + linux_helper_radix_tree_iter_frame_vector_pop(&iter->frames); + return err; +} + +struct drgn_error *linux_helper_radix_tree_iter_next(struct linux_helper_radix_tree_iter *iter, + struct linux_helper_radix_tree_iter_entry **ret) +{ + bool entry_populated = false; + struct drgn_error *err = NULL; + struct drgn_object node; + drgn_object_init(&node, iter->prog); + if (!iter->started) { + iter->started = true; + err = radix_tree_iter_handle_node(iter, &iter->root, 0, &entry_populated); + } + + while (!err && !entry_populated && iter->frames.size) { + struct linux_helper_radix_tree_iter_frame *frame = + &iter->frames.data[iter->frames.size - 1]; + if (frame->next_slot <= iter->RADIX_TREE_MAP_MASK) { + err = drgn_object_subscript(&node, &frame->slots, frame->next_slot); + if (!err) + err = radix_tree_iter_handle_node(iter, &node, + frame->index + (frame->next_slot++ + << frame->shift), + &entry_populated); + } else { + drgn_object_deinit(&frame->slots); + linux_helper_radix_tree_iter_frame_vector_pop(&iter->frames); + } + } + if (!err) + *ret = entry_populated ? &iter->entry : NULL; + drgn_object_deinit(&node); + return err; +} + +struct drgn_error *linux_helper_idr_iter_init(struct linux_helper_idr_iter *iter, + const struct drgn_object *idr) +{ + struct drgn_error *err; + struct drgn_object idr_rt, idr_base; + drgn_object_init(&idr_rt, drgn_object_program(idr)); + drgn_object_init(&idr_base, drgn_object_program(idr)); + + err = drgn_object_member(&idr_base, idr, "idr_base"); + if (!err) { + err = drgn_object_read_unsigned(&idr_base, &iter->base); + if (err) + goto out; + } else if (err->code == DRGN_ERROR_LOOKUP) { + drgn_error_destroy(err); + iter->base = 0; + } else { + goto out; + } + + err = drgn_object_member(&idr_rt, idr, "idr_rt"); + if (err) + goto out; + err = drgn_object_address_of(&idr_rt, &idr_rt); + if (err) + goto out; + err = linux_helper_radix_tree_iter_init(&iter->iter, &idr_rt); +out: + drgn_object_deinit(&idr_rt); + drgn_object_deinit(&idr_base); + return err; +} + +void linux_helper_idr_iter_deinit(struct linux_helper_idr_iter *iter) +{ + linux_helper_radix_tree_iter_deinit(&iter->iter); +} + +struct drgn_error *linux_helper_idr_iter_next(struct linux_helper_idr_iter *iter, + struct linux_helper_radix_tree_iter_entry **ret) +{ + struct drgn_error *err = linux_helper_radix_tree_iter_next(&iter->iter, ret); + if (!err && *ret) + (*ret)->index += iter->base; + return err; +} + +// See `find_pid_in_pid_hash` +static struct drgn_error *pid_iter_init_pid_hash(struct drgn_program *prog, + const struct drgn_object *ns, + struct linux_helper_pid_iter *iter) +{ + struct drgn_error *err; + drgn_object_init(&iter->pid_hash, prog); + drgn_object_init(&iter->pos, prog); + drgn_object_init(&iter->ns, prog); + drgn_object_init(&iter->entry, prog); + err = drgn_program_find_object(prog, "pid_hash", NULL, DRGN_FIND_OBJECT_VARIABLE, + &iter->pid_hash); + if (err) + goto out; + struct drgn_qualified_type void_star_type; + err = drgn_program_find_type(prog, "void *", NULL, &void_star_type); + if (err) + goto out; + err = drgn_object_set_unsigned(&iter->pos, void_star_type, 0, 0); + if (err) + goto out; + err = drgn_object_copy(&iter->ns, ns); + if (err) + goto out; + uint64_t ns_level; + err = pid_hash_init(prog, ns, &iter->upid_type, &iter->index, &ns_level); + if (err) + goto out; + snprintf(iter->member_specifier, sizeof(iter->member_specifier), "numbers[%" PRIu64 "]", + ns_level); + err = drgn_program_find_type(prog, "struct pid", NULL, &iter->pid_type); + if (err) + goto out; + err = drgn_program_find_type(prog, "struct upid", NULL, &iter->upid_type); + if (err) + goto out; +out: + if (err) { + drgn_object_deinit(&iter->pid_hash); + drgn_object_deinit(&iter->pos); + drgn_object_deinit(&iter->ns); + drgn_object_deinit(&iter->entry); + } + return err; +} + +struct drgn_error *linux_helper_pid_iter_init(struct linux_helper_pid_iter *iter, + const struct drgn_object *ns) +{ + struct drgn_program *prog = drgn_object_program(ns); + struct drgn_error *err; + struct drgn_object idr; + drgn_object_init(&idr, prog); + + err = drgn_object_member_dereference(&idr, ns, "idr"); + if (!err) { + iter->has_idr = true; + err = drgn_program_find_type(prog, "struct pid *", NULL, &iter->pid_type); + if (!err) + err = linux_helper_idr_iter_init(&iter->iter, &idr); + } else if (err->code == DRGN_ERROR_LOOKUP) { + iter->has_idr = false; + drgn_error_destroy(err); + err = pid_iter_init_pid_hash(prog, ns, iter); + } + + drgn_object_deinit(&idr); + return err; +} + +void linux_helper_pid_iter_deinit(struct linux_helper_pid_iter *iter) +{ + if (iter->has_idr) { + linux_helper_idr_iter_deinit(&iter->iter); + } else { + drgn_object_deinit(&iter->pid_hash); + drgn_object_deinit(&iter->pos); + drgn_object_deinit(&iter->ns); + drgn_object_deinit(&iter->entry); + } +} + +struct drgn_error *linux_helper_pid_iter_next(struct linux_helper_pid_iter *iter, + struct drgn_object **ret) +{ + if (iter->has_idr) { + struct linux_helper_radix_tree_iter_entry *entry; + struct drgn_error *err = linux_helper_idr_iter_next(&iter->iter, &entry); + if (err) + return err; + if (!entry) { + *ret = NULL; + return NULL; + } + err = drgn_object_cast(&entry->node, iter->pid_type, &entry->node); + if (!err) + *ret = &entry->node; + return err; + } + + struct drgn_error *err = NULL; + struct drgn_object upid, upid_ns; + drgn_object_init(&upid, drgn_object_program(&iter->ns)); + drgn_object_init(&upid_ns, drgn_object_program(&iter->ns)); + + for (;;) { + for (;;) { + bool is_truthy; + err = drgn_object_bool(&iter->pos, &is_truthy); + if (err) + goto out; + if (is_truthy) + break; + if (iter->index == 0) { + *ret = NULL; + goto out; + } + err = drgn_object_subscript(&iter->pos, &iter->pid_hash, --iter->index); + if (err) + goto out; + err = drgn_object_member(&iter->pos, &iter->pos, "first"); + if (err) + goto out; + err = drgn_object_bool(&iter->pos, &is_truthy); + if (err) + goto out; + } + err = drgn_object_container_of(&upid, &iter->pos, iter->upid_type, "pid_chain"); + if (err) + goto out; + err = drgn_object_member_dereference(&iter->pos, &iter->pos, "next"); + if (err) + goto out; + err = drgn_object_member_dereference(&upid_ns, &upid, "ns"); + if (err) + goto out; + int ns_cmp_result; + err = drgn_object_cmp(&upid_ns, &iter->ns, &ns_cmp_result); + if (err) + goto out; + if (ns_cmp_result == 0) { + err = drgn_object_container_of(&iter->entry, &upid, iter->pid_type, + iter->member_specifier); + if (!err) + *ret = &iter->entry; + goto out; + } + } + +out: + drgn_object_deinit(&upid); + drgn_object_deinit(&upid_ns); + return err; +} + +struct drgn_error *linux_helper_task_iter_init(struct linux_helper_task_iter *iter, + const struct drgn_object *ns) +{ + struct drgn_program *prog = drgn_object_program(ns); + struct drgn_error *err = linux_helper_pid_iter_init(&iter->iter, ns); + if (err) + return err; + struct drgn_object PIDTYPE_PID; + drgn_object_init(&PIDTYPE_PID, prog); + err = drgn_program_find_object(prog, "PIDTYPE_PID", NULL, DRGN_FIND_OBJECT_CONSTANT, + &PIDTYPE_PID); + if (!err) + err = drgn_object_read_unsigned(&PIDTYPE_PID, &iter->PIDTYPE_PID); + if (err) + linux_helper_pid_iter_deinit(&iter->iter); + drgn_object_deinit(&PIDTYPE_PID); + return err; +} + +struct drgn_error *linux_helper_task_iter_next(struct linux_helper_task_iter *iter, + struct drgn_object **ret) +{ + struct drgn_error *err; + bool value_is_truthy; + do { + err = linux_helper_pid_iter_next(&iter->iter, ret); + if (err || !*ret) + return err; + err = linux_helper_pid_task(*ret, *ret, iter->PIDTYPE_PID); + if (err) + return err; + err = drgn_object_bool(*ret, &value_is_truthy); + } while (!err && !value_is_truthy); + return err; +} + +void linux_helper_task_iter_deinit(struct linux_helper_task_iter *iter) +{ + linux_helper_pid_iter_deinit(&iter->iter); +} diff --git a/libdrgn/python/drgnpy.h b/libdrgn/python/drgnpy.h index e5aa04bd..e4ffba92 100644 --- a/libdrgn/python/drgnpy.h +++ b/libdrgn/python/drgnpy.h @@ -97,6 +97,14 @@ typedef struct { struct pyobjectp_set objects; } Program; +typedef struct _GenericIterator { + PyObject_HEAD + Program *prog; + void *iter; + PyObject *(*next)(struct _GenericIterator *); + void (*iter_deinit)(void *); +} GenericIterator; + typedef struct { PyObject_HEAD const struct drgn_register *reg; @@ -167,6 +175,7 @@ extern PyObject *TypeKind_class; extern PyTypeObject DrgnObject_type; extern PyTypeObject DrgnType_type; extern PyTypeObject FaultError_type; +extern PyTypeObject GenericIterator_type; extern PyTypeObject Language_type; extern PyTypeObject ObjectIterator_type; extern PyTypeObject Platform_type; @@ -297,5 +306,17 @@ PyObject *drgnpy_linux_helper_kaslr_offset(PyObject *self, PyObject *args, PyObject *kwds); PyObject *drgnpy_linux_helper_pgtable_l5_enabled(PyObject *self, PyObject *args, PyObject *kwds); +GenericIterator *drgnpy_linux_helper_for_each_task(PyObject *self, + PyObject *args, + PyObject *kwds); +GenericIterator *drgnpy_linux_helper_for_each_pid(PyObject *self, + PyObject *args, + PyObject *kwds); +GenericIterator *drgnpy_linux_helper_idr_for_each(PyObject *self, + PyObject *args, + PyObject *kwds); +GenericIterator *drgnpy_linux_helper_radix_tree_for_each(PyObject *self, + PyObject *args, + PyObject *kwds); #endif /* DRGNPY_H */ diff --git a/libdrgn/python/helpers.c b/libdrgn/python/helpers.c index 2f16a20d..d4444f1d 100644 --- a/libdrgn/python/helpers.c +++ b/libdrgn/python/helpers.c @@ -249,3 +249,251 @@ PyObject *drgnpy_linux_helper_pgtable_l5_enabled(PyObject *self, PyObject *args, return PyErr_Format(PyExc_ValueError, "not Linux kernel"); Py_RETURN_BOOL(prog->prog.vmcoreinfo.pgtable_l5_enabled); } + +static void GenericIterator_dealloc(GenericIterator *self) +{ + if (self->iter) { + self->iter_deinit(self->iter); + free(self->iter); + } + Py_XDECREF(self->prog); + Py_TYPE(self)->tp_free((PyObject *)self); +} + +static PyObject *GenericIterator_next(GenericIterator *self) +{ + return self->next(self); +} + +PyTypeObject GenericIterator_type = { + PyVarObject_HEAD_INIT(NULL, 0) + .tp_name = "_drgn._GenericIterator", + .tp_basicsize = sizeof(GenericIterator), + .tp_dealloc = (destructor)GenericIterator_dealloc, + .tp_flags = Py_TPFLAGS_DEFAULT, + .tp_iter = PyObject_SelfIter, + .tp_iternext = (iternextfunc)GenericIterator_next, +}; + +static PyObject *for_each_task_next(GenericIterator *self) +{ + struct drgn_error *err; + struct drgn_object *entry; + err = linux_helper_task_iter_next(self->iter, &entry); + if (err) + return set_drgn_error(err); + if (!entry) + return NULL; + DrgnObject *ret = DrgnObject_alloc(self->prog); + if (!ret) + return NULL; + err = drgn_object_copy(&ret->obj, entry); + if (err) { + Py_DECREF(ret); + return set_drgn_error(err); + } + return (PyObject *)ret; +} + +GenericIterator *drgnpy_linux_helper_for_each_task(PyObject *self, + PyObject *args, + PyObject *kwds) +{ + static char *keywords[] = {"prog_or_ns", NULL}; + struct drgn_error *err = NULL; + struct prog_or_ns_arg prog_or_ns; + + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O&:for_each_task", + keywords, &prog_or_pid_ns_converter, + &prog_or_ns)) + return NULL; + + GenericIterator *iterator = + (GenericIterator *)GenericIterator_type.tp_alloc( + &GenericIterator_type, 0); + if (!iterator) + goto out; + iterator->prog = prog_or_ns.prog; + Py_INCREF(iterator->prog); + iterator->next = for_each_task_next; + iterator->iter_deinit = (void (*)(void *))linux_helper_task_iter_deinit; + iterator->iter = malloc(sizeof(struct linux_helper_task_iter)); + if (!iterator->iter) { + PyErr_NoMemory(); + Py_DECREF(iterator); + iterator = NULL; + goto out; + } + err = linux_helper_task_iter_init(iterator->iter, prog_or_ns.ns); + if (err) { + set_drgn_error(err); + Py_DECREF(iterator); + iterator = NULL; + } +out: + prog_or_ns_cleanup(&prog_or_ns); + return iterator; +} + +static PyObject *for_each_pid_next(GenericIterator *self) +{ + struct drgn_error *err; + struct drgn_object *entry; + err = linux_helper_pid_iter_next(self->iter, &entry); + if (err) + return set_drgn_error(err); + if (!entry) + return NULL; + DrgnObject *ret = DrgnObject_alloc(self->prog); + if (!ret) + return NULL; + err = drgn_object_copy(&ret->obj, entry); + if (err) { + Py_DECREF(ret); + return set_drgn_error(err); + } + return (PyObject *)ret; +} + +GenericIterator * +drgnpy_linux_helper_for_each_pid(PyObject *self, PyObject *args, PyObject *kwds) +{ + static char *keywords[] = {"prog_or_ns", NULL}; + struct drgn_error *err = NULL; + struct prog_or_ns_arg prog_or_ns; + + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O&:for_each_pid", + keywords, &prog_or_pid_ns_converter, + &prog_or_ns)) + return NULL; + + GenericIterator *iterator = + (GenericIterator *)GenericIterator_type.tp_alloc( + &GenericIterator_type, 0); + if (!iterator) + goto out; + iterator->prog = prog_or_ns.prog; + Py_INCREF(iterator->prog); + iterator->next = for_each_pid_next; + iterator->iter_deinit = (void (*)(void *))linux_helper_pid_iter_deinit; + iterator->iter = malloc(sizeof(struct linux_helper_pid_iter)); + if (!iterator->iter) { + PyErr_NoMemory(); + Py_DECREF(iterator); + iterator = NULL; + goto out; + } + err = linux_helper_pid_iter_init(iterator->iter, prog_or_ns.ns); + if (err) { + set_drgn_error(err); + Py_DECREF(iterator); + iterator = NULL; + } +out: + prog_or_ns_cleanup(&prog_or_ns); + return iterator; +} + +static PyObject *idr_iter_entry_wrap(struct linux_helper_radix_tree_iter_entry *entry, + Program *prog) +{ + DrgnObject *node = DrgnObject_alloc(prog); + if (!node) + return NULL; + struct drgn_error *err = drgn_object_copy(&node->obj, &entry->node); + if (err) { + Py_DECREF(node); + return set_drgn_error(err); + } + PyObject *ret = + Py_BuildValue("KO", (unsigned long long)entry->index, node); + Py_DECREF(node); + return ret; +} + +static PyObject *idr_for_each_next(GenericIterator *self) +{ + struct linux_helper_radix_tree_iter_entry *entry; + struct drgn_error *err = linux_helper_idr_iter_next(self->iter, &entry); + if (err) + return set_drgn_error(err); + if (!entry) + return NULL; + return idr_iter_entry_wrap(entry, self->prog); +} + +GenericIterator * +drgnpy_linux_helper_idr_for_each(PyObject *self, PyObject *args, PyObject *kwds) +{ + static char *keywords[] = {"idr", NULL}; + struct drgn_error *err; + DrgnObject *idr; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!:idr_for_each", + keywords, &DrgnObject_type, &idr)) + return NULL; + + GenericIterator *iterator = + (GenericIterator *)GenericIterator_type.tp_alloc( + &GenericIterator_type, 0); + if (!iterator) + return NULL; + iterator->prog = DrgnObject_prog(idr); + Py_INCREF(iterator->prog); + iterator->next = idr_for_each_next; + iterator->iter_deinit = (void (*)(void *))linux_helper_idr_iter_deinit; + iterator->iter = malloc(sizeof(struct linux_helper_idr_iter)); + if (!iterator->iter) { + Py_DECREF(iterator); + return (GenericIterator *)PyErr_NoMemory(); + } + err = linux_helper_idr_iter_init(iterator->iter, &idr->obj); + if (err) { + Py_DECREF(iterator); + return set_drgn_error(err); + } + return iterator; +} + +static PyObject *radix_tree_for_each_next(GenericIterator *self) +{ + struct linux_helper_radix_tree_iter_entry *entry; + struct drgn_error *err = linux_helper_radix_tree_iter_next(self->iter, &entry); + if (err) + return set_drgn_error(err); + if (!entry) + return NULL; + return idr_iter_entry_wrap(entry, self->prog); +} + +GenericIterator *drgnpy_linux_helper_radix_tree_for_each(PyObject *self, + PyObject *args, + PyObject *kwds) +{ + static char *keywords[] = {"root", NULL}; + struct drgn_error *err; + DrgnObject *root; + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!:radix_tree_for_each", + keywords, &DrgnObject_type, &root)) + return NULL; + + GenericIterator *iterator = + (GenericIterator *)GenericIterator_type.tp_alloc( + &GenericIterator_type, 0); + if (!iterator) + return NULL; + iterator->prog = DrgnObject_prog(root); + Py_INCREF(iterator->prog); + iterator->next = radix_tree_for_each_next; + iterator->iter_deinit = (void (*)(void *))linux_helper_radix_tree_iter_deinit; + iterator->iter = malloc(sizeof(struct linux_helper_radix_tree_iter)); + if (!iterator->iter) { + Py_DECREF(iterator); + return (GenericIterator *)PyErr_NoMemory(); + } + err = linux_helper_radix_tree_iter_init(iterator->iter, &root->obj); + if (err) { + Py_DECREF(iterator); + return set_drgn_error(err); + } + return iterator; +} diff --git a/libdrgn/python/module.c b/libdrgn/python/module.c index 7fd1d2c3..79128bf1 100644 --- a/libdrgn/python/module.c +++ b/libdrgn/python/module.c @@ -141,6 +141,18 @@ static PyMethodDef drgn_methods[] = { {"_linux_helper_pgtable_l5_enabled", (PyCFunction)drgnpy_linux_helper_pgtable_l5_enabled, METH_VARARGS | METH_KEYWORDS}, + {"_linux_helper_for_each_task", + (PyCFunction)drgnpy_linux_helper_for_each_task, + METH_VARARGS | METH_KEYWORDS, drgn__linux_helper_for_each_task_DOC}, + {"_linux_helper_for_each_pid", + (PyCFunction)drgnpy_linux_helper_for_each_pid, + METH_VARARGS | METH_KEYWORDS, drgn__linux_helper_for_each_pid_DOC}, + {"_linux_helper_idr_for_each", + (PyCFunction)drgnpy_linux_helper_idr_for_each, + METH_VARARGS | METH_KEYWORDS, drgn__linux_helper_idr_for_each_DOC}, + {"_linux_helper_radix_tree_for_each", + (PyCFunction)drgnpy_linux_helper_radix_tree_for_each, + METH_VARARGS | METH_KEYWORDS, drgn__linux_helper_radix_tree_for_each_DOC}, {}, }; @@ -230,6 +242,7 @@ DRGNPY_PUBLIC PyMODINIT_FUNC PyInit__drgn(void) add_type(m, &StackTrace_type) || add_type(m, &Symbol_type) || add_type(m, &DrgnType_type) || + add_type(m, &GenericIterator_type) || add_type(m, &TypeEnumerator_type) || add_type(m, &TypeMember_type) || add_type(m, &TypeParameter_type) || diff --git a/tests/helpers/linux/test_pid.py b/tests/helpers/linux/test_pid.py index cc2056b5..92664346 100644 --- a/tests/helpers/linux/test_pid.py +++ b/tests/helpers/linux/test_pid.py @@ -1,6 +1,7 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # SPDX-License-Identifier: GPL-3.0-or-later +from multiprocessing import Barrier, Process import os from drgn.helpers.linux.pid import find_pid, find_task, for_each_pid, for_each_task @@ -30,5 +31,23 @@ class TestPid(LinuxHelperTestCase): self.assertEqual(task.comm.string_(), comm) def test_for_each_task(self): - pid = os.getpid() - self.assertTrue(any(task.pid == pid for task in for_each_task(self.prog))) + NUM_PROCS = 12 + barrier = Barrier(NUM_PROCS + 1) + + def proc_func(): + barrier.wait() + + try: + procs = [Process(target=proc_func) for _ in range(NUM_PROCS)] + for proc in procs: + proc.start() + pids = {task.pid.value_() for task in for_each_task(self.prog)} + for proc in procs: + self.assertIn(proc.pid, pids) + self.assertIn(os.getpid(), pids) + barrier.wait() + except: + barrier.abort() + for proc in procs: + proc.terminate() + raise