diff --git a/examples/execve2.py b/examples/execve2.py index 237f2b5..d93759f 100644 --- a/examples/execve2.py +++ b/examples/execve2.py @@ -28,10 +28,7 @@ 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" +@bpf +@bpfglobal +def LICENSE() -> str: + return "GPL" diff --git a/pythonbpf/functions_pass.py b/pythonbpf/functions_pass.py index b4fc417..a66a9d4 100644 --- a/pythonbpf/functions_pass.py +++ b/pythonbpf/functions_pass.py @@ -12,6 +12,8 @@ def get_probe_string(func_node): # For helper functions, we return "helper" for decorator in func_node.decorator_list: + if isinstance(decorator, ast.Name) and decorator.id == "bpfglobal": + return None if isinstance(decorator, ast.Call) and isinstance(decorator.func, ast.Name): if decorator.func.id == "section" and len(decorator.args) == 1: arg = decorator.args[0] @@ -134,7 +136,9 @@ def process_bpf_chunk(func_node, module, return_type, map_sym_tab): param = func.args[0] param.add_attribute("nocapture") - func.section = get_probe_string(func_node) + probe_string = get_probe_string(func_node) + if probe_string is not None: + func.section = probe_string block = func.append_basic_block(name="entry") builder = ir.IRBuilder(block) @@ -151,6 +155,9 @@ def func_proc(tree, module, chunks, map_sym_tab): if isinstance(decorator, ast.Name) and decorator.id == "map": is_global = True break + elif isinstance(decorator, ast.Name) and decorator.id == "bpfglobal": + is_global = True + break if is_global: continue func_type = get_probe_string(func_node) diff --git a/pythonbpf/license_pass.py b/pythonbpf/license_pass.py index 0f986a0..ddab38c 100644 --- a/pythonbpf/license_pass.py +++ b/pythonbpf/license_pass.py @@ -18,20 +18,28 @@ def emit_license(module: ir.Module, license_str: str): return gvar def license_processing(tree, module): - """Process the LICENSE assignment in the given AST tree and return the section name""" + """Process the LICENSE function decorated with @bpf and @bpfglobal and return the section name""" count = 0 for node in tree.body: - if isinstance(node, ast.Assign): - for target in node.targets: - if isinstance(target, ast.Name) and target.id == "LICENSE": - if count == 0: - count += 1 - if isinstance(node.value, ast.Constant) and isinstance(node.value.value, str): - emit_license(module, node.value.value) - return "LICENSE" - else: - print("ERROR: LICENSE must be a string literal") - return None + if isinstance(node, ast.FunctionDef) and node.name == "LICENSE": + # check decorators + decorators = [dec.id for dec in node.decorator_list if isinstance(dec, ast.Name)] + if "bpf" in decorators and "bpfglobal" in decorators: + if count == 0: + count += 1 + # check function body has a return string + if ( + len(node.body) == 1 + and isinstance(node.body[0], ast.Return) + and isinstance(node.body[0].value, ast.Constant) + and isinstance(node.body[0].value.value, str) + ): + emit_license(module, node.body[0].value.value) + return "LICENSE" else: - print("ERROR: LICENSE already assigned") + print("ERROR: LICENSE() must return a string literal") return None + else: + print("ERROR: LICENSE already defined") + return None + return None diff --git a/pythonbpf/type_deducer.py b/pythonbpf/type_deducer.py index 5d63864..49dd224 100644 --- a/pythonbpf/type_deducer.py +++ b/pythonbpf/type_deducer.py @@ -14,6 +14,8 @@ def ctypes_to_ir(ctype: str): "c_float": ir.FloatType(), "c_double": ir.DoubleType(), "c_void_p": ir.IntType(64), + # Not so sure about this one + "str": ir.PointerType(ir.IntType(8)) } if ctype in mapping: return mapping[ctype]