diff --git a/examples/execve2.py b/examples/execve2.py index 73e71ec..ca22e8f 100644 --- a/examples/execve2.py +++ b/examples/execve2.py @@ -1,12 +1,12 @@ -from pythonbpf.decorators import bpf, bpfglobal, section +from pythonbpf.decorators import bpf, map, section, bpfglobal from ctypes import c_void_p, c_int64, c_int32, c_uint64 from pythonbpf.helpers import bpf_ktime_get_ns from pythonbpf.maps import HashMap @bpf -@bpfglobal -def last(): +@map +def last() -> HashMap: return HashMap(key_type=c_uint64, value_type=c_uint64, max_entries=1) @@ -25,5 +25,9 @@ def hello_again(ctx: c_void_p) -> c_int64: ts = bpf_ktime_get_ns() return c_int64(0) +# @bpf +# @bpfglobal +# def LICENSE() -> str: +# return "GPL" LICENSE = "GPL" diff --git a/pythonbpf/decorators.py b/pythonbpf/decorators.py index c7a1071..6bf3b24 100644 --- a/pythonbpf/decorators.py +++ b/pythonbpf/decorators.py @@ -9,6 +9,10 @@ def bpfglobal(func): func._is_bpfglobal = True return func +def map(func): + """Decorator to mark a function as a BPF map.""" + func._is_map = True + return func def section(name: str): def wrapper(fn): diff --git a/pythonbpf/functions_pass.py b/pythonbpf/functions_pass.py index 4598db9..08b068a 100644 --- a/pythonbpf/functions_pass.py +++ b/pythonbpf/functions_pass.py @@ -90,19 +90,11 @@ def process_bpf_chunk(func_node, module, return_type): def create_bpf_map(module, map_name, map_params): """Create a BPF map in the module with the given parameters""" - type_mapping = { - 'c_uint32': ir.IntType(32), - 'c_uint64': ir.IntType(64), - 'c_int32': ir.IntType(32), - 'c_int64': ir.IntType(64), - # Add more mappings as needed - } - key_type_str = map_params.get('key_type', 'c_uint32') value_type_str = map_params.get('value_type', 'c_uint32') - key_type = type_mapping.get(key_type_str, ir.IntType(32)) - value_type = type_mapping.get(value_type_str, ir.IntType(32)) + key_type = ctypes_to_ir(key_type_str) + value_type = ctypes_to_ir(value_type_str) map_struct_type = ir.LiteralStructType([ ir.PointerType(), # type diff --git a/pythonbpf/globals_pass.py b/pythonbpf/globals_pass.py index 43cc535..6275196 100644 --- a/pythonbpf/globals_pass.py +++ b/pythonbpf/globals_pass.py @@ -42,5 +42,8 @@ def globals_processing(tree, module: ir.Module): elif isinstance(dec, ast.Name) and dec.id == "bpfglobal": collected.append(node.name) + + elif isinstance(dec, ast.Name) and dec.id == "map": + collected.append(node.name) emit_globals(module, collected)