Merge remote-tracking branch 'origin/upstreams/develop' into develop

PR URL: https://www.github.com/delphix/drgn/pull/57

contrib/vmmap.py

@@ -9,6 +9,7 @@
 
 from drgn.helpers.linux.device import MAJOR, MINOR
 from drgn.helpers.linux.fs import d_path
+from drgn.helpers.linux.mm import for_each_vma
 from drgn.helpers.linux.pid import find_task
 
 if len(sys.argv) != 2:
@@ -19,21 +20,12 @@
 if not task:
     sys.exit(f"Cannot find task {pid}")
 
-try:
-    vma = task.mm.mmap
-except AttributeError:
-    sys.exit('maple tree VMA mmap is not supported yet (v6.1+)')
-
 FLAGS = ((0x1, "r"), (0x2, "w"), (0x4, "x"))
 PAGE_SHIFT = prog["PAGE_SHIFT"]
 
 print("Start        End          Flgs   Offset Dev   Inode            File path")
 
-# Starting with 763ecb035029f500d7e6d ("mm: remove the vma linked list") (in v6.1),
-# the VMA mmap linked list is replaced with maple tree which is not supported right now:
-# https://github.com/osandov/drgn/issues/261
-
-while vma:
+for vma in for_each_vma(task.mm):
     flags = "".join([v if f & vma.vm_flags else "-" for f, v in FLAGS])
     flags += "s" if vma.vm_flags & 0x8 else "p"
     print(f"{vma.vm_start.value_():0x}-{vma.vm_end.value_():0x} {flags} ",
@@ -53,5 +45,3 @@
         pgoff = 0
 
     print(f"{pgoff:08x} {major:02x}:{minor:02x} {inode:<16} {path}")
-
-    vma = vma.vm_next

drgn/helpers/linux/mapletree.py

@@ -0,0 +1,163 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# SPDX-License-Identifier: LGPL-2.1-or-later
+
+"""
+Maple Trees
+-----------
+
+The ``drgn.helpers.linux.mapletree`` module provides helpers for working with
+maple trees from :linux:`include/linux/maple_tree.h`.
+
+Maple trees were introduced in Linux 6.1.
+"""
+
+import collections
+import operator
+from typing import Iterator, Tuple
+
+from drgn import NULL, IntegerLike, Object, Program, sizeof
+from drgn.helpers.linux.xarray import _XA_ZERO_ENTRY, _xa_is_node, xa_is_zero
+
+__all__ = (
+    "mt_for_each",
+    "mtree_load",
+)
+
+
+def _ulong_max(prog: Program) -> int:
+    return (1 << (8 * sizeof(prog.type("unsigned long")))) - 1
+
+
+# Combination of mte_to_node(), mte_node_type(), ma_data_end(), and
+# ma_is_leaf().
+def _mte_to_node(
+    prog: Program, entry_value: int, max: int
+) -> Tuple[Object, Object, Object, int, bool]:
+    MAPLE_NODE_MASK = 255
+    MAPLE_NODE_TYPE_MASK = 0xF
+    MAPLE_NODE_TYPE_SHIFT = 0x3
+    maple_leaf_64 = 1
+    maple_range_64 = 2
+    maple_arange_64 = 3
+
+    node = Object(prog, "struct maple_node *", entry_value & ~MAPLE_NODE_MASK)
+    type = (entry_value >> MAPLE_NODE_TYPE_SHIFT) & MAPLE_NODE_TYPE_MASK
+    if type == maple_arange_64:
+        m = node.ma64
+        pivots = m.pivot
+        slots = m.slot
+        end = m.meta.end.value_()
+    elif type == maple_range_64 or type == maple_leaf_64:
+        m = node.mr64
+        pivots = m.pivot
+        slots = m.slot
+        pivot = pivots[len(pivots) - 1].value_()
+        if not pivot:
+            end = m.meta.end.value_()
+        elif pivot == max:
+            end = len(pivots) - 1
+        else:
+            end = len(pivots)
+    else:
+        raise NotImplementedError(f"unknown maple_type {type}")
+
+    return node, pivots, slots, end, type < maple_range_64
+
+
+def mtree_load(mt: Object, index: IntegerLike, *, advanced: bool = False) -> Object:
+    """
+    Look up the entry at a given index in a maple tree.
+
+    >>> entry = mtree_load(task.mm.mm_mt.address_of_(), 0x55d65cfaa000)
+    >>> cast("struct vm_area_struct *", entry)
+    *(struct vm_area_struct *)0xffff97ad82bfc930 = {
+        ...
+    }
+
+    :param mt: ``struct maple_tree *``
+    :param index: Entry index.
+    :param advanced: Whether to return nodes only visible to the maple tree
+        advanced API. If ``False``, zero entries (see
+        :func:`~drgn.helpers.linux.xarray.xa_is_zero()`) will be returned as
+        ``NULL``.
+    :return: ``void *`` found entry, or ``NULL`` if not found.
+    """
+    prog = mt.prog_
+    index = operator.index(index)
+    entry = mt.ma_root.read_()
+    entry_value = entry.value_()
+    if _xa_is_node(entry_value):
+        max = _ulong_max(prog)
+        while True:
+            node, pivots, slots, end, leaf = _mte_to_node(prog, entry_value, max)
+
+            for offset in range(end):
+                pivot = pivots[offset].value_()
+                if pivot >= index:
+                    max = pivot
+                    break
+            else:
+                offset = end
+
+            entry_value = slots[offset].value_()
+            if leaf:
+                if not advanced and entry_value == _XA_ZERO_ENTRY:
+                    return NULL(prog, "void *")
+                return Object(prog, "void *", entry_value)
+    elif entry_value and index == 0:
+        return entry
+    else:
+        return NULL(prog, "void *")
+
+
+def mt_for_each(
+    mt: Object, *, advanced: bool = False
+) -> Iterator[Tuple[int, int, Object]]:
+    """
+    Iterate over all of the entries and their ranges in a maple tree.
+
+    >>> for first_index, last_index, entry in mt_for_each(task.mm.mm_mt.address_of_()):
+    ...     print(hex(first_index), hex(last_index), entry)
+    ...
+    0x55d65cfaa000 0x55d65cfaafff (void *)0xffff97ad82bfc930
+    0x55d65cfab000 0x55d65cfabfff (void *)0xffff97ad82bfc0a8
+    0x55d65cfac000 0x55d65cfacfff (void *)0xffff97ad82bfc000
+    0x55d65cfad000 0x55d65cfadfff (void *)0xffff97ad82bfcb28
+    ...
+
+    :param mt: ``struct maple_tree *``
+    :param advanced: Whether to return nodes only visible to the maple tree
+        advanced API. If ``False``, zero entries (see
+        :func:`~drgn.helpers.linux.xarray.xa_is_zero()`) will be skipped.
+    :return: Iterator of (first_index, last_index, ``void *``) tuples. Both
+        indices are inclusive.
+    """
+    entry = mt.ma_root.read_()
+    entry_value = entry.value_()
+    if _xa_is_node(entry_value):
+        prog = mt.prog_
+        queue = collections.deque(((entry_value, 0, _ulong_max(prog)),))
+        while queue:
+            entry_value, min, max = queue.popleft()
+            node, pivots, slots, end, leaf = _mte_to_node(prog, entry_value, max)
+
+            if leaf:
+                prev = min
+                for offset in range(end):
+                    pivot = pivots[offset].value_()
+                    slot = slots[offset].read_()
+                    if slot and (advanced or not xa_is_zero(slot)):
+                        yield (prev, pivot, slot)
+                    prev = pivot + 1
+                slot = slots[end].read_()
+                if slot and (advanced or not xa_is_zero(slot)):
+                    yield (prev, max, slot)
+            else:
+                prev = min
+                for offset in range(end):
+                    pivot = pivots[offset].value_()
+                    queue.append((slots[offset].value_(), prev, pivot))
+                    prev = pivot + 1
+                queue.append((slots[end].value_(), prev, max))
+    elif entry_value:
+        yield (0, 0, entry)

drgn/helpers/linux/mm.py

@@ -20,6 +20,7 @@
 )
 from drgn import IntegerLike, Object, Program, cast
 from drgn.helpers.common.format import decode_enum_type_flags
