diff --git a/examples/struct_and_perf.py b/examples/struct_and_perf.py index 2cfe233..a94a198 100644 --- a/examples/struct_and_perf.py +++ b/examples/struct_and_perf.py @@ -27,10 +27,7 @@ def hello(ctx: c_void_p) -> c_int32: dataobj.pid = pid() dataobj.ts = ktime() # dataobj.comm = strobj - print( - f"clone called at {dataobj.ts} by pid" - f"{dataobj.pid}, comm {strobj}" - ) + print(f"clone called at {dataobj.ts} by pid" f"{dataobj.pid}, comm {strobj}") events.output(dataobj) return c_int32(0) diff --git a/pythonbpf/functions_pass.py b/pythonbpf/functions_pass.py index 3e0fbf4..bbfafa0 100644 --- a/pythonbpf/functions_pass.py +++ b/pythonbpf/functions_pass.py @@ -83,19 +83,16 @@ def handle_assign( elif isinstance(rval, ast.Constant): if isinstance(rval.value, bool): if rval.value: - builder.store(ir.Constant(ir.IntType(1), 1), - local_sym_tab[var_name][0]) + builder.store(ir.Constant(ir.IntType(1), 1), local_sym_tab[var_name][0]) else: - builder.store(ir.Constant(ir.IntType(1), 0), - local_sym_tab[var_name][0]) + builder.store(ir.Constant(ir.IntType(1), 0), local_sym_tab[var_name][0]) print(f"Assigned constant {rval.value} to {var_name}") elif isinstance(rval.value, int): # Assume c_int64 for now # var = builder.alloca(ir.IntType(64), name=var_name) # var.align = 8 builder.store( - ir.Constant(ir.IntType(64), - rval.value), local_sym_tab[var_name][0] + ir.Constant(ir.IntType(64), rval.value), local_sym_tab[var_name][0] ) # local_sym_tab[var_name] = var print(f"Assigned constant {rval.value} to {var_name}") @@ -110,8 +107,7 @@ def handle_assign( global_str.linkage = "internal" global_str.global_constant = True global_str.initializer = str_const - str_ptr = builder.bitcast( - global_str, ir.PointerType(ir.IntType(8))) + str_ptr = builder.bitcast(global_str, ir.PointerType(ir.IntType(8))) builder.store(str_ptr, local_sym_tab[var_name][0]) print(f"Assigned string constant '{rval.value}' to {var_name}") else: @@ -130,8 +126,7 @@ def handle_assign( # var = builder.alloca(ir_type, name=var_name) # var.align = ir_type.width // 8 builder.store( - ir.Constant( - ir_type, rval.args[0].value), local_sym_tab[var_name][0] + ir.Constant(ir_type, rval.args[0].value), local_sym_tab[var_name][0] ) print( f"Assigned {call_type} constant " @@ -177,8 +172,7 @@ def handle_assign( ir_type = struct_info.ir_type # var = builder.alloca(ir_type, name=var_name) # Null init - builder.store(ir.Constant(ir_type, None), - local_sym_tab[var_name][0]) + builder.store(ir.Constant(ir_type, None), local_sym_tab[var_name][0]) local_var_metadata[var_name] = call_type print(f"Assigned struct {call_type} to {var_name}") # local_sym_tab[var_name] = var @@ -249,8 +243,7 @@ def handle_cond(func, module, builder, cond, local_sym_tab, map_sym_tab): print(f"Undefined variable {cond.id} in condition") return None elif isinstance(cond, ast.Compare): - lhs = eval_expr(func, module, builder, cond.left, - local_sym_tab, map_sym_tab)[0] + lhs = eval_expr(func, module, builder, cond.left, local_sym_tab, map_sym_tab)[0] if len(cond.ops) != 1 or len(cond.comparators) != 1: print("Unsupported complex comparison") return None @@ -303,8 +296,7 @@ def handle_if( else: else_block = None - cond = handle_cond(func, module, builder, stmt.test, - local_sym_tab, map_sym_tab) + cond = handle_cond(func, module, builder, stmt.test, local_sym_tab, map_sym_tab) if else_block: builder.cbranch(cond, then_block, else_block) else: @@ -449,8 +441,7 @@ def allocate_mem( ir_type = ctypes_to_ir(call_type) var = builder.alloca(ir_type, name=var_name) var.align = ir_type.width // 8 - print( - f"Pre-allocated variable {var_name} of type {call_type}") + print(f"Pre-allocated variable {var_name} of type {call_type}") elif HelperHandlerRegistry.has_handler(call_type): # Assume return type is int64 for now ir_type = ir.IntType(64) @@ -671,8 +662,7 @@ def infer_return_type(func_node: ast.FunctionDef): if found_type is None: found_type = t elif found_type != t: - raise ValueError("Conflicting return types:" - f"{found_type} vs {t}") + raise ValueError("Conflicting return types:" f"{found_type} vs {t}") return found_type or "None" @@ -709,8 +699,7 @@ def assign_string_to_array(builder, target_array_ptr, source_string_ptr, array_l char = builder.load(src_ptr) # Store character in target - dst_ptr = builder.gep( - target_array_ptr, [ir.Constant(ir.IntType(32), 0), idx]) + dst_ptr = builder.gep(target_array_ptr, [ir.Constant(ir.IntType(32), 0), idx]) builder.store(char, dst_ptr) # Increment counter @@ -721,6 +710,5 @@ def assign_string_to_array(builder, target_array_ptr, source_string_ptr, array_l # Ensure null termination last_idx = ir.Constant(ir.IntType(32), array_length - 1) - null_ptr = builder.gep( - target_array_ptr, [ir.Constant(ir.IntType(32), 0), last_idx]) + null_ptr = builder.gep(target_array_ptr, [ir.Constant(ir.IntType(32), 0), last_idx]) builder.store(ir.Constant(ir.IntType(8), 0), null_ptr) diff --git a/pythonbpf/helper/__init__.py b/pythonbpf/helper/__init__.py index 5e538d3..1728eb8 100644 --- a/pythonbpf/helper/__init__.py +++ b/pythonbpf/helper/__init__.py @@ -1,2 +1,4 @@ from .helper_utils import HelperHandlerRegistry from .bpf_helper_handler import handle_helper_call + +__all__ = [HelperHandlerRegistry, handle_helper_call] diff --git a/pythonbpf/helper/bpf_helper_handler.py b/pythonbpf/helper/bpf_helper_handler.py index 69ca22b..87c0230 100644 --- a/pythonbpf/helper/bpf_helper_handler.py +++ b/pythonbpf/helper/bpf_helper_handler.py @@ -1,10 +1,14 @@ import ast from llvmlite import ir from enum import Enum -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) +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): @@ -18,9 +22,16 @@ class BPFHelperID(Enum): @HelperHandlerRegistry.register("ktime") -def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, - local_sym_tab=None, struct_sym_tab=None, - local_var_metadata=None): +def bpf_ktime_get_ns_emitter( + call, + map_ptr, + module, + builder, + func, + local_sym_tab=None, + struct_sym_tab=None, + local_var_metadata=None, +): """ Emit LLVM IR for bpf_ktime_get_ns helper function call. """ @@ -34,27 +45,34 @@ def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, @HelperHandlerRegistry.register("lookup") -def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func, - local_sym_tab=None, struct_sym_tab=None, - local_var_metadata=None): +def bpf_map_lookup_elem_emitter( + call, + map_ptr, + module, + builder, + func, + local_sym_tab=None, + struct_sym_tab=None, + local_var_metadata=None, +): """ Emit LLVM IR for bpf_map_lookup_elem helper function call. """ if not call.args or len(call.args) != 1: - raise ValueError("Map lookup expects exactly one argument (key), got " - f"{len(call.args)}") + raise ValueError( + "Map lookup expects exactly one argument (key), got " f"{len(call.args)}" + ) key_ptr = get_or_create_ptr_from_arg(call.args[0], builder, local_sym_tab) map_void_ptr = builder.bitcast(map_ptr, ir.PointerType()) fn_type = ir.FunctionType( ir.PointerType(), # Return type: void* [ir.PointerType(), ir.PointerType()], # Args: (void*, void*) - var_arg=False + var_arg=False, ) fn_ptr_type = ir.PointerType(fn_type) - fn_addr = ir.Constant(ir.IntType( - 64), BPFHelperID.BPF_MAP_LOOKUP_ELEM.value) + fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_MAP_LOOKUP_ELEM.value) fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False) @@ -63,33 +81,46 @@ def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func, @HelperHandlerRegistry.register("print") -def bpf_printk_emitter(call, map_ptr, module, builder, func, - local_sym_tab=None, struct_sym_tab=None, - local_var_metadata=None): +def bpf_printk_emitter( + call, + map_ptr, + module, + builder, + func, + local_sym_tab=None, + struct_sym_tab=None, + local_var_metadata=None, +): """Emit LLVM IR for bpf_printk helper function call.""" if not hasattr(func, "_fmt_counter"): func._fmt_counter = 0 if not call.args: - raise ValueError( - "bpf_printk expects at least one argument (format string)") + raise ValueError("bpf_printk expects at least one argument (format string)") args = [] if isinstance(call.args[0], ast.JoinedStr): - args = handle_fstring_print(call.args[0], module, builder, func, - local_sym_tab, struct_sym_tab, - local_var_metadata) - elif (isinstance(call.args[0], ast.Constant) and - isinstance(call.args[0].value, str)): + args = handle_fstring_print( + call.args[0], + module, + builder, + func, + local_sym_tab, + struct_sym_tab, + local_var_metadata, + ) + elif isinstance(call.args[0], ast.Constant) and isinstance(call.args[0].value, str): # TODO: We are only supporting single arguments for now. # In case of multiple args, the first one will be taken. args = simple_string_print(call.args[0].value, module, builder, func) else: raise NotImplementedError( - "Only simple strings or f-strings are supported in bpf_printk.") + "Only simple strings or f-strings are supported in bpf_printk." + ) fn_type = ir.FunctionType( - ir.IntType(64), [ir.PointerType(), ir.IntType(32)], var_arg=True) + 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) @@ -99,18 +130,25 @@ def bpf_printk_emitter(call, map_ptr, module, builder, func, @HelperHandlerRegistry.register("update") -def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func, - local_sym_tab=None, struct_sym_tab=None, - local_var_metadata=None): +def bpf_map_update_elem_emitter( + call, + map_ptr, + module, + builder, + func, + local_sym_tab=None, + struct_sym_tab=None, + local_var_metadata=None, +): """ Emit LLVM IR for bpf_map_update_elem helper function call. Expected call signature: map.update(key, value, flags=0) """ - if (not call.args or - len(call.args) < 2 or - len(call.args) > 3): - raise ValueError("Map update expects 2 or 3 args (key, value, flags), " - f"got {len(call.args)}") + if not call.args or len(call.args) < 2 or len(call.args) > 3: + raise ValueError( + "Map update expects 2 or 3 args (key, value, flags), " + f"got {len(call.args)}" + ) key_arg = call.args[0] value_arg = call.args[1] @@ -124,12 +162,11 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func, fn_type = ir.FunctionType( ir.IntType(64), [ir.PointerType(), ir.PointerType(), ir.PointerType(), ir.IntType(64)], - var_arg=False + var_arg=False, ) fn_ptr_type = ir.PointerType(fn_type) - fn_addr = ir.Constant(ir.IntType( - 64), BPFHelperID.BPF_MAP_UPDATE_ELEM.value) + fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_MAP_UPDATE_ELEM.value) fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) if isinstance(flags_val, int): @@ -138,22 +175,31 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func, flags_const = flags_val result = builder.call( - fn_ptr, [map_void_ptr, key_ptr, value_ptr, flags_const], tail=False) + fn_ptr, [map_void_ptr, key_ptr, value_ptr, flags_const], tail=False + ) return result, None @HelperHandlerRegistry.register("delete") -def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func, - local_sym_tab=None, struct_sym_tab=None, - local_var_metadata=None): +def bpf_map_delete_elem_emitter( + call, + map_ptr, + module, + builder, + func, + local_sym_tab=None, + struct_sym_tab=None, + local_var_metadata=None, +): """ Emit LLVM IR for bpf_map_delete_elem helper function call. Expected call signature: map.delete(key) """ if not call.args or len(call.args) != 1: - raise ValueError("Map delete expects exactly one argument (key), got " - f"{len(call.args)}") + raise ValueError( + "Map delete expects exactly one argument (key), got " f"{len(call.args)}" + ) key_ptr = get_or_create_ptr_from_arg(call.args[0], builder, local_sym_tab) map_void_ptr = builder.bitcast(map_ptr, ir.PointerType()) @@ -161,12 +207,11 @@ def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func, fn_type = ir.FunctionType( ir.IntType(64), # Return type: int64 (status code) [ir.PointerType(), ir.PointerType()], # Args: (void*, void*) - var_arg=False + var_arg=False, ) fn_ptr_type = ir.PointerType(fn_type) - fn_addr = ir.Constant(ir.IntType( - 64), BPFHelperID.BPF_MAP_DELETE_ELEM.value) + fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_MAP_DELETE_ELEM.value) fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False) @@ -175,15 +220,21 @@ def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func, @HelperHandlerRegistry.register("pid") -def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func, - local_sym_tab=None, struct_sym_tab=None, - local_var_metadata=None): +def bpf_get_current_pid_tgid_emitter( + call, + map_ptr, + module, + builder, + func, + local_sym_tab=None, + struct_sym_tab=None, + local_var_metadata=None, +): """ Emit LLVM IR for bpf_get_current_pid_tgid helper function call. """ # func is an arg to just have a uniform signature with other emitters - helper_id = ir.Constant(ir.IntType( - 64), BPFHelperID.BPF_GET_CURRENT_PID_TGID.value) + helper_id = ir.Constant(ir.IntType(64), BPFHelperID.BPF_GET_CURRENT_PID_TGID.value) fn_type = ir.FunctionType(ir.IntType(64), [], var_arg=False) fn_ptr_type = ir.PointerType(fn_type) fn_ptr = builder.inttoptr(helper_id, fn_ptr_type) @@ -196,18 +247,26 @@ def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func, @HelperHandlerRegistry.register("output") -def bpf_perf_event_output_handler(call, map_ptr, module, builder, func, - local_sym_tab=None, struct_sym_tab=None, - local_var_metadata=None): +def bpf_perf_event_output_handler( + call, + map_ptr, + module, + builder, + func, + local_sym_tab=None, + struct_sym_tab=None, + local_var_metadata=None, +): if len(call.args) != 1: - raise ValueError("Perf event output expects exactly one argument, " - f"got {len(call.args)}") + raise ValueError( + "Perf event output expects exactly one argument, " f"got {len(call.args)}" + ) data_arg = call.args[0] ctx_ptr = func.args[0] # First argument to the function is ctx - data_ptr, size_val = get_data_ptr_and_size(data_arg, local_sym_tab, - struct_sym_tab, - local_var_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) @@ -216,36 +275,56 @@ def bpf_perf_event_output_handler(call, map_ptr, module, builder, func, 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 + [ + 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_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) + 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, - local_sym_tab=None, map_sym_tab=None, - struct_sym_tab=None, local_var_metadata=None): +def handle_helper_call( + call, + module, + builder, + func, + local_sym_tab=None, + map_sym_tab=None, + struct_sym_tab=None, + local_var_metadata=None, +): """Process a BPF helper function call and emit the appropriate LLVM IR.""" + # Helper function to get map pointer and invoke handler def invoke_helper(method_name, map_ptr=None): handler = HelperHandlerRegistry.get_handler(method_name) if not handler: raise NotImplementedError( - f"Helper function '{method_name}' is not implemented.") - return handler(call, map_ptr, module, builder, func, - local_sym_tab, struct_sym_tab, local_var_metadata) + f"Helper function '{method_name}' is not implemented." + ) + return handler( + call, + map_ptr, + module, + builder, + func, + local_sym_tab, + struct_sym_tab, + local_var_metadata, + ) # Handle direct function calls (e.g., print(), ktime()) if isinstance(call.func, ast.Name): @@ -265,7 +344,8 @@ def handle_helper_call(call, module, builder, func, map_name = value.id else: raise NotImplementedError( - f"Unsupported map access pattern: {ast.dump(value)}") + f"Unsupported map access pattern: {ast.dump(value)}" + ) # Verify map exists and get pointer if not map_sym_tab or map_name not in map_sym_tab: diff --git a/pythonbpf/helper/helper_utils.py b/pythonbpf/helper/helper_utils.py index adb766e..ed3791b 100644 --- a/pythonbpf/helper/helper_utils.py +++ b/pythonbpf/helper/helper_utils.py @@ -8,14 +8,17 @@ logger = logging.getLogger(__name__) class HelperHandlerRegistry: """Registry for BPF helpers""" + _handlers = {} @classmethod def register(cls, helper_name): """Decorator to register a handler function for a helper""" + def decorator(func): cls._handlers[helper_name] = func return func + return decorator @classmethod @@ -55,7 +58,8 @@ def get_or_create_ptr_from_arg(arg, builder, local_sym_tab): ptr = create_int_constant_ptr(arg.value, builder) else: raise NotImplementedError( - "Only simple variable names are supported as args in map helpers.") + "Only simple variable names are supported as args in map helpers." + ) return ptr @@ -69,13 +73,13 @@ def get_flags_val(arg, builder, local_sym_tab): flags_ptr = local_sym_tab[arg.id][0] return builder.load(flags_ptr) else: - raise ValueError( - f"Variable '{arg.id}' not found in local symbol table") + raise ValueError(f"Variable '{arg.id}' not found in local symbol table") elif isinstance(arg, ast.Constant) and isinstance(arg.value, int): return arg.value raise NotImplementedError( - "Only var names or int consts are supported as map helpers flags.") + "Only var names or int consts are supported as map helpers flags." + ) def simple_string_print(string_value, module, builder, func): @@ -87,9 +91,15 @@ def simple_string_print(string_value, module, builder, func): return args -def handle_fstring_print(joined_str, module, builder, func, - local_sym_tab=None, struct_sym_tab=None, - local_var_metadata=None): +def handle_fstring_print( + joined_str, + module, + builder, + func, + local_sym_tab=None, + struct_sym_tab=None, + local_var_metadata=None, +): """Handle f-string formatting for bpf_printk emitter.""" fmt_parts = [] exprs = [] @@ -100,25 +110,34 @@ def handle_fstring_print(joined_str, module, builder, func, if isinstance(value, ast.Constant): _process_constant_in_fstring(value, fmt_parts, exprs) elif isinstance(value, ast.FormattedValue): - _process_fval(value, fmt_parts, exprs, - local_sym_tab, struct_sym_tab, - local_var_metadata) + _process_fval( + value, + fmt_parts, + exprs, + local_sym_tab, + struct_sym_tab, + local_var_metadata, + ) else: - raise NotImplementedError( - f"Unsupported f-string value type: {type(value)}") + raise NotImplementedError(f"Unsupported f-string value type: {type(value)}") fmt_str = "".join(fmt_parts) args = simple_string_print(fmt_str, module, builder, func) # NOTE: Process expressions (limited to 3 due to BPF constraints) if len(exprs) > 3: - logger.warning( - "bpf_printk supports up to 3 args, extra args will be ignored.") + logger.warning("bpf_printk supports up to 3 args, extra args will be ignored.") for expr in exprs[:3]: - arg_value = _prepare_expr_args(expr, func, module, builder, - local_sym_tab, struct_sym_tab, - local_var_metadata) + arg_value = _prepare_expr_args( + expr, + func, + module, + builder, + local_sym_tab, + struct_sym_tab, + local_var_metadata, + ) args.append(arg_value) return args @@ -133,24 +152,31 @@ def _process_constant_in_fstring(cst, fmt_parts, exprs): exprs.append(ir.Constant(ir.IntType(64), cst.value)) else: raise NotImplementedError( - f"Unsupported constant type in f-string: {type(cst.value)}") + f"Unsupported constant type in f-string: {type(cst.value)}" + ) -def _process_fval(fval, fmt_parts, exprs, - local_sym_tab, struct_sym_tab, - local_var_metadata): +def _process_fval( + fval, fmt_parts, exprs, local_sym_tab, struct_sym_tab, local_var_metadata +): """Process formatted values in f-string.""" logger.debug(f"Processing formatted value: {ast.dump(fval)}") if isinstance(fval.value, ast.Name): _process_name_in_fval(fval.value, fmt_parts, exprs, local_sym_tab) elif isinstance(fval.value, ast.Attribute): - _process_attr_in_fval(fval.value, fmt_parts, exprs, - local_sym_tab, struct_sym_tab, - local_var_metadata) + _process_attr_in_fval( + fval.value, + fmt_parts, + exprs, + local_sym_tab, + struct_sym_tab, + local_var_metadata, + ) else: raise NotImplementedError( - f"Unsupported formatted value in f-string: {type(fval.value)}") + f"Unsupported formatted value in f-string: {type(fval.value)}" + ) def _process_name_in_fval(name_node, fmt_parts, exprs, local_sym_tab): @@ -160,34 +186,39 @@ def _process_name_in_fval(name_node, fmt_parts, exprs, local_sym_tab): _populate_fval(var_type, name_node, fmt_parts, exprs) -def _process_attr_in_fval(attr_node, fmt_parts, exprs, - local_sym_tab, struct_sym_tab, - local_var_metadata): +def _process_attr_in_fval( + attr_node, fmt_parts, exprs, local_sym_tab, struct_sym_tab, local_var_metadata +): """Process attribute nodes in formatted values.""" - if (isinstance(attr_node.value, ast.Name) and - local_sym_tab and attr_node.value.id in local_sym_tab): + if ( + isinstance(attr_node.value, ast.Name) + and local_sym_tab + and attr_node.value.id in local_sym_tab + ): var_name = attr_node.value.id field_name = attr_node.attr if not local_var_metadata or var_name not in local_var_metadata: raise ValueError( - f"Metadata for '{var_name}' not found in local var metadata") + f"Metadata for '{var_name}' not found in local var metadata" + ) var_type = local_var_metadata[var_name] if var_type not in struct_sym_tab: raise ValueError( - f"Struct '{var_type}' for '{var_name}' not in symbol table") + f"Struct '{var_type}' for '{var_name}' not in symbol table" + ) struct_info = struct_sym_tab[var_type] if field_name not in struct_info.fields: - raise ValueError( - f"Field '{field_name}' not found in struct '{var_type}'") + raise ValueError(f"Field '{field_name}' not found in struct '{var_type}'") field_type = struct_info.field_type(field_name) _populate_fval(field_type, attr_node, fmt_parts, exprs) else: raise NotImplementedError( - "Only simple attribute on local vars is supported in f-strings.") + "Only simple attribute on local vars is supported in f-strings." + ) def _populate_fval(ftype, node, fmt_parts, exprs): @@ -202,14 +233,14 @@ def _populate_fval(ftype, node, fmt_parts, exprs): exprs.append(node) else: raise NotImplementedError( - f"Unsupported integer width in f-string: {ftype.width}") + f"Unsupported integer width in f-string: {ftype.width}" + ) elif ftype == ir.PointerType(ir.IntType(8)): # NOTE: We assume i8* is a string fmt_parts.append("%s") exprs.append(node) else: - raise NotImplementedError( - f"Unsupported field type in f-string: {ftype}") + raise NotImplementedError(f"Unsupported field type in f-string: {ftype}") def _create_format_string_global(fmt_str, func, module, builder): @@ -218,11 +249,11 @@ def _create_format_string_global(fmt_str, func, module, builder): func._fmt_counter += 1 fmt_gvar = ir.GlobalVariable( - module, ir.ArrayType(ir.IntType(8), len(fmt_str)), name=fmt_name) + module, ir.ArrayType(ir.IntType(8), len(fmt_str)), name=fmt_name + ) fmt_gvar.global_constant = True fmt_gvar.initializer = ir.Constant( - ir.ArrayType(ir.IntType(8), len(fmt_str)), - bytearray(fmt_str.encode("utf8")) + ir.ArrayType(ir.IntType(8), len(fmt_str)), bytearray(fmt_str.encode("utf8")) ) fmt_gvar.linkage = "internal" fmt_gvar.align = 1 @@ -230,13 +261,20 @@ def _create_format_string_global(fmt_str, func, module, builder): return builder.bitcast(fmt_gvar, ir.PointerType()) -def _prepare_expr_args(expr, func, module, builder, - local_sym_tab, struct_sym_tab, - local_var_metadata): +def _prepare_expr_args( + expr, func, module, builder, local_sym_tab, struct_sym_tab, local_var_metadata +): """Evaluate and prepare an expression to use as an arg for bpf_printk.""" - val, _ = eval_expr(func, module, builder, expr, - local_sym_tab, None, struct_sym_tab, - local_var_metadata) + val, _ = eval_expr( + func, + module, + builder, + expr, + local_sym_tab, + None, + struct_sym_tab, + local_var_metadata, + ) if val: if isinstance(val.type, ir.PointerType): @@ -246,19 +284,19 @@ def _prepare_expr_args(expr, func, module, builder, val = builder.sext(val, ir.IntType(64)) else: logger.warning( - "Only int and ptr supported in bpf_printk args. " - "Others default to 0.") + "Only int and ptr supported in bpf_printk args. " "Others default to 0." + ) val = ir.Constant(ir.IntType(64), 0) return val else: logger.warning( "Failed to evaluate expression for bpf_printk argument. " - "It will be converted to 0.") + "It will be converted to 0." + ) return ir.Constant(ir.IntType(64), 0) -def get_data_ptr_and_size(data_arg, local_sym_tab, struct_sym_tab, - local_var_metadata): +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 @@ -266,7 +304,8 @@ def get_data_ptr_and_size(data_arg, local_sym_tab, struct_sym_tab, data_ptr = local_sym_tab[data_name][0] else: raise ValueError( - f"Data variable {data_name} not found in local symbol table.") + f"Data variable {data_name} not found in local symbol table." + ) # Check if data_name is a struct if local_var_metadata and data_name in local_var_metadata: @@ -277,12 +316,14 @@ def get_data_ptr_and_size(data_arg, local_sym_tab, struct_sym_tab, return data_ptr, size_val else: raise ValueError( - f"Struct {data_type} for {data_name} not in symbol table.") + f"Struct {data_type} for {data_name} not in symbol table." + ) else: raise ValueError( f"Metadata for variable {data_name} " - "not found in local variable metadata.") + "not found in local variable metadata." + ) else: raise NotImplementedError( - "Only simple object names are supported " - "as data in perf event output.") + "Only simple object names are supported " "as data in perf event output." + )