diff --git a/pythonbpf/helper/bpf_helper_handler.py b/pythonbpf/helper/bpf_helper_handler.py index 6a0cf14..5cf133f 100644 --- a/pythonbpf/helper/bpf_helper_handler.py +++ b/pythonbpf/helper/bpf_helper_handler.py @@ -2,7 +2,7 @@ import ast from llvmlite import ir from pythonbpf.expr_pass import eval_expr from enum import Enum -from .helper_utils import HelperHandlerRegistry, get_or_create_ptr_from_arg, get_flags_val, _handle_fstring_print, _simple_string_print +from .helper_utils import HelperHandlerRegistry, get_or_create_ptr_from_arg, get_flags_val, _handle_fstring_print, _simple_string_print, _get_data_ptr_and_size class BPFHelperID(Enum): @@ -201,50 +201,31 @@ def bpf_perf_event_output_handler(call, map_ptr, module, builder, func, data_arg = call.args[0] ctx_ptr = func.args[0] # First argument to the function is ctx - if isinstance(data_arg, ast.Name): - data_name = data_arg.id - if local_sym_tab and data_name in local_sym_tab: - data_ptr = local_sym_tab[data_name][0] - else: - raise ValueError( - f"Data variable {data_name} not found in local symbol table.") - # Check is data_name is a struct - if local_var_metadata and data_name in local_var_metadata: - data_type = local_var_metadata[data_name] - if data_type in struct_sym_tab: - struct_info = struct_sym_tab[data_type] - size_val = ir.Constant(ir.IntType(64), struct_info.size) - else: - raise ValueError( - f"Struct type {data_type} for variable {data_name} not found in struct symbol table.") - else: - raise ValueError( - f"Metadata for variable {data_name} not found in local variable metadata.") + data_ptr, size_val = _get_data_ptr_and_size(data_arg, local_sym_tab, + struct_sym_tab, + local_var_metadata) - # BPF_F_CURRENT_CPU is -1 in 32 bit - flags_val = ir.Constant(ir.IntType(64), 0xFFFFFFFF) + # BPF_F_CURRENT_CPU is -1 in 32 bit + flags_val = ir.Constant(ir.IntType(64), 0xFFFFFFFF) - map_void_ptr = builder.bitcast(map_ptr, ir.PointerType()) - data_void_ptr = builder.bitcast(data_ptr, ir.PointerType()) - fn_type = ir.FunctionType( - ir.IntType(64), - [ir.PointerType(ir.IntType(8)), ir.PointerType(), ir.IntType(64), - ir.PointerType(), ir.IntType(64)], - var_arg=False - ) - fn_ptr_type = ir.PointerType(fn_type) + map_void_ptr = builder.bitcast(map_ptr, ir.PointerType()) + data_void_ptr = builder.bitcast(data_ptr, ir.PointerType()) + fn_type = ir.FunctionType( + ir.IntType(64), + [ir.PointerType(ir.IntType(8)), ir.PointerType(), ir.IntType(64), + ir.PointerType(), ir.IntType(64)], + var_arg=False + ) + fn_ptr_type = ir.PointerType(fn_type) - # helper id - fn_addr = ir.Constant(ir.IntType( - 64), BPFHelperID.BPF_PERF_EVENT_OUTPUT.value) - fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) + # helper id + fn_addr = ir.Constant(ir.IntType(64), + BPFHelperID.BPF_PERF_EVENT_OUTPUT.value) + fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) - result = builder.call( - fn_ptr, [ctx_ptr, map_void_ptr, flags_val, data_void_ptr, size_val], tail=False) - return result, None - else: - raise NotImplementedError( - "Only simple object names are supported as data in perf event output.") + result = builder.call( + fn_ptr, [ctx_ptr, map_void_ptr, flags_val, data_void_ptr, size_val], tail=False) + return result, None def handle_helper_call(call, module, builder, func, diff --git a/pythonbpf/helper/helper_utils.py b/pythonbpf/helper/helper_utils.py index 1d0fb73..5c9248e 100644 --- a/pythonbpf/helper/helper_utils.py +++ b/pythonbpf/helper/helper_utils.py @@ -168,7 +168,7 @@ def _process_attr_in_fval(attr_node, fmt_parts, exprs, raise ValueError( f"Variable metadata for '{var_name}' not found in local variable metadata") - var_type = local_sym_tab[var_name][1] + var_type = local_var_metadata[var_name] if var_type not in struct_sym_tab: raise ValueError( f"Struct type '{var_type}' for variable '{var_name}' not found in struct symbol table") @@ -251,12 +251,30 @@ def _prepare_expr_args(expr, func, module, builder, return ir.Constant(ir.IntType(64), 0) -def _call_bpf_printk_helper(args, builder): - """Call the BPF_PRINTK helper function with the provided arguments.""" - fn_type = ir.FunctionType( - ir.IntType(64), [ir.PointerType(), ir.IntType(32)], var_arg=True) - fn_ptr_type = ir.PointerType(fn_type) - fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_PRINTK.value) - fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) +def _get_data_ptr_and_size(data_arg, local_sym_tab, struct_sym_tab, + local_var_metadata): + """Extract data pointer and size information for perf event output.""" + if isinstance(data_arg, ast.Name): + data_name = data_arg.id + if local_sym_tab and data_name in local_sym_tab: + data_ptr = local_sym_tab[data_name][0] + else: + raise ValueError( + f"Data variable {data_name} not found in local symbol table.") - return builder.call(fn_ptr, args, tail=True) + # Check if data_name is a struct + if local_var_metadata and data_name in local_var_metadata: + data_type = local_var_metadata[data_name] + if data_type in struct_sym_tab: + struct_info = struct_sym_tab[data_type] + size_val = ir.Constant(ir.IntType(64), struct_info.size) + return data_ptr, size_val + else: + raise ValueError( + f"Struct type {data_type} for variable {data_name} not found in struct symbol table.") + else: + raise ValueError( + f"Metadata for variable {data_name} not found in local variable metadata.") + else: + raise NotImplementedError( + "Only simple object names are supported as data in perf event output.")