+from drgn.helpers.linux.mapletree import mt_for_each
 
 __all__ = (
     "PFN_PHYS",
@@ -39,6 +40,7 @@
     "follow_pfn",
     "follow_phys",
     "for_each_page",
+    "for_each_vma",
     "page_size",
     "page_to_pfn",
     "page_to_phys",
@@ -1312,6 +1314,40 @@ def environ(task: Object) -> List[bytes]:
     return access_remote_vm(mm, env_start, env_end - env_start).split(b"\0")[:-1]
 
 
+def for_each_vma(mm: Object) -> Iterator[Object]:
+    """
+    Iterate over every virtual memory area (VMA) in a virtual address space.
+
+    >>> for vma in for_each_vma(task.mm):
+    ...     print(vma)
+    ...
+    *(struct vm_area_struct *)0xffff97ad82bfc930 = {
+        ...
+    }
+    *(struct vm_area_struct *)0xffff97ad82bfc0a8 = {
+        ...
+    }
+    ...
+
+    :param mm: ``struct mm_struct *``
+    :return: Iterator of ``struct vm_area_struct *`` objects.
+    """
+    try:
+        # Since Linux kernel commit 763ecb035029 ("mm: remove the vma linked
+        # list") (in v6.1), VMAs are stored in a maple tree.
+        mt = mm.mm_mt.address_of_()
+    except AttributeError:
+        # Before that, they are in a linked list.
+        vma = mm.mmap
+        while vma:
+            yield vma
+            vma = vma.vm_next
+    else:
+        type = mm.prog_.type("struct vm_area_struct *")
+        for _, _, entry in mt_for_each(mt):
+            yield cast(type, entry)
+
+
 def totalram_pages(prog: Program) -> int:
     """
     Return the total number of RAM memory pages.

drgn/helpers/linux/xarray.py

@@ -46,6 +46,10 @@
 _XA_ZERO_ENTRY = 1030  # xa_mk_internal(257)
 
 
+def _xa_is_node(entry_value: int) -> bool:
+    return (entry_value & 3) == 2 and entry_value > 4096
+
+
 def xa_load(xa: Object, index: IntegerLike, *, advanced: bool = False) -> Object:
     """
     Look up the entry at a given index in an XArray.
@@ -142,12 +146,12 @@ def to_node(entry_value: int) -> Object:
 
         # Return > 0 if xa_is_node(), < 0 if xa_is_sibling(), and 0 otherwise.
         def is_internal(slots: Optional[Object], entry_value: int) -> int:
-            if (entry_value & 3) == 2:
-                if entry_value > 4096:
-                    return 1
-                elif entry_value < 256:
-                    return -1
-            return 0
+            if _xa_is_node(entry_value):
+                return 1
+            elif (entry_value & 3) == 2 and entry_value < 256:
+                return -1
+            else:
+                return 0
 
         # xa_to_node()
         def to_node(entry_value: int) -> Object:

tests/linux_kernel/__init__.py

@@ -301,6 +301,45 @@ def iter_mounts(pid="self"):
 _mlock.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
 
 
+_MAPS_RE = re.compile(
+    rb"(?P<start>[0-9a-f]+)-(?P<end>[0-9a-f]+) (?P<flags>\S+) (?P<offset>[0-9a-f]+) (?P<dev_major>[0-9a-f]+):(?P<dev_minor>[0-9a-f]+) (?P<ino>[0-9]+)\s*(?P<path>.*)"
+)
+
+
+class VmMap(NamedTuple):
+    start: int
+    end: int
+    read: bool
+    write: bool
+    execute: bool
+    shared: bool
+    offset: int
+    dev: int
+    ino: int
+    path: str
+
+
+def iter_maps(pid="self"):
+    with open(f"/proc/{pid}/maps", "rb") as f:
+        for line in f:
+            match = _MAPS_RE.match(line)
+            flags = match["flags"]
+            yield VmMap(
+                start=int(match["start"], 16),
+                end=int(match["end"], 16),
+                read=b"r" in flags,
+                write=b"w" in flags,
+                execute=b"x" in flags,
+                shared=b"s" in flags,
+                offset=int(match["offset"], 16),
+                dev=os.makedev(
+                    int(match["dev_major"], 16), int(match["dev_minor"], 16)
+                ),
+                ino=int(match["ino"]),
+                path=os.fsdecode(match["path"]),
+            )
+
+
 def mlock(addr, len):
     _check_ctypes_syscall(_mlock(addr, len))