Merge pull request #5 from pythonbpf/struct_refactor

Struct refactor

.gitignore

@@ -6,3 +6,4 @@
 __pycache__/
 *.ll
 *.o
+.ipynb_checkpoints/

examples/execve5.py

@@ -26,10 +26,10 @@ def hello(ctx: c_void_p) -> c_int32:
     ts = ktime()
     process_id = pid()
     strobj = "hellohellohello"
-    dataobj.pid = process_id
-    dataobj.ts = ts
+    dataobj.pid = pid()
+    dataobj.ts = ktime()
     # dataobj.comm = strobj
-    print(f"clone called at {ts} by pid {process_id}, comm {strobj}")
+    print(f"clone called at {dataobj.ts} by pid {dataobj.pid}, comm {strobj}")
     events.output(dataobj)
     return c_int32(0)
 

pythonbpf/bpf_helper_handler.py

@@ -3,7 +3,7 @@ from llvmlite import ir
 from .expr_pass import eval_expr
 
 
-def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, local_var_metadata=None):
+def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
     """
     Emit LLVM IR for bpf_ktime_get_ns helper function call.
     """
@@ -63,7 +63,7 @@ def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func, local_sym_
     return result, ir.PointerType()
 
 
-def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, local_var_metadata=None):
+def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
     if not hasattr(func, "_fmt_counter"):
         func._fmt_counter = 0
 
@@ -101,10 +101,42 @@ def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None,
                         else:
                             raise NotImplementedError(
                                 "Only integer and pointer types are supported in formatted values.")
-                        print("Formatted value variable:", var_ptr, var_type)
                     else:
                         raise ValueError(
                             f"Variable {value.value.id} not found in local symbol table.")
+                elif isinstance(value.value, ast.Attribute):
+                    # object field access from struct
+                    if isinstance(value.value.value, ast.Name) and local_sym_tab and value.value.value.id in local_sym_tab:
+                        var_name = value.value.value.id
+                        field_name = value.value.attr
+                        if local_var_metadata and var_name in local_var_metadata:
+                            var_type = local_var_metadata[var_name]
+                            if var_type in struct_sym_tab:
+                                struct_info = struct_sym_tab[var_type]
+                                if field_name in struct_info.fields:
+                                    field_type = struct_info.field_type(
+                                        field_name)
+                                    if isinstance(field_type, ir.IntType):
+                                        fmt_parts.append("%lld")
+                                        exprs.append(value.value)
+                                    elif field_type == ir.PointerType(ir.IntType(8)):
+                                        fmt_parts.append("%s")
+                                        exprs.append(value.value)
+                                    else:
+                                        raise NotImplementedError(
+                                            "Only integer and pointer types are supported in formatted values.")
+                                else:
+                                    raise ValueError(
+                                        f"Field {field_name} not found in struct {var_type}.")
+                            else:
+                                raise ValueError(
+                                    f"Struct type {var_type} for variable {var_name} not found in struct symbol table.")
+                        else:
+                            raise ValueError(
+                                f"Metadata for variable {var_name} not found in local variable metadata.")
+                    else:
+                        raise ValueError(
+                            f"Variable {value.value.value.id} not found in local symbol table.")
                 else:
                     raise NotImplementedError(
                         "Only simple variable names are supported in formatted values.")
@@ -119,12 +151,12 @@ def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None,
         fmt_gvar = ir.GlobalVariable(
             module, ir.ArrayType(ir.IntType(8), len(fmt_str)), name=fmt_name)
         fmt_gvar.global_constant = True
-        fmt_gvar.initializer = ir.Constant(
+        fmt_gvar.initializer = ir.Constant(                 # type: ignore
             ir.ArrayType(ir.IntType(8), len(fmt_str)),
             bytearray(fmt_str.encode("utf8"))
         )
         fmt_gvar.linkage = "internal"
-        fmt_gvar.align = 1
+        fmt_gvar.align = 1              # type: ignore
 
         fmt_ptr = builder.bitcast(fmt_gvar, ir.PointerType())
 
@@ -136,8 +168,9 @@ def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None,
                 "Warning: bpf_printk supports up to 3 arguments, extra arguments will be ignored.")
 
         for expr in exprs[:3]:
