From b2a57edf111f2a4191a44fb9439033a47157f3e3 Mon Sep 17 00:00:00 2001 From: Pragyansh Chaturvedi Date: Mon, 13 Oct 2025 19:01:01 +0530 Subject: [PATCH] Simplify maps_pass --- pythonbpf/maps/map_debug_info.py | 93 ++++++++++++++ pythonbpf/maps/map_types.py | 39 ++++++ pythonbpf/maps/maps_pass.py | 201 +++++-------------------------- 3 files changed, 161 insertions(+), 172 deletions(-) create mode 100644 pythonbpf/maps/map_debug_info.py create mode 100644 pythonbpf/maps/map_types.py diff --git a/pythonbpf/maps/map_debug_info.py b/pythonbpf/maps/map_debug_info.py new file mode 100644 index 0000000..b1a25d9 --- /dev/null +++ b/pythonbpf/maps/map_debug_info.py @@ -0,0 +1,93 @@ +from pythonbpf.debuginfo import DebugInfoGenerator +from .map_types import BPFMapType + + +def create_map_debug_info(module, map_global, map_name, map_params): + """Generate debug info metadata for BPF maps HASH and PERF_EVENT_ARRAY""" + 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.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 + ) + 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 + + +def create_ringbuf_debug_info(module, map_global, map_name, map_params): + """Generate debug information metadata for BPF RINGBUF map""" + generator = DebugInfoGenerator(module) + + int_type = generator.get_int32_type() + + type_array = generator.create_array_type( + int_type, map_params.get("type", BPFMapType.RINGBUF).value + ) + type_ptr = generator.create_pointer_type(type_array, 64) + type_member = generator.create_struct_member("type", type_ptr, 0) + + max_entries_array = generator.create_array_type(int_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, 64 + ) + + elements_arr = [type_member, max_entries_member] + + struct_type = generator.create_struct_type(elements_arr, 128, is_distinct=True) + + global_var = generator.create_global_var_debug_info( + map_name, struct_type, is_local=False + ) + map_global.set_metadata("dbg", global_var) + return global_var diff --git a/pythonbpf/maps/map_types.py b/pythonbpf/maps/map_types.py new file mode 100644 index 0000000..3a8dcdd --- /dev/null +++ b/pythonbpf/maps/map_types.py @@ -0,0 +1,39 @@ +from enum import Enum + + +class BPFMapType(Enum): + UNSPEC = 0 + HASH = 1 + ARRAY = 2 + PROG_ARRAY = 3 + PERF_EVENT_ARRAY = 4 + PERCPU_HASH = 5 + PERCPU_ARRAY = 6 + STACK_TRACE = 7 + CGROUP_ARRAY = 8 + LRU_HASH = 9 + LRU_PERCPU_HASH = 10 + LPM_TRIE = 11 + ARRAY_OF_MAPS = 12 + HASH_OF_MAPS = 13 + DEVMAP = 14 + SOCKMAP = 15 + CPUMAP = 16 + XSKMAP = 17 + SOCKHASH = 18 + CGROUP_STORAGE_DEPRECATED = 19 + CGROUP_STORAGE = 19 + REUSEPORT_SOCKARRAY = 20 + PERCPU_CGROUP_STORAGE_DEPRECATED = 21 + PERCPU_CGROUP_STORAGE = 21 + QUEUE = 22 + STACK = 23 + SK_STORAGE = 24 + DEVMAP_HASH = 25 + STRUCT_OPS = 26 + RINGBUF = 27 + INODE_STORAGE = 28 + TASK_STORAGE = 29 + BLOOM_FILTER = 30 + USER_RINGBUF = 31 + CGRP_STORAGE = 32 diff --git a/pythonbpf/maps/maps_pass.py b/pythonbpf/maps/maps_pass.py index 95748a8..8459848 100644 --- a/pythonbpf/maps/maps_pass.py +++ b/pythonbpf/maps/maps_pass.py @@ -1,10 +1,11 @@ import ast +import logging from logging import Logger from llvmlite import ir -from enum import Enum + from .maps_utils import MapProcessorRegistry -from pythonbpf.debuginfo import DebugInfoGenerator -import logging +from .map_types import BPFMapType +from .map_debug_info import create_map_debug_info, create_ringbuf_debug_info logger: Logger = logging.getLogger(__name__) @@ -26,44 +27,6 @@ def is_map(func_node): ) -class BPFMapType(Enum): - UNSPEC = 0 - HASH = 1 - ARRAY = 2 - PROG_ARRAY = 3 - PERF_EVENT_ARRAY = 4 - PERCPU_HASH = 5 - PERCPU_ARRAY = 6 - STACK_TRACE = 7 - CGROUP_ARRAY = 8 - LRU_HASH = 9 - LRU_PERCPU_HASH = 10 - LPM_TRIE = 11 - ARRAY_OF_MAPS = 12 - HASH_OF_MAPS = 13 - DEVMAP = 14 - SOCKMAP = 15 - CPUMAP = 16 - XSKMAP = 17 - SOCKHASH = 18 - CGROUP_STORAGE_DEPRECATED = 19 - CGROUP_STORAGE = 19 - REUSEPORT_SOCKARRAY = 20 - PERCPU_CGROUP_STORAGE_DEPRECATED = 21 - PERCPU_CGROUP_STORAGE = 21 - QUEUE = 22 - STACK = 23 - SK_STORAGE = 24 - DEVMAP_HASH = 25 - STRUCT_OPS = 26 - RINGBUF = 27 - INODE_STORAGE = 28 - TASK_STORAGE = 29 - BLOOM_FILTER = 30 - USER_RINGBUF = 31 - CGRP_STORAGE = 32 - - def create_bpf_map(module, map_name, map_params): """Create a BPF map in the module with given parameters and debug info""" @@ -84,114 +47,37 @@ def create_bpf_map(module, map_name, map_params): return map_global -def create_map_debug_info(module, map_global, map_name, map_params): - """Generate debug info metadata for BPF maps HASH and PERF_EVENT_ARRAY""" - generator = DebugInfoGenerator(module) +def _parse_map_params(rval, expected_args=None): + """Parse map parameters from call arguments and keywords.""" - 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 - ) - value_ptr = generator.create_pointer_type( - array_type if "value_size" in map_params else ulong_type, 64 - ) + params = {} - elements_arr = [] + # Parse positional arguments + if expected_args: + for i, arg_name in enumerate(expected_args): + if i < len(rval.args): + arg = rval.args[i] + if isinstance(arg, ast.Name): + params[arg_name] = arg.id + elif isinstance(arg, ast.Constant): + params[arg_name] = arg.value - # 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 + # Parse keyword arguments (override positional) + for keyword in rval.keywords: + if isinstance(keyword.value, ast.Name): + params[keyword.arg] = keyword.value.id + elif isinstance(keyword.value, ast.Constant): + params[keyword.arg] = keyword.value.value - 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 - - -def create_ringbuf_debug_info(module, map_global, map_name, map_params): - """Generate debug information metadata for BPF RINGBUF map""" - generator = DebugInfoGenerator(module) - - int_type = generator.get_int32_type() - - type_array = generator.create_array_type( - int_type, map_params.get("type", BPFMapType.RINGBUF).value - ) - type_ptr = generator.create_pointer_type(type_array, 64) - type_member = generator.create_struct_member("type", type_ptr, 0) - - max_entries_array = generator.create_array_type(int_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, 64 - ) - - elements_arr = [type_member, max_entries_member] - - struct_type = generator.create_struct_type(elements_arr, 128, is_distinct=True) - - global_var = generator.create_global_var_debug_info( - map_name, struct_type, is_local=False - ) - map_global.set_metadata("dbg", global_var) - return global_var + return params @MapProcessorRegistry.register("RingBuf") def process_ringbuf_map(map_name, rval, module): """Process a BPF_RINGBUF map declaration""" logger.info(f"Processing Ringbuf: {map_name}") - map_params = {"type": BPFMapType.RINGBUF} - - # Parse max_entries if present - if len(rval.args) >= 1 and isinstance(rval.args[0], ast.Constant): - const_val = rval.args[0].value - if isinstance(const_val, int): - map_params["max_entries"] = const_val - - for keyword in rval.keywords: - if keyword.arg == "max_entries" and isinstance(keyword.value, ast.Constant): - const_val = keyword.value.value - if isinstance(const_val, int): - map_params["max_entries"] = const_val + map_params = _parse_map_params(rval, expected_args=["max_entries"]) + map_params["type"] = BPFMapType.RINGBUF logger.info(f"Ringbuf map parameters: {map_params}") @@ -204,27 +90,8 @@ def process_ringbuf_map(map_name, rval, module): 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 + map_params = _parse_map_params(rval, expected_args=["key", "value", "max_entries"]) + map_params["type"] = BPFMapType.HASH logger.info(f"Map parameters: {map_params}") map_global = create_bpf_map(module, map_name, map_params) @@ -237,18 +104,8 @@ def process_hash_map(map_name, rval, module): 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 + map_params = _parse_map_params(rval, expected_args=["key_size", "value_size"]) + map_params["type"] = BPFMapType.PERF_EVENT_ARRAY logger.info(f"Map parameters: {map_params}") map_global = create_bpf_map(module, map_name, map_params)