mm, slab: Fix test failures on kernels with SLOB

In the Fixes commit, as part of the refactor, the
_identify_kernel_address() function started open-coding the
get_slab_object_info() function. Unfortunately, this resulted in the
function failing when CONFIG_SLOB was enabled, since the AttributeError
for accessing "slab.slab_cache" was not caught.

Open-coding was done to avoid the duplicate bounds check for
in_direct_map. In fact, in_direct_map seems like a useful helper itself,
and by factoring it out of _find_containing_slab() and into its callers,
we can simplify things a bit. Since _find_containing_slab() can handle
SLOB, this fixes the test failure too.

Fixes: 48e0c51 ("helper.common.memory: recognize vmap addresses in identify_address()")

Signed-off-by: Stephen Brennan <[email protected]>

Author: brenns10

drgn/helpers/common/memory.py

@@ -14,26 +14,12 @@
 from typing import Any, Dict, Optional
 
 import drgn
-from drgn import (
-    FaultError,
-    IntegerLike,
-    Object,
-    PlatformFlags,
-    Program,
-    SymbolKind,
-    cast,
-)
+from drgn import FaultError, IntegerLike, Object, PlatformFlags, Program, SymbolKind
 from drgn.helpers.common.format import escape_ascii_string
 from drgn.helpers.common.prog import takes_program_or_default
-from drgn.helpers.linux.mm import (
-    PageSlab,
-    compound_head,
-    find_vmap_area,
-    pfn_to_virt,
-    virt_to_page,
-)
+from drgn.helpers.linux.mm import find_vmap_area, in_direct_map
 from drgn.helpers.linux.pid import for_each_task
-from drgn.helpers.linux.slab import _get_slab_cache_helper, _get_slab_type
+from drgn.helpers.linux.slab import _find_containing_slab, _get_slab_cache_helper
 
 __all__ = (
     "identify_address",
@@ -111,27 +97,17 @@ def _identify_kernel_address(
     prog: Program, addr: int, cache: Optional[Dict[Any, Any]] = None
 ) -> Optional[str]:
     try:
-        direct_map_start = pfn_to_virt(prog["min_low_pfn"]).value_()
-        direct_map_end = (pfn_to_virt(prog["max_low_pfn"]) + prog["PAGE_SIZE"]).value_()
-        in_direct_map = direct_map_start <= addr < direct_map_end
+        direct_map = in_direct_map(prog, addr)
     except NotImplementedError:
         # Virtual address translation isn't implemented for this
         # architecture.
-        in_direct_map = False
-    if in_direct_map:
-        page = virt_to_page(prog, addr)
+        direct_map = False
 
-        try:
-            head_page = compound_head(page)
-            is_slab = PageSlab(head_page)
-        except FaultError:
-            return None
-
-        if is_slab:
-            slab = cast(_get_slab_type(prog), head_page)
-            slab_info = _get_slab_cache_helper(slab.slab_cache).object_info(
-                head_page, slab, addr
-            )
+    if direct_map:
+        result = _find_containing_slab(prog, addr)
+        if result is not None:
+            slab_cache, page, slab = result
+            slab_info = _get_slab_cache_helper(slab_cache).object_info(page, slab, addr)
             if slab_info:
                 cache_name = escape_ascii_string(
                     slab_info.slab_cache.name.string_(), escape_backslash=True

drgn/helpers/linux/mm.py

@@ -1377,3 +1377,15 @@ def totalram_pages(prog: Program) -> int:
         return prog["_totalram_pages"].counter.value_()
     except KeyError:
         return prog["totalram_pages"].value_()
+
+
+@takes_program_or_default
+def in_direct_map(prog: Program, addr: IntegerLike) -> bool:
+    """
+    Return True if an address is within the kernel's direct memory mapping
+    :param addr: address to check
+    """
+    addr = operator.index(addr)
+    start_addr = pfn_to_virt(prog["min_low_pfn"]).value_()
+    end_addr = (pfn_to_virt(prog["max_low_pfn"]) + prog["PAGE_SIZE"]).value_()
+    return start_addr <= addr < end_addr

drgn/helpers/linux/slab.py

@@ -39,8 +39,8 @@
     _get_PageSlab_impl,
     compound_head,
     for_each_page,
+    in_direct_map,
     page_to_virt,
-    pfn_to_virt,
     virt_to_page,
 )
 from drgn.helpers.linux.percpu import per_cpu_ptr
@@ -464,12 +464,6 @@ def slab_cache_for_each_allocated_object(
 def _find_containing_slab(
     prog: Program, addr: int
 ) -> Optional[Tuple[Object, Object, Object]]:
-    start_addr = pfn_to_virt(prog["min_low_pfn"]).value_()
-    end_addr = (pfn_to_virt(prog["max_low_pfn"]) + prog["PAGE_SIZE"]).value_()
-    if addr < start_addr or addr >= end_addr:
-        # Not a directly mapped address
-        return None
-
     page = virt_to_page(prog, addr)
 
     try:
@@ -516,6 +510,8 @@ def slab_object_info(prog: Program, addr: IntegerLike) -> "Optional[SlabObjectIn
         if not.
     """
     addr = operator.index(addr)
+    if not in_direct_map(prog, addr):
+        return None
     result = _find_containing_slab(prog, addr)
     if result is None:
         return None
@@ -568,6 +564,8 @@ def find_containing_slab_cache(prog: Program, addr: IntegerLike) -> Object:
     :return: ``struct kmem_cache *`` containing *addr*, or ``NULL`` if *addr*
         is not from a slab cache.
     """
+    if not in_direct_map(prog, addr):
+        return NULL(prog, "struct kmem_cache *")
     result = _find_containing_slab(prog, operator.index(addr))
     if result is None:
         return NULL(prog, "struct kmem_cache *")