+            print(f"{ast.dump(expr)}")
             val, _ = eval_expr(func, module, builder,
-                               expr, local_sym_tab, None)
+                               expr, local_sym_tab, None, struct_sym_tab, local_var_metadata)
             if val:
                 if isinstance(val.type, ir.PointerType):
                     val = builder.ptrtoint(val, ir.IntType(64))
@@ -339,7 +372,7 @@ def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func, local_sym_
     return result, None
 
 
-def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, local_var_metadata=None):
+def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
     """
     Emit LLVM IR for bpf_get_current_pid_tgid helper function call.
     """
@@ -375,7 +408,7 @@ def bpf_perf_event_output_handler(call, map_ptr, module, builder, func, local_sy
             data_type = local_var_metadata[data_name]
             if data_type in struct_sym_tab:
                 struct_info = struct_sym_tab[data_type]
-                size_val = ir.Constant(ir.IntType(64), struct_info["size"])
+                size_val = ir.Constant(ir.IntType(64), struct_info.size)
             else:
                 raise ValueError(
                     f"Struct type {data_type} for variable {data_name} not found in struct symbol table.")
@@ -420,11 +453,12 @@ helper_func_list = {
 
 
 def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
+    print(local_var_metadata)
     if isinstance(call.func, ast.Name):
         func_name = call.func.id
         if func_name in helper_func_list:
             # it is not a map method call
-            return helper_func_list[func_name](call, None, module, builder, func, local_sym_tab)
+            return helper_func_list[func_name](call, None, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata)
         else:
             raise NotImplementedError(
                 f"Function {func_name} is not implemented as a helper function.")

pythonbpf/codegen.py

@@ -3,13 +3,14 @@ from llvmlite import ir
 from .license_pass import license_processing
 from .functions_pass import func_proc
 from .maps_pass import maps_proc
-from .structs_pass import structs_proc
+from .structs.structs_pass import structs_proc
 from .globals_pass import globals_processing
 import os
 import subprocess
 import inspect
 from pathlib import Path
 from pylibbpf import BpfProgram
+import tempfile
 
 
 def find_bpf_chunks(tree):
@@ -122,14 +123,17 @@ def compile():
 
 def BPF() -> BpfProgram:
     caller_frame = inspect.stack()[1]
-    caller_file = Path(caller_frame.filename).resolve()
-    ll_file = Path("/tmp") / caller_file.with_suffix(".ll").name
-    o_file = Path("/tmp") / caller_file.with_suffix(".o").name
-    compile_to_ir(str(caller_file), str(ll_file))
-
+    src = inspect.getsource(caller_frame.frame)
+    with tempfile.NamedTemporaryFile(mode="w+", delete=True, suffix=".py") as f, \
+        tempfile.NamedTemporaryFile(mode="w+", delete=True, suffix=".ll") as inter, \
+        tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".o") as obj_file:
+        f.write(src)
+        f.flush()
+        source = f.name
+        compile_to_ir(source, str(inter.name))
         subprocess.run([
             "llc", "-march=bpf", "-filetype=obj", "-O2",
-        str(ll_file), "-o", str(o_file)
+            str(inter.name), "-o", str(obj_file.name)
         ], check=True)
 
-    return BpfProgram(str(o_file))
+        return BpfProgram(str(obj_file.name))

pythonbpf/expr_pass.py

@@ -4,6 +4,7 @@ from llvmlite import ir
 
 def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab, structs_sym_tab=None, local_var_metadata=None):
     print(f"Evaluating expression: {ast.dump(expr)}")
+    print(local_var_metadata)
     if isinstance(expr, ast.Name):
         if expr.id in local_sym_tab:
             var = local_sym_tab[expr.id][0]
