diff --git a/pythonbpf/allocation_pass.py b/pythonbpf/allocation_pass.py index 49c787f..ad49118 100644 --- a/pythonbpf/allocation_pass.py +++ b/pythonbpf/allocation_pass.py @@ -1,10 +1,12 @@ import ast +import ctypes import logging from llvmlite import ir from dataclasses import dataclass from typing import Any from pythonbpf.helper import HelperHandlerRegistry +from pythonbpf.vmlinux_parser.dependency_node import Field from .expr import VmlinuxHandlerRegistry from pythonbpf.type_deducer import ctypes_to_ir @@ -60,21 +62,11 @@ def handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab): continue var_name = target.id - # Skip if already allocated if var_name in local_sym_tab: logger.debug(f"Variable {var_name} already allocated, skipping") continue - # When allocating a variable, check if it's a vmlinux struct type - if isinstance( - stmt.value, ast.Name - ) and VmlinuxHandlerRegistry.is_vmlinux_struct(stmt.value.id): - # Handle vmlinux struct allocation - # This requires more implementation - print(stmt.value) - pass - # Determine type and allocate based on rval if isinstance(rval, ast.Call): _allocate_for_call(builder, var_name, rval, local_sym_tab, structs_sym_tab) @@ -248,9 +240,40 @@ def _allocate_for_attribute(builder, var_name, rval, local_sym_tab, structs_sym_ logger.error(f"Struct variable '{struct_var}' not found") return - struct_type = local_sym_tab[struct_var].metadata + struct_type: type = local_sym_tab[struct_var].metadata if not struct_type or struct_type not in structs_sym_tab: - logger.error(f"Struct type '{struct_type}' not found") + if VmlinuxHandlerRegistry.is_vmlinux_struct(struct_type.__name__): + # Handle vmlinux struct field access + vmlinux_struct_name = struct_type.__name__ + if not VmlinuxHandlerRegistry.has_field(vmlinux_struct_name, field_name): + logger.error(f"Field '{field_name}' not found in vmlinux struct '{vmlinux_struct_name}'") + return + + field_type: tuple[ir.GlobalVariable, Field] = VmlinuxHandlerRegistry.get_field_type(vmlinux_struct_name, field_name) + field_ir, field = field_type + #TODO: For now, we only support integer type allocations. + + # loaded_value = builder.load(field_ir, align=8) + # #TODO: fatal flaw that this always assumes first argument of function to be the context of what this gets. + # base_ptr = builder.function.args[0] + # gep_result = builder.gep( + # base_ptr, + # [loaded_value], + # inbounds=False, # Not using inbounds GEP + # ) + # print("DEBB", loaded_value, base_ptr, gep_result) + # Use i64 for allocation since that's what the global variable contains + + actual_ir_type = ir.IntType(64) + + # Allocate with the actual IR type, not the GlobalVariable + var = _allocate_with_type(builder, var_name, actual_ir_type) + local_sym_tab[var_name] = LocalSymbol(var, actual_ir_type, field) + + logger.info(f"Pre-allocated {var_name} from vmlinux struct {vmlinux_struct_name}.{field_name}") + return + else: + logger.error(f"Struct type '{struct_type}' not found") return struct_info = structs_sym_tab[struct_type] diff --git a/pythonbpf/expr/vmlinux_registry.py b/pythonbpf/expr/vmlinux_registry.py index 7c25095..5bae334 100644 --- a/pythonbpf/expr/vmlinux_registry.py +++ b/pythonbpf/expr/vmlinux_registry.py @@ -52,3 +52,18 @@ class VmlinuxHandlerRegistry: if cls._handler is None: return None return cls._handler.get_vmlinux_struct_type(name) + + @classmethod + def has_field(cls, vmlinux_struct_name, field_name): + """Check if a vmlinux struct has a specific field""" + if cls._handler is None: + return False + return cls._handler.has_field(vmlinux_struct_name, field_name) + + @classmethod + def get_field_type(cls, vmlinux_struct_name, field_name): + """Get the type of a field in a vmlinux struct""" + if cls._handler is None: + return None + assert isinstance(cls._handler, VmlinuxHandler) + return cls._handler.get_field_type(vmlinux_struct_name, field_name) diff --git a/pythonbpf/functions/functions_pass.py b/pythonbpf/functions/functions_pass.py index d8c3b6e..a5449ae 100644 --- a/pythonbpf/functions/functions_pass.py +++ b/pythonbpf/functions/functions_pass.py @@ -20,7 +20,7 @@ from pythonbpf.assign_pass import ( from pythonbpf.allocation_pass import ( handle_assign_allocation, allocate_temp_pool, - create_targets_and_rvals, + create_targets_and_rvals, LocalSymbol, ) from .return_utils import handle_none_return, handle_xdp_return, is_xdp_name @@ -347,11 +347,12 @@ def process_func_body( resolved_type = VmlinuxHandlerRegistry.get_struct_type( context_type_name ) - context_type = {"type": ir.PointerType(resolved_type), "ptr": True} + context_type = LocalSymbol(None, ir.PointerType(resolved_type), resolved_type) else: try: resolved_type = ctypes_to_ir(context_type_name) - context_type = {"type": ir.PointerType(resolved_type), "ptr": True} + logger.error("THIS SHOULD NOT HAPPEN. I THINK. PROBABLY.") + context_type = LocalSymbol(None, ir.PointerType(resolved_type), resolved_type) except Exception: raise TypeError(f"Type '{context_type_name}' not declared") diff --git a/pythonbpf/vmlinux_parser/vmlinux_exports_handler.py b/pythonbpf/vmlinux_parser/vmlinux_exports_handler.py index e244cfc..f68822a 100644 --- a/pythonbpf/vmlinux_parser/vmlinux_exports_handler.py +++ b/pythonbpf/vmlinux_parser/vmlinux_exports_handler.py @@ -98,3 +98,23 @@ class VmlinuxHandler: # Return pointer to field and field type return None return None + + def has_field(self, struct_name, field_name): + """Check if a vmlinux struct has a specific field""" + if self.is_vmlinux_struct(struct_name): + python_type = self.vmlinux_symtab[struct_name]["python_type"] + return hasattr(python_type, field_name) + return False + + def get_field_type(self, vmlinux_struct_name, field_name): + """Get the type of a field in a vmlinux struct""" + if self.is_vmlinux_struct(vmlinux_struct_name): + python_type = self.vmlinux_symtab[vmlinux_struct_name]["python_type"] + if hasattr(python_type, field_name): + return self.vmlinux_symtab[vmlinux_struct_name]["members"][field_name] + else: + raise ValueError( + f"Field {field_name} not found in vmlinux struct {vmlinux_struct_name}" + ) + else: + raise ValueError(f"{vmlinux_struct_name} is not a vmlinux struct") diff --git a/tests/c-form/Makefile b/tests/c-form/Makefile index 6be0827..64ff900 100644 --- a/tests/c-form/Makefile +++ b/tests/c-form/Makefile @@ -1,5 +1,5 @@ BPF_CLANG := clang -CFLAGS := -O2 -emit-llvm -target bpf -c +CFLAGS := -O0 -emit-llvm -target bpf -c SRC := $(wildcard *.bpf.c) LL := $(SRC:.bpf.c=.bpf.ll) diff --git a/tests/failing_tests/vmlinux/struct_field_access.py b/tests/failing_tests/vmlinux/struct_field_access.py index 4ff90c2..1d6e85e 100644 --- a/tests/failing_tests/vmlinux/struct_field_access.py +++ b/tests/failing_tests/vmlinux/struct_field_access.py @@ -17,7 +17,7 @@ def hello_world(ctx: struct_trace_event_raw_sys_enter) -> c_int64: a = 2 + TASK_COMM_LEN + TASK_COMM_LEN b = ctx.id print(f"Hello, World{TASK_COMM_LEN} and {a}") - print(f"This is context field {b}") + # print(f"This is context field {b}") return c_int64(TASK_COMM_LEN + 2) diff --git a/tests/passing_tests/assign/comprehensive.py b/tests/passing_tests/assign/comprehensive.py index aeb6898..47a38ca 100644 --- a/tests/passing_tests/assign/comprehensive.py +++ b/tests/passing_tests/assign/comprehensive.py @@ -1,4 +1,4 @@ -from pythonbpf import bpf, map, section, bpfglobal, compile, struct +from pythonbpf import bpf, map, section, bpfglobal, compile, struct, compile_to_ir from ctypes import c_void_p, c_int64, c_int32, c_uint64 from pythonbpf.maps import HashMap from pythonbpf.helper import ktime @@ -70,5 +70,5 @@ def hello_world(ctx: c_void_p) -> c_int64: def LICENSE() -> str: return "GPL" - -compile() +compile_to_ir("comprehensive.py", "comprehensive.ll") +# compile()