Add passing test hash_map_struct.py for using structs as hashmap key/val types

Add support for userspace+kernelspace stack trace example using blazesym

pythonbpf/codegen.py

@@ -86,7 +86,7 @@ def processor(source_code, filename, module):
     license_processing(tree, module)
     globals_processing(tree, module)
     structs_sym_tab = structs_proc(tree, module, bpf_chunks)
-    map_sym_tab = maps_proc(tree, module, bpf_chunks)
+    map_sym_tab = maps_proc(tree, module, bpf_chunks, structs_sym_tab)
     func_proc(tree, module, bpf_chunks, map_sym_tab, structs_sym_tab)
 
     globals_list_creation(tree, module)

pythonbpf/maps/map_debug_info.py

@@ -1,22 +1,31 @@
+import logging
+from llvmlite import ir
 from pythonbpf.debuginfo import DebugInfoGenerator
 from .map_types import BPFMapType
 
+logger: logging.Logger = logging.getLogger(__name__)
 
-def create_map_debug_info(module, map_global, map_name, map_params):
+
+def create_map_debug_info(module, map_global, map_name, map_params, structs_sym_tab):
     """Generate debug info metadata for BPF maps HASH and PERF_EVENT_ARRAY"""
     generator = DebugInfoGenerator(module)
-
+    logger.info(f"Creating debug info for map {map_name} with params {map_params}")
     uint_type = generator.get_uint32_type()
-    ulong_type = generator.get_uint64_type()
     array_type = generator.create_array_type(
         uint_type, map_params.get("type", BPFMapType.UNSPEC).value
     )
     type_ptr = generator.create_pointer_type(array_type, 64)
     key_ptr = generator.create_pointer_type(
-        array_type if "key_size" in map_params else ulong_type, 64
+        array_type
+        if "key_size" in map_params
+        else _get_key_val_dbg_type(map_params.get("key"), generator, structs_sym_tab),
+        64,
     )
     value_ptr = generator.create_pointer_type(
-        array_type if "value_size" in map_params else ulong_type, 64
+        array_type
+        if "value_size" in map_params
+        else _get_key_val_dbg_type(map_params.get("value"), generator, structs_sym_tab),
+        64,
     )
 
     elements_arr = []
@@ -64,7 +73,13 @@ def create_map_debug_info(module, map_global, map_name, map_params):
     return global_var
 
 
-def create_ringbuf_debug_info(module, map_global, map_name, map_params):
+# TODO: This should not be exposed outside of the module.
+# Ideally we should expose a single create_map_debug_info function that handles all map types.
+# We can probably use a registry pattern to register different map types and their debug info generators.
+# map_params["type"] will be used to determine which generator to use.
+def create_ringbuf_debug_info(
+    module, map_global, map_name, map_params, structs_sym_tab
+):
     """Generate debug information metadata for BPF RINGBUF map"""
     generator = DebugInfoGenerator(module)
 
@@ -91,3 +106,65 @@ def create_ringbuf_debug_info(module, map_global, map_name, map_params):
     )
     map_global.set_metadata("dbg", global_var)
     return global_var
+
+
+def _get_key_val_dbg_type(name, generator, structs_sym_tab):
+    """Get the debug type for key/value based on type object"""
+
+    if not name:
+        logger.warn("No name provided for key/value type, defaulting to uint64")
+        return generator.get_uint64_type()
+
+    type_obj = structs_sym_tab.get(name)
+    if type_obj:
+        return _get_struct_debug_type(type_obj, generator, structs_sym_tab)
+
+    # Fallback to basic types
+    logger.info(f"No struct named {name}, falling back to basic type")
+
+    # NOTE: Only handling int and long for now
+    if name in ["c_int32", "c_uint32"]:
+        return generator.get_uint32_type()
+
+    # Default fallback for now
+    return generator.get_uint64_type()
+
+
+def _get_struct_debug_type(struct_obj, generator, structs_sym_tab):
+    """Recursively create debug type for struct"""
+    elements_arr = []
+    for fld in struct_obj.fields.keys():
+        fld_type = struct_obj.field_type(fld)
+        if isinstance(fld_type, ir.IntType):
+            if fld_type.width == 32:
+                fld_dbg_type = generator.get_uint32_type()
+            else:
+                # NOTE: Assuming 64-bit for all other int types
+                fld_dbg_type = generator.get_uint64_type()
+        elif isinstance(fld_type, ir.ArrayType):
+            # NOTE: Array types have u8 elements only for now
+            # Debug info generation should fail for other types
+            elem_type = fld_type.element
+            if isinstance(elem_type, ir.IntType) and elem_type.width == 8:
+                char_type = generator.get_uint8_type()
+                fld_dbg_type = generator.create_array_type(char_type, fld_type.count)
+            else:
+                logger.warning(
+                    f"Array element type {str(elem_type)} not supported for debug info, skipping"
+                )
+                continue
+        else:
+            # NOTE: Only handling int and char arrays for now
+            logger.warning(
+                f"Field type {str(fld_type)} not supported for debug info, skipping"
+            )
+            continue
+
+        member = generator.create_struct_member(
+            fld, fld_dbg_type, struct_obj.field_size(fld)
+        )
+        elements_arr.append(member)
+    struct_type = generator.create_struct_type(
+        elements_arr, struct_obj.size, is_distinct=True
+    )
+    return struct_type

pythonbpf/maps/maps_pass.py

@@ -12,13 +12,15 @@ from pythonbpf.expr.vmlinux_registry import VmlinuxHandlerRegistry
 logger: Logger = logging.getLogger(__name__)
 
 
