make license handling easier

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)

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

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]