TEMPORARY: reintroduce Python implementations of linux helpers so we can test against them

Author: Svetlitski

drgn/helpers/linux/idr.py

@@ -11,12 +11,32 @@
 IDRs were not based on radix trees.
 """
 
+from typing import Iterator, Tuple
+
 from _drgn import (
     _linux_helper_idr_find as idr_find,
     _linux_helper_idr_for_each as idr_for_each,
 )
+from drgn import Object
+from drgn.helpers.linux.radixtree import old_radix_tree_for_each
 
 __all__ = (
     "idr_find",
     "idr_for_each",
+    "old_idr_for_each",
 )
+
+
+def old_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 old_radix_tree_for_each(idr.idr_rt.address_of_()):
+        yield index + base, entry

drgn/helpers/linux/pid.py

@@ -9,18 +9,71 @@
 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_pid_task as pid_task,
     _linux_helper_for_each_task as for_each_task,
     _linux_helper_for_each_pid as for_each_pid,
 )
+from drgn import Object, Program, cast, container_of
+from drgn.helpers.linux.idr import old_idr_for_each
+from drgn.helpers.linux.list import hlist_for_each_entry
 
 __all__ = (
     "find_pid",
     "find_task",
     "for_each_pid",
     "for_each_task",
     "pid_task",
+    "old_for_each_pid",
+    "old_for_each_task",
 )
+
+
+def old_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 old_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 old_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 old_for_each_pid(prog_or_ns):
+        task = pid_task(pid, PIDTYPE_PID)
+        if task:
+            yield task

drgn/helpers/linux/radixtree.py

@@ -9,9 +9,58 @@
 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,
     _linux_helper_radix_tree_for_each as radix_tree_for_each,
 )
+from drgn import Object, cast
+
+__all__ = (
+    "radix_tree_for_each",
+    "radix_tree_lookup",
+    "old_radix_tree_for_each",
+)
+
+_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 old_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)
 
-__all__ = ("radix_tree_for_each", "radix_tree_lookup")
+    yield from aux(node, 0)

tests/helpers/linux/test_c_helpers.py

@@ -0,0 +1,36 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from drgn.helpers.linux.pid import (
+    for_each_pid,
+    for_each_task,
+    old_for_each_pid,
+    old_for_each_task,
+)
+from drgn.helpers.linux.idr import idr_for_each, old_idr_for_each
+from drgn.helpers.linux.radixtree import radix_tree_for_each, old_radix_tree_for_each
+from tests.helpers.linux import LinuxHelperTestCase
+
+
+class TestCHelpers(LinuxHelperTestCase):
+    def verify_iterators_match(self, old_iter, other_iter, arg):
+        self.assertTrue(sum(1 for _ in old_iter(arg)) >= 100)
+        self.assertTrue(all(x == y for x, y in zip(old_iter(arg), other_iter(arg))))
+
+    def test_for_each_task(self):
+        self.verify_iterators_match(old_for_each_task, for_each_task, self.prog)
+
+    def test_for_each_pid(self):
+        self.verify_iterators_match(old_for_each_pid, for_each_pid, self.prog)
+
+    def test_idr_for_each(self):
+        if hasattr(self.prog["init_pid_ns"], "idr"):
+            idr = self.prog["init_pid_ns"].idr
+            self.verify_iterators_match(old_idr_for_each, idr_for_each, idr)
+
+    def test_radix_tree_for_each(self):
+        if hasattr(self.prog["init_pid_ns"], "idr"):
+            root = self.prog["init_pid_ns"].idr.idr_rt.address_of_()
+            self.verify_iterators_match(
+                old_radix_tree_for_each, radix_tree_for_each, root
+            )