drgn.helpers.linux.device: add for_each_registered_{chr,blk}dev()

Author: osandov

drgn/helpers/linux/device.py

@@ -10,13 +10,17 @@
 """
 
 import operator
+from typing import Iterator, Tuple
 
-from drgn import IntegerLike
+from drgn import IntegerLike, Object, Program, cast
+from drgn.helpers.common.prog import takes_program_or_default
 
 __all__ = (
     "MAJOR",
     "MINOR",
     "MKDEV",
+    "for_each_registered_blkdev",
+    "for_each_registered_chrdev",
 )
 
 
@@ -51,3 +55,33 @@ def MKDEV(major: IntegerLike, minor: IntegerLike) -> int:
     :param minor: Device minor ID.
     """
     return (operator.index(major) << _MINORBITS) | operator.index(minor)
+
+
+@takes_program_or_default
+def for_each_registered_chrdev(
+    prog: Program,
+) -> Iterator[Tuple[int, int, bytes, Object]]:
+    cdev_map_probes = prog["cdev_map"].probes
+    for cd in prog["chrdevs"]:
+        while cd := cd.read_():
+            major = cd.major.value_()
+            dev = MKDEV(major, cd.baseminor)
+            cdev = cd.cdev.read_()
+            if not cdev:
+                probe = cdev_map_probes[major].read_()
+                while next := probe.next.read_():
+                    if probe.dev.value_() == dev:
+                        cdev = cast("struct cdev *", probe.data)
+                        break
+                    probe = next
+
+            yield dev, cd.minorct.value_(), cd.name.string_(), cdev
+            cd = cd.next
+
+
+@takes_program_or_default
+def for_each_registered_blkdev(prog: Program) -> Iterator[Tuple[int, bytes]]:
+    for name in prog["major_names"]:
+        while name := name.read_():
+            yield name.major.value_(), name.name.string_()
+            name = name.next

tests/linux_kernel/helpers/test_device.py

@@ -0,0 +1,57 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# SPDX-License-Identifier: LGPL-2.1-or-later
+
+
+from drgn.helpers.linux.device import (
+    MAJOR,
+    for_each_registered_blkdev,
+    for_each_registered_chrdev,
+)
+from tests.linux_kernel import LinuxKernelTestCase
+
+
+class TestDevice(LinuxKernelTestCase):
+    def test_for_each_registered_chrdev(self):
+        expected = []
+        with open("/proc/devices", "rb") as f:
+            ignore = True
+            for line in f:
+                line = line.rstrip(b"\n")
+                if line == b"Character devices:":
+                    ignore = False
+                elif ignore:
+                    pass
+                elif not line:
+                    break
+                else:
+                    tokens = line.split(maxsplit=1)
+                    expected.append((int(tokens[0]), tokens[1]))
+
+        self.assertCountEqual(
+            [
+                (MAJOR(dev), name)
+                for dev, _, name, _ in for_each_registered_chrdev(self.prog)
+            ],
+            expected,
+        )
+
+    def test_for_each_registered_blkdev(self):
+        expected = []
+        with open("/proc/devices", "rb") as f:
+            ignore = True
+            for line in f:
+                line = line.rstrip(b"\n")
+                if line == b"Block devices:":
+                    ignore = False
+                elif ignore:
+                    pass
+                elif not line:
+                    break
+                else:
+                    tokens = line.split(maxsplit=1)
+                    expected.append((int(tokens[0]), tokens[1]))
+
+        self.assertCountEqual(
+            [(major, name) for major, name in for_each_registered_blkdev(self.prog)],
+            expected,
+        )