JakeHillion/drgn
1
0
Fork 0
mirror of https://github.com/JakeHillion/drgn.git synced 2024-12-22 09:13:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Rewrite linux helper iterators in C

Browse Source 
In preparation for introducing an API to represent threads, the linux
helper iterators, radix_tree_for_each, idr_for_each, for_each_pid, and
for_each_task have been rewritten in C. This will allow them to be
accessed from libdrgn, which will be necessary for the threads API.

Signed-off-by: Kevin Svetlitski <svetlitski@fb.com>
This commit is contained in:
Kevin Svetlitski 2021-11-18 16:46:59 -08:00 committed by Omar Sandoval
parent 0f68cd44e2
commit 2b47583c73
10 changed files with 937 additions and 207 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
_drgn.pyi
View File

@ -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.
"""
...

24
drgn/helpers/linux/idr.py
View File

@ -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

53
drgn/helpers/linux/pid.py
View File

@ -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

50
drgn/helpers/linux/radixtree.py
View File

@ -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)

98
libdrgn/helpers.h
View File

@ -15,6 +15,8 @@
#include <stddef.h>
#include <stdint.h>
#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 */

576
libdrgn/linux_kernel_helpers.c
View File

@ -6,10 +6,13 @@
#include <inttypes.h>
#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);
}

21
libdrgn/python/drgnpy.h
View File

@ -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 */

248
libdrgn/python/helpers.c
View File

@ -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;
}

13
libdrgn/python/module.c
View File

@ -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) ||

23
tests/helpers/linux/test_pid.py
View File

@ -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