-def maps_proc(tree, module, chunks):
+def maps_proc(tree, module, chunks, structs_sym_tab):
     """Process all functions decorated with @map to find BPF maps"""
     map_sym_tab = {}
     for func_node in chunks:
         if is_map(func_node):
             logger.info(f"Found BPF map: {func_node.name}")
-            map_sym_tab[func_node.name] = process_bpf_map(func_node, module)
+            map_sym_tab[func_node.name] = process_bpf_map(
+                func_node, module, structs_sym_tab
+            )
     return map_sym_tab
 
 
@@ -60,7 +62,8 @@ def _parse_map_params(rval, expected_args=None):
             if i < len(rval.args):
                 arg = rval.args[i]
                 if isinstance(arg, ast.Name):
-                    params[arg_name] = arg.id
+                    result = _get_vmlinux_enum(handler, arg.id)
+                    params[arg_name] = result if result is not None else arg.id
                 elif isinstance(arg, ast.Constant):
                     params[arg_name] = arg.value
 
@@ -68,19 +71,21 @@ def _parse_map_params(rval, expected_args=None):
     for keyword in rval.keywords:
         if isinstance(keyword.value, ast.Name):
             name = keyword.value.id
-            if handler and handler.is_vmlinux_enum(name):
-                result = handler.get_vmlinux_enum_value(name)
-                params[keyword.arg] = result if result is not None else name
-            else:
-                params[keyword.arg] = name
+            result = _get_vmlinux_enum(handler, name)
+            params[keyword.arg] = result if result is not None else name
         elif isinstance(keyword.value, ast.Constant):
             params[keyword.arg] = keyword.value.value
 
     return params
 
 
+def _get_vmlinux_enum(handler, name):
+    if handler and handler.is_vmlinux_enum(name):
+        return handler.get_vmlinux_enum_value(name)
+
+
 @MapProcessorRegistry.register("RingBuffer")
-def process_ringbuf_map(map_name, rval, module):
+def process_ringbuf_map(map_name, rval, module, structs_sym_tab):
     """Process a BPF_RINGBUF map declaration"""
     logger.info(f"Processing Ringbuf: {map_name}")
     map_params = _parse_map_params(rval, expected_args=["max_entries"])
@@ -100,12 +105,14 @@ def process_ringbuf_map(map_name, rval, module):
     logger.info(f"Ringbuf map parameters: {map_params}")
 
     map_global = create_bpf_map(module, map_name, map_params)
-    create_ringbuf_debug_info(module, map_global.sym, map_name, map_params)
+    create_ringbuf_debug_info(
+        module, map_global.sym, map_name, map_params, structs_sym_tab
+    )
     return map_global
 
 
 @MapProcessorRegistry.register("HashMap")
-def process_hash_map(map_name, rval, module):
+def process_hash_map(map_name, rval, module, structs_sym_tab):
     """Process a BPF_HASH map declaration"""
     logger.info(f"Processing HashMap: {map_name}")
     map_params = _parse_map_params(rval, expected_args=["key", "value", "max_entries"])
@@ -114,12 +121,12 @@ def process_hash_map(map_name, rval, module):
     logger.info(f"Map parameters: {map_params}")
     map_global = create_bpf_map(module, map_name, map_params)
     # Generate debug info for BTF
-    create_map_debug_info(module, map_global.sym, map_name, map_params)
+    create_map_debug_info(module, map_global.sym, map_name, map_params, structs_sym_tab)
     return map_global
 
 
 @MapProcessorRegistry.register("PerfEventArray")
-def process_perf_event_map(map_name, rval, module):
+def process_perf_event_map(map_name, rval, module, structs_sym_tab):
     """Process a BPF_PERF_EVENT_ARRAY map declaration"""
     logger.info(f"Processing PerfEventArray: {map_name}")
     map_params = _parse_map_params(rval, expected_args=["key_size", "value_size"])
@@ -132,7 +139,7 @@ def process_perf_event_map(map_name, rval, module):
     return map_global
 
 
-def process_bpf_map(func_node, module):
+def process_bpf_map(func_node, module, structs_sym_tab):
     """Process a BPF map (a function decorated with @map)"""
     map_name = func_node.name
     logger.info(f"Processing BPF map: {map_name}")
@@ -151,7 +158,7 @@ def process_bpf_map(func_node, module):
     if isinstance(rval, ast.Call) and isinstance(rval.func, ast.Name):
         handler = MapProcessorRegistry.get_processor(rval.func.id)
         if handler:
-            return handler(map_name, rval, module)
+            return handler(map_name, rval, module, structs_sym_tab)
         else:
             logger.warning(f"Unknown map type {rval.func.id}, defaulting to HashMap")
             return process_hash_map(map_name, rval, module)

tests/passing_tests/hash_map_struct.py

@@ -0,0 +1,42 @@
+from pythonbpf import bpf, section, struct, bpfglobal, compile, map
+from pythonbpf.maps import HashMap
+from pythonbpf.helper import pid
+from ctypes import c_void_p, c_int64
+
+
+@bpf
+@struct
+class val_type:
+    counter: c_int64
+    shizzle: c_int64
+
+
+@bpf
+@map
+def last() -> HashMap:
+    return HashMap(key=val_type, value=c_int64, max_entries=16)
+
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_clone")
+def hello_world(ctx: c_void_p) -> c_int64:
+    obj = val_type()
+    obj.counter, obj.shizzle = 42, 96
+    t = last.lookup(obj)
+    if t:
+        print(f"Found existing entry: counter={obj.counter}, pid={t}")
+        last.delete(obj)
+        return 0  # type: ignore [return-value]
+    val = pid()
+    last.update(obj, val)
+    print(f"Map updated!, {obj.counter}, {obj.shizzle}, {val}")
+    return 0  # type: ignore [return-value]
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+compile()