import ast from llvmlite import ir from enum import Enum from .maps_utils import MapProcessorRegistry from ..debuginfo import DebugInfoGenerator import logging logger = logging.getLogger(__name__) def maps_proc(tree, module, chunks): """Process all functions decorated with @map to find BPF maps""" map_sym_tab = {} for func_node in chunks: if is_map(func_node): print(f"Found BPF map: {func_node.name}") map_sym_tab[func_node.name] = process_bpf_map(func_node, module) return map_sym_tab def is_map(func_node): return any( isinstance(decorator, ast.Name) and decorator.id == "map" for decorator in func_node.decorator_list ) class BPFMapType(Enum): HASH = 1 PERF_EVENT_ARRAY = 4 def create_bpf_map(module, map_name, map_params): """Create a BPF map in the module with given parameters and debug info""" # Create the anonymous struct type for BPF map map_struct_type = ir.LiteralStructType( [ir.PointerType() for _ in range(len(map_params))] ) # Create the global variable map_global = ir.GlobalVariable(module, map_struct_type, name=map_name) map_global.linkage = "dso_local" map_global.global_constant = False map_global.initializer = ir.Constant(map_struct_type, None) map_global.section = ".maps" map_global.align = 8 logger.info(f"Created BPF map: {map_name} with params {map_params}") return map_global def create_map_debug_info(module, map_global, map_name, map_params): """Generate debug information metadata for BPF map""" generator = DebugInfoGenerator(module) 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.HASH).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 ) value_ptr = generator.create_pointer_type( array_type if "value_size" in map_params else ulong_type, 64 ) elements_arr = [] # Create struct members # scope field does not appear for some reason cnt = 0 for elem in map_params: if elem == "max_entries": continue if elem == "type": ptr = type_ptr elif "key" in elem: ptr = key_ptr else: ptr = value_ptr # TODO: the best way to do this is not 64, but get the size each time. this will not work for structs. member = generator.create_struct_member(elem, ptr, cnt * 64) elements_arr.append(member) cnt += 1 if "max_entries" in map_params: max_entries_array = generator.create_array_type( uint_type, map_params["max_entries"] ) max_entries_ptr = generator.create_pointer_type(max_entries_array, 64) max_entries_member = generator.create_struct_member( "max_entries", max_entries_ptr, cnt * 64 ) elements_arr.append(max_entries_member) # Create the struct type struct_type = generator.create_struct_type( elements_arr, 64 * len(elements_arr), is_distinct=True ) # Create global variable debug info global_var = generator.create_global_var_debug_info( map_name, struct_type, is_local=False ) # Attach debug info to the global variable map_global.set_metadata("dbg", global_var) return global_var @MapProcessorRegistry.register("HashMap") def process_hash_map(map_name, rval, module): """Process a BPF_HASH map declaration""" logger.info(f"Processing HashMap: {map_name}") map_params = {"type": BPFMapType.HASH} # Assuming order: key_type, value_type, max_entries if len(rval.args) >= 1 and isinstance(rval.args[0], ast.Name): map_params["key"] = rval.args[0].id if len(rval.args) >= 2 and isinstance(rval.args[1], ast.Name): map_params["value"] = rval.args[1].id if len(rval.args) >= 3 and isinstance(rval.args[2], ast.Constant): const_val = rval.args[2].value if isinstance(const_val, (int, str)): # safe check map_params["max_entries"] = const_val for keyword in rval.keywords: if keyword.arg == "key" and isinstance(keyword.value, ast.Name): map_params["key"] = keyword.value.id elif keyword.arg == "value" and isinstance(keyword.value, ast.Name): map_params["value"] = keyword.value.id elif keyword.arg == "max_entries" and isinstance(keyword.value, ast.Constant): const_val = keyword.value.value if isinstance(const_val, (int, str)): map_params["max_entries"] = const_val 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, map_name, map_params) return map_global @MapProcessorRegistry.register("PerfEventArray") def process_perf_event_map(map_name, rval, module): """Process a BPF_PERF_EVENT_ARRAY map declaration""" logger.info(f"Processing PerfEventArray: {map_name}") map_params = {"type": BPFMapType.PERF_EVENT_ARRAY} if len(rval.args) >= 1 and isinstance(rval.args[0], ast.Name): map_params["key_size"] = rval.args[0].id if len(rval.args) >= 2 and isinstance(rval.args[1], ast.Name): map_params["value_size"] = rval.args[1].id for keyword in rval.keywords: if keyword.arg == "key_size" and isinstance(keyword.value, ast.Name): map_params["key_size"] = keyword.value.id elif keyword.arg == "value_size" and isinstance(keyword.value, ast.Name): map_params["value_size"] = keyword.value.id 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, map_name, map_params) return map_global def process_bpf_map(func_node, module): """Process a BPF map (a function decorated with @map)""" map_name = func_node.name logger.info(f"Processing BPF map: {map_name}") # For now, assume single return statement return_stmt = None for stmt in func_node.body: if isinstance(stmt, ast.Return): return_stmt = stmt break if return_stmt is None: raise ValueError("BPF map must have a return statement") rval = return_stmt.value 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) else: logger.warning( f"Unknown map type " f"{rval.func.id}, defaulting to HashMap" ) return process_hash_map(map_name, rval, module) else: raise ValueError("Function under @map must return a map")