Use pointer arithmetic to resolve vmlinux struct fields

pythonbpf/allocation_pass.py

@@ -245,13 +245,17 @@ def _allocate_for_attribute(builder, var_name, rval, local_sym_tab, structs_sym_
 
             # 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]
+            base_ptr = builder.function.args[0]
+            local_sym_tab[
+                struct_var
+            ].var = base_ptr  # This is repurposing of var to store the pointer of the base type
+            local_sym_tab[struct_var].ir_type = field_ir
             # gep_result = builder.gep(
             #     base_ptr,
             #     [loaded_value],
             #     inbounds=False,  # Not using inbounds GEP
             # )
-            # print("DEBB", loaded_value, base_ptr, gep_result)
+            # print("DEBUG", loaded_value, base_ptr, gep_result)
             # Use i64 for allocation since that's what the global variable contains
 
             actual_ir_type = ir.IntType(64)

pythonbpf/vmlinux_parser/assignment_info.py

@@ -1,12 +1,11 @@
 from enum import Enum, auto
-from typing import Any, Dict, List, Optional, TypedDict
+from typing import Any, Dict, List, Optional
 from dataclasses import dataclass
 import llvmlite.ir as ir
 
 from pythonbpf.vmlinux_parser.dependency_node import Field
 
 
-@dataclass
 class AssignmentType(Enum):
     CONSTANT = auto()
     STRUCT = auto()
@@ -16,7 +15,7 @@ class AssignmentType(Enum):
 
 
 @dataclass
-class FunctionSignature(TypedDict):
+class FunctionSignature:
     return_type: str
     param_types: List[str]
     varargs: bool
@@ -24,7 +23,7 @@ class FunctionSignature(TypedDict):
 
 # Thew name of the assignment will be in the dict that uses this class
 @dataclass
-class AssignmentInfo(TypedDict):
+class AssignmentInfo:
     value_type: AssignmentType
     python_type: type
     value: Optional[Any]

pythonbpf/vmlinux_parser/vmlinux_exports_handler.py

@@ -37,16 +37,16 @@ class VmlinuxHandler:
         """Check if name is a vmlinux enum constant"""
         return (
             name in self.vmlinux_symtab
-            and self.vmlinux_symtab[name]["value_type"] == AssignmentType.CONSTANT
+            and self.vmlinux_symtab[name].value_type == AssignmentType.CONSTANT
         )
 
     def get_vmlinux_struct_type(self, name):
         """Check if name is a vmlinux struct type"""
         if (
             name in self.vmlinux_symtab
-            and self.vmlinux_symtab[name]["value_type"] == AssignmentType.STRUCT
+            and self.vmlinux_symtab[name].value_type == AssignmentType.STRUCT
         ):
-            return self.vmlinux_symtab[name]["python_type"]
+            return self.vmlinux_symtab[name].python_type
         else:
             raise ValueError(f"{name} is not a vmlinux struct type")
 
@@ -54,13 +54,13 @@ class VmlinuxHandler:
         """Check if name is a vmlinux struct"""
         return (
             name in self.vmlinux_symtab
-            and self.vmlinux_symtab[name]["value_type"] == AssignmentType.STRUCT
+            and self.vmlinux_symtab[name].value_type == AssignmentType.STRUCT
         )
 
     def handle_vmlinux_enum(self, name):
         """Handle vmlinux enum constants by returning LLVM IR constants"""
         if self.is_vmlinux_enum(name):
-            value = self.vmlinux_symtab[name]["value"]
+            value = self.vmlinux_symtab[name].value
             logger.info(f"Resolving vmlinux enum {name} = {value}")
             return ir.Constant(ir.IntType(64), value), ir.IntType(64)
         return None
@@ -68,7 +68,7 @@ class VmlinuxHandler:
     def get_vmlinux_enum_value(self, name):
         """Handle vmlinux enum constants by returning LLVM IR constants"""
         if self.is_vmlinux_enum(name):
-            value = self.vmlinux_symtab[name]["value"]
+            value = self.vmlinux_symtab[name].value
             logger.info(f"The value of vmlinux enum {name} = {value}")
             return value
         return None
@@ -98,24 +98,54 @@ class VmlinuxHandler:
                 python_type.__name__, field_name
             )
 
+            # Load the offset value
+            offset = builder.load(globvar_ir)
+
+            # Convert pointer to integer
+            i64_type = ir.IntType(64)
+            ptr_as_int = builder.ptrtoint(var_info.var, i64_type)
+
+            # Add the offset
+            field_addr = builder.add(ptr_as_int, offset)
+
+            # Convert back to pointer
+            i8_ptr_type = ir.IntType(8).as_pointer()
+            field_ptr_i8 = builder.inttoptr(field_addr, i8_ptr_type)
+            logger.info(f"field_ptr_i8: {field_ptr_i8}")
+
             # Return pointer to field and field type
-            return None
+            return field_ptr_i8, field_data
         else:
             raise RuntimeError("Variable accessed not found in symbol table")
 
     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"]
+            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"]
+            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]
+                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")
+
+    def get_field_index(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 list(
+                    self.vmlinux_symtab[vmlinux_struct_name].members.keys()
+                ).index(field_name)
             else:
                 raise ValueError(
                     f"Field {field_name} not found in vmlinux struct {vmlinux_struct_name}"

tests/c-form/struct_field_tests.bpf.c

@@ -21,7 +21,21 @@ int handle_setuid_entry(struct trace_event_raw_sys_enter *ctx) {
   // Access each argument separately with clear variable assignments
   long int arg0 = ctx->id;
   bpf_printk("args[0]: %d", arg0);
-
+  /*
+   * the IR to aim for is
+   * %2 = alloca ptr, align 8
+   * store ptr %0, ptr %2, align 8
+   * Above, %0 is the arg pointer
+   * %5 = load ptr, ptr %2, align 8
+   * %6 = getelementptr inbounds %struct.trace_event_raw_sys_enter, ptr %5, i32 0, i32 2
+   * %7 = load i64, ptr @"llvm.trace_event_raw_sys_enter:0:16$0:2:0", align 8
+   * %8 = bitcast ptr %5 to ptr
+   * %9 = getelementptr i8, ptr %8, i64 %7
+   * %10 = bitcast ptr %9 to ptr
+   * %11 = call ptr @llvm.bpf.passthrough.p0.p0(i32 0, ptr %10)
+   * %12 = load i64, ptr %11, align 8, !dbg !101
+   *
+   */
   return 0;
 }