@@ -66,6 +67,23 @@ def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab, structs_s
                     if method_name in helper_func_list:
                         return handle_helper_call(
                             expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
+    elif isinstance(expr, ast.Attribute):
+        if isinstance(expr.value, ast.Name):
+            var_name = expr.value.id
+            attr_name = expr.attr
+            if var_name in local_sym_tab:
+                var_ptr, var_type = local_sym_tab[var_name]
+                print(f"Loading attribute "
+                      f"{attr_name} from variable {var_name}")
+                print(f"Variable type: {var_type}, Variable ptr: {var_ptr}")
+                print(local_var_metadata)
+                if local_var_metadata and var_name in local_var_metadata:
+                    metadata = structs_sym_tab[local_var_metadata[var_name]]
+                    if attr_name in metadata.fields:
+                        gep = metadata.gep(builder, var_ptr, attr_name)
+                        val = builder.load(gep)
+                        field_type = metadata.field_type(attr_name)
+                        return val, field_type
     print("Unsupported expression evaluation")
     return None
 
@@ -73,6 +91,7 @@ def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab, structs_s
 def handle_expr(func, module, builder, expr, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata):
     """Handle expression statements in the function body."""
     print(f"Handling expression: {ast.dump(expr)}")
+    print(local_var_metadata)
     call = expr.value
     if isinstance(call, ast.Call):
         eval_expr(func, module, builder, call, local_sym_tab,

pythonbpf/functions_pass.py

@@ -49,16 +49,12 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc
             struct_type = local_var_metadata[var_name]
             struct_info = structs_sym_tab[struct_type]
 
-            if field_name in struct_info["fields"]:
-                field_idx = struct_info["fields"][field_name]
-                struct_ptr = local_sym_tab[var_name][0]
-                field_ptr = builder.gep(
-                    struct_ptr, [ir.Constant(ir.IntType(32), 0),
-                                 ir.Constant(ir.IntType(32), field_idx)],
-                    inbounds=True)
+            if field_name in struct_info.fields:
+                field_ptr = struct_info.gep(
+                    builder, local_sym_tab[var_name][0], field_name)
                 val = eval_expr(func, module, builder, rval,
                                 local_sym_tab, map_sym_tab, structs_sym_tab)
-                if isinstance(struct_info["field_types"][field_idx], ir.ArrayType) and val[1] == ir.PointerType(ir.IntType(8)):
+                if isinstance(struct_info.field_type(field_name), ir.ArrayType) and val[1] == ir.PointerType(ir.IntType(8)):
                     # TODO: Figure it out, not a priority rn
                     # Special case for string assignment to char array
                     # str_len = struct_info["field_types"][field_idx].count
@@ -138,7 +134,7 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc
                 print(f"Dereferenced and assigned to {var_name}")
             elif call_type in structs_sym_tab and len(rval.args) == 0:
                 struct_info = structs_sym_tab[call_type]
-                ir_type = struct_info["type"]
+                ir_type = struct_info.ir_type
                 # var = builder.alloca(ir_type, name=var_name)
                 # Null init
                 builder.store(ir.Constant(ir_type, None),
@@ -282,6 +278,7 @@ def handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab, structs_s
 def process_stmt(func, module, builder, stmt, local_sym_tab, map_sym_tab, structs_sym_tab, did_return, ret_type=ir.IntType(64)):
     print(f"Processing statement: {ast.dump(stmt)}")
     if isinstance(stmt, ast.Expr):
+        print(local_var_metadata)
         handle_expr(func, module, builder, stmt, local_sym_tab,
                     map_sym_tab, structs_sym_tab, local_var_metadata)
     elif isinstance(stmt, ast.Assign):
@@ -363,7 +360,7 @@ def allocate_mem(module, builder, body, func, ret_type, map_sym_tab, local_sym_t
                             f"Pre-allocated variable {var_name} for deref")
                     elif call_type in structs_sym_tab:
                         struct_info = structs_sym_tab[call_type]
-                        ir_type = struct_info["type"]
+                        ir_type = struct_info.ir_type
                         var = builder.alloca(ir_type, name=var_name)
                         local_var_metadata[var_name] = call_type
                         print(
@@ -547,6 +544,8 @@ def infer_return_type(func_node: ast.FunctionDef):
     return found_type or "None"
 
 # For string assignment to fixed-size arrays
+
+
 def assign_string_to_array(builder, target_array_ptr, source_string_ptr, array_length):
     """
     Copy a string (i8*) to a fixed-size array ([N x i8]*)
@@ -566,7 +565,8 @@ def assign_string_to_array(builder, target_array_ptr, source_string_ptr, array_l
     # Copy loop
     builder.position_at_end(copy_block)
     idx = builder.load(i)
-    in_bounds = builder.icmp_unsigned('<', idx, ir.Constant(ir.IntType(32), array_length))
+    in_bounds = builder.icmp_unsigned(
+        '<', idx, ir.Constant(ir.IntType(32), array_length))
     builder.cbranch(in_bounds, copy_block, end_block)
 
     with builder.if_then(in_bounds):
@@ -575,7 +575,8 @@ def assign_string_to_array(builder, target_array_ptr, source_string_ptr, array_l
         char = builder.load(src_ptr)
 
         # Store character in target
-        dst_ptr = builder.gep(target_array_ptr, [ir.Constant(ir.IntType(32), 0), idx])
+        dst_ptr = builder.gep(
+            target_array_ptr, [ir.Constant(ir.IntType(32), 0), idx])
         builder.store(char, dst_ptr)
 
         # Increment counter
@@ -586,5 +587,6 @@ def assign_string_to_array(builder, target_array_ptr, source_string_ptr, array_l
 
     # Ensure null termination
     last_idx = ir.Constant(ir.IntType(32), array_length - 1)
-    null_ptr = builder.gep(target_array_ptr, [ir.Constant(ir.IntType(32), 0), last_idx])
+    null_ptr = builder.gep(
+        target_array_ptr, [ir.Constant(ir.IntType(32), 0), last_idx])
     builder.store(ir.Constant(ir.IntType(8), 0), null_ptr)

pythonbpf/structs/struct_type.py

@@ -0,0 +1,31 @@
+from llvmlite import ir
+
+
+class StructType:
+    def __init__(self, ir_type, fields, size):
+        self.ir_type = ir_type
+        self.fields = fields
+        self.size = size
+
+    def field_idx(self, field_name):
+        return list(self.fields.keys()).index(field_name)
+
+    def field_type(self, field_name):
+        return self.fields[field_name]
+
+    def gep(self, builder, ptr, field_name):
+        idx = self.field_idx(field_name)
+        return builder.gep(ptr, [ir.Constant(ir.IntType(32), 0),
+                                 ir.Constant(ir.IntType(32), idx)],
+                           inbounds=True)
+
+    def field_size(self, field_name):
+        fld = self.fields[field_name]
+        if isinstance(fld, ir.ArrayType):
+            return fld.count * (fld.element.width // 8)
+        elif isinstance(fld, ir.IntType):
+            return fld.width // 8
+        elif isinstance(fld, ir.PointerType):
+            return 8
+
+        raise TypeError(f"Unsupported field type: {fld}")

pythonbpf/structs/structs_pass.py

@@ -0,0 +1,97 @@
+import ast
+import logging
+from llvmlite import ir
+from pythonbpf.type_deducer import ctypes_to_ir
+from .struct_type import StructType
+
+logger = logging.getLogger(__name__)
+
+# TODO: Shall we allow the following syntax:
+# struct MyStruct:
+#     field1: int
+#     field2: str(32)
+# Where int is mapped to c_uint64?
+# Shall we just int64, int32 and uint32 similarly?
+
+
+def structs_proc(tree, module, chunks):
+    """ Process all class definitions to find BPF structs """
+    structs_sym_tab = {}
+    for cls_node in chunks:
+        if is_bpf_struct(cls_node):
+            print(f"Found BPF struct: {cls_node.name}")
+            struct_info = process_bpf_struct(cls_node, module)
+            structs_sym_tab[cls_node.name] = struct_info
+    return structs_sym_tab
+
+
+def is_bpf_struct(cls_node):
+    return any(
+        isinstance(decorator, ast.Name) and decorator.id == "struct"
+        for decorator in cls_node.decorator_list
+    )
+
+
+def process_bpf_struct(cls_node, module):
+    """ Process a single BPF struct definition """
+
+    fields = parse_struct_fields(cls_node)
+    field_types = list(fields.values())
+    total_size = calc_struct_size(field_types)
+    struct_type = ir.LiteralStructType(field_types)
+    logger.info(f"Created struct {cls_node.name} with fields {fields.keys()}")
+    return StructType(struct_type, fields, total_size)
+
+
+def parse_struct_fields(cls_node):
+    """ Parse fields of a struct class node """
+    fields = {}
+
+    for item in cls_node.body:
+        if isinstance(item, ast.AnnAssign) and \
+           isinstance(item.target, ast.Name):
+            fields[item.target.id] = get_type_from_ann(item.annotation)
+        else:
+            logger.error(f"Unsupported struct field: {ast.dump(item)}")
+            raise TypeError(f"Unsupported field in {ast.dump(cls_node)}")
+    return fields
+
+
+def get_type_from_ann(annotation):
+    """ Convert an AST annotation node to an LLVM IR type for struct fields"""
+    if isinstance(annotation, ast.Call) and \
+       isinstance(annotation.func, ast.Name):
+        if annotation.func.id == "str":
+            # Char array
+            # Assumes constant integer argument
+            length = annotation.args[0].value
+            return ir.ArrayType(ir.IntType(8), length)
+    elif isinstance(annotation, ast.Name):
+        # Int type, written as c_int64, c_uint32, etc.
+        return ctypes_to_ir(annotation.id)
+
+    raise TypeError(f"Unsupported annotation type: {ast.dump(annotation)}")
+
+
+def calc_struct_size(field_types):
+    """ Calculate total size of the struct with alignment and padding """
+    curr_offset = 0
+    for ftype in field_types:
+        if isinstance(ftype, ir.IntType):
+            fsize = ftype.width // 8
+            alignment = fsize
+        elif isinstance(ftype, ir.ArrayType):
+            fsize = ftype.count * (ftype.element.width // 8)
+            alignment = ftype.element.width // 8
+        elif isinstance(ftype, ir.PointerType):
+            # We won't encounter this rn, but for the future
+            fsize = 8
+            alignment = 8
+        else:
+            raise TypeError(f"Unsupported field type: {ftype}")
+
+        padding = (alignment - (curr_offset % alignment)) % alignment
+        curr_offset += padding + fsize
+
+    final_padding = (8 - (curr_offset % 8)) % 8
+    return curr_offset + final_padding

pythonbpf/structs_pass.py

@@ -1,69 +0,0 @@
-import ast
-from llvmlite import ir
-from .type_deducer import ctypes_to_ir
-from . import dwarf_constants as dc
-
-structs_sym_tab = {}
-
-
-def structs_proc(tree, module, chunks):
-    for cls_node in chunks:
-        # Check if this class is a struct
-        is_struct = False
-        for decorator in cls_node.decorator_list:
-            if isinstance(decorator, ast.Name) and decorator.id == "struct":
-                is_struct = True
-                break
-        if is_struct:
-            print(f"Found BPF struct: {cls_node.name}")
-            process_bpf_struct(cls_node, module)
-            continue
-    return structs_sym_tab
-
-
-def process_bpf_struct(cls_node, module):
-    struct_name = cls_node.name
-    field_names = []
-    field_types = []
-
-    for item in cls_node.body:
-        if isinstance(item, ast.AnnAssign) and isinstance(item.target, ast.Name):
-            print(f"Field: {item.target.id}, Type: "
-                  f"{ast.dump(item.annotation)}")
-            field_names.append(item.target.id)
-            if isinstance(item.annotation, ast.Call) and isinstance(item.annotation.func, ast.Name) and item.annotation.func.id == "str":
-                # This is a char array with fixed length
-                # TODO: For now assuming str is always called with constant
-                field_types.append(ir.ArrayType(
-                    ir.IntType(8), item.annotation.args[0].value))
-            else:
-                field_types.append(ctypes_to_ir(item.annotation.id))
-
-    curr_offset = 0
-    for ftype in field_types:
-        if isinstance(ftype, ir.IntType):
-            fsize = ftype.width // 8
-            alignment = fsize
-        elif isinstance(ftype, ir.ArrayType):
-            fsize = ftype.count * (ftype.element.width // 8)
-            alignment = ftype.element.width // 8
-        elif isinstance(ftype, ir.PointerType):
-            fsize = 8
-            alignment = 8
-        else:
-            print(f"Unsupported field type in struct {struct_name}")
-            return
-        padding = (alignment - (curr_offset % alignment)) % alignment
-        curr_offset += padding
-        curr_offset += fsize
-    final_padding = (8 - (curr_offset % 8)) % 8
-    total_size = curr_offset + final_padding
-
-    struct_type = ir.LiteralStructType(field_types)
-    structs_sym_tab[struct_name] = {
-        "type": struct_type,
-        "fields": {name: idx for idx, name in enumerate(field_names)},
-        "size": total_size,
-        "field_types": field_types,
-    }
-    print(f"Created struct {struct_name} with fields {field_names}")