Pass down structs_sym_tab to maps_debug_info, allow vmlinux enums to be used in an indexed format for map declaration

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/helper/__init__.py

@@ -12,6 +12,7 @@ from .helpers import (
     smp_processor_id,
     uid,
     skb_store_bytes,
+    get_stack,
     XDP_DROP,
     XDP_PASS,
 )
@@ -83,6 +84,7 @@ __all__ = [
     "smp_processor_id",
     "uid",
     "skb_store_bytes",
+    "get_stack",
     "XDP_DROP",
     "XDP_PASS",
 ]

pythonbpf/helper/bpf_helper_handler.py

@@ -32,6 +32,7 @@ class BPFHelperID(Enum):
     BPF_GET_CURRENT_UID_GID = 15
     BPF_GET_CURRENT_COMM = 16
     BPF_PERF_EVENT_OUTPUT = 25
+    BPF_GET_STACK = 67
     BPF_PROBE_READ_KERNEL_STR = 115
     BPF_RINGBUF_OUTPUT = 130
     BPF_RINGBUF_RESERVE = 131
@@ -805,7 +806,10 @@ def bpf_skb_store_bytes_emitter(
         flags_val = get_flags_val(call.args[3], builder, local_sym_tab)
     else:
         flags_val = 0
-    flags = ir.Constant(ir.IntType(64), flags_val)
+    if isinstance(flags_val, int):
+        flags = ir.Constant(ir.IntType(64), flags_val)
+    else:
+        flags = flags_val
     fn_type = ir.FunctionType(
         ir.IntType(64),
         args_signature,
@@ -921,8 +925,8 @@ def bpf_ringbuf_submit_emitter(
     )
 
     flags_const = get_flags_val(flags_arg, builder, local_sym_tab)
-    if not flags_arg:
+    if isinstance(flags_const, int):
-        flags_const = ir.Constant(ir.IntType(64), 0)
+        flags_const = ir.Constant(ir.IntType(64), flags_const)
 
     map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
     fn_type = ir.FunctionType(
@@ -940,6 +944,60 @@ def bpf_ringbuf_submit_emitter(
     return result, None
 
 
+@HelperHandlerRegistry.register(
+    "get_stack",
+    param_types=[ir.PointerType(ir.IntType(8)), ir.IntType(64)],
+    return_type=ir.IntType(64),
+)
+def bpf_get_stack_emitter(
+    call,
+    map_ptr,
+    module,
+    builder,
+    func,
+    local_sym_tab=None,
+    struct_sym_tab=None,
+    map_sym_tab=None,
+):
+    """
+    Emit LLVM IR for bpf_get_stack helper function call.
+    """
+    if len(call.args) not in (1, 2):
+        raise ValueError(
+            f"get_stack expects atmost two arguments (buf, flags), got {len(call.args)}"
+        )
+    ctx_ptr = func.args[0]  # First argument to the function is ctx
+    buf_arg = call.args[0]
+    flags_arg = call.args[1] if len(call.args) == 2 else None
+    buf_ptr, buf_size = get_buffer_ptr_and_size(
+        buf_arg, builder, local_sym_tab, struct_sym_tab
+    )
+    flags_val = get_flags_val(flags_arg, builder, local_sym_tab)
+    if isinstance(flags_val, int):
+        flags_val = ir.Constant(ir.IntType(64), flags_val)
+
+    buf_void_ptr = builder.bitcast(buf_ptr, ir.PointerType())
+    fn_type = ir.FunctionType(
+        ir.IntType(64),
+        [
+            ir.PointerType(ir.IntType(8)),
+            ir.PointerType(),
+            ir.IntType(64),
+            ir.IntType(64),
+        ],
+        var_arg=False,
+    )
+    fn_ptr_type = ir.PointerType(fn_type)
+    fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_GET_STACK.value)
+    fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
+    result = builder.call(
+        fn_ptr,
+        [ctx_ptr, buf_void_ptr, ir.Constant(ir.IntType(64), buf_size), flags_val],
+        tail=False,
+    )
+    return result, ir.IntType(64)
+
+
 def handle_helper_call(
     call,
     module,

pythonbpf/helper/helpers.py

@@ -52,6 +52,11 @@ def skb_store_bytes(offset, from_buf, size, flags=0):
     return ctypes.c_int64(0)
 
 
+def get_stack(buf, flags=0):
+    """get the current stack trace"""
+    return ctypes.c_int64(0)
+
+
 XDP_ABORTED = ctypes.c_int64(0)
 XDP_DROP = ctypes.c_int64(1)
 XDP_PASS = ctypes.c_int64(2)

pythonbpf/maps/map_debug_info.py

@@ -2,7 +2,7 @@ from pythonbpf.debuginfo import DebugInfoGenerator
 from .map_types import BPFMapType
 
 
-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)
 
@@ -64,7 +64,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)
 

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 = _get_vmlinux_enum(handler, name)
-                result = handler.get_vmlinux_enum_value(name)
+            params[keyword.arg] = result if result is not None else name
-                params[keyword.arg] = result if result is not None else name
-            else:
-                params[keyword.arg] = 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"])
@@ -105,7 +110,7 @@ def process_ringbuf_map(map_name, rval, module):
 
 
 @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"])
@@ -119,7 +124,7 @@ def process_hash_map(map_name, rval, module):
 
 
 @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 +137,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 +156,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)