fix CO-RE read for cast structs

Author: varun-r-mallya

pythonbpf/assign_pass.py

@@ -1,5 +1,6 @@
 import ast
 import logging
+from inspect import isclass
 
 from llvmlite import ir
 from pythonbpf.expr import eval_expr
@@ -149,11 +150,26 @@ def handle_variable_assignment(
         return False
 
     val, val_type = val_result
-    logger.info(f"Evaluated value for {var_name}: {val} of type {val_type}, {var_type}")
+    logger.info(f"Evaluated value for {var_name}: {val} of type {val_type}, expected {var_type}")
+
     if val_type != var_type:
-        # if isclass(val_type) and (val_type.__module__ == "vmlinux"):
-        #     logger.info("Handling typecast to vmlinux struct")
-        #     print(val_type, var_type)
+        # Handle vmlinux struct pointers - they're represented as Python classes but are i64 pointers
+        if isclass(val_type) and (val_type.__module__ == "vmlinux"):
+            logger.info("Handling vmlinux struct pointer assignment")
+            # vmlinux struct pointers: val is a pointer, need to convert to i64
+            if isinstance(var_type, ir.IntType) and var_type.width == 64:
+                # Convert pointer to i64 using ptrtoint
+                if isinstance(val.type, ir.PointerType):
+                    val = builder.ptrtoint(val, ir.IntType(64))
+                    logger.info(f"Converted vmlinux struct pointer to i64 using ptrtoint")
+                builder.store(val, var_ptr)
+                logger.info(f"Assigned vmlinux struct pointer to {var_name} (i64)")
+                return True
+            else:
+                logger.error(
+                    f"Type mismatch: vmlinux struct pointer requires i64, got {var_type}"
+                )
+                return False
         if isinstance(val_type, Field):
             logger.info("Handling assignment to struct field")
             # Special handling for struct_xdp_md i32 fields that are zero-extended to i64

pythonbpf/helper/__init__.py

@@ -1,6 +1,6 @@
 from .helper_registry import HelperHandlerRegistry
 from .helper_utils import reset_scratch_pool
-from .bpf_helper_handler import handle_helper_call, emit_probe_read_kernel_str_call
+from .bpf_helper_handler import handle_helper_call, emit_probe_read_kernel_str_call, emit_probe_read_kernel_call
 from .helpers import (
     ktime,
     pid,
@@ -74,6 +74,7 @@ def helper_call_handler(
     "reset_scratch_pool",
     "handle_helper_call",
     "emit_probe_read_kernel_str_call",
+    "emit_probe_read_kernel_call",
     "ktime",
     "pid",
     "deref",

pythonbpf/vmlinux_parser/vmlinux_exports_handler.py

@@ -115,11 +115,11 @@ def handle_vmlinux_struct_field(
                 struct_name = python_type.__name__
                 globvar_ir, field_data = self.get_field_type(struct_name, field_name)
 
-                # Handle cast struct field access (use standard GEP)
+                # Handle cast struct field access (use bpf_probe_read_kernel)
                 # Load the struct pointer from the local variable
                 struct_ptr = builder.load(var_info.var)
 
-                # Use standard GEP for non-context struct field access
+                # Use bpf_probe_read_kernel for non-context struct field access
                 field_value = self.load_struct_field(
                     builder, struct_ptr, globvar_ir, field_data, struct_name
                 )
@@ -133,7 +133,7 @@ def load_struct_field(
         builder, struct_ptr_int, offset_global, field_data, struct_name=None
     ):
         """
-        Generate LLVM IR to load a field from a regular (non-context) struct using standard GEP.
+        Generate LLVM IR to load a field from a regular (non-context) struct using bpf_probe_read_kernel.
 
         Args:
             builder: llvmlite IRBuilder instance
@@ -159,7 +159,8 @@ def load_struct_field(
             inbounds=False,
         )
 
-        # Determine the appropriate IR type based on field information
+        # Determine the appropriate field size based on field information
+        field_size_bytes = 8  # Default to 8 bytes (64-bit)
         int_width = 64  # Default to 64-bit
         needs_zext = False
 
@@ -172,7 +173,7 @@ def load_struct_field(
 
                     if field_size_bits in [8, 16, 32, 64]:
                         int_width = field_size_bits
-                        logger.info(f"Determined field size: {int_width} bits")
+                        logger.info(f"Determined field size: {int_width} bits ({field_size_bytes} bytes)")
 
                         # Special handling for struct_xdp_md i32 fields
                         if struct_name == "struct_xdp_md" and int_width == 32:
@@ -195,16 +196,31 @@ def load_struct_field(
                     field_data.ctype_complex_type, ctypes._Pointer
                 ):
                     int_width = 64  # Pointers are always 64-bit
+                    field_size_bytes = 8
                     logger.info("Field is a pointer type, using 64 bits")
                 else:
                     logger.warning("Complex vmlinux field type, using default 64 bits")
 
-        # Bitcast to appropriate pointer type based on determined width
-        ptr_type = ir.PointerType(ir.IntType(int_width))
-        typed_ptr = builder.bitcast(field_ptr, ptr_type)
+        # Allocate local storage for the field value
+        local_storage = builder.alloca(ir.IntType(int_width))
+        local_storage_i8_ptr = builder.bitcast(local_storage, i8_ptr_type)
+
+        # Use bpf_probe_read_kernel to safely read the field
+        # This generates:
+        #   %gep = getelementptr i8, ptr %struct_ptr, i64 %offset (already done above as field_ptr)
+        #   %passed = tail call ptr @llvm.bpf.passthrough.p0.p0(i32 2, ptr %gep)
+        #   %result = call i64 inttoptr (i64 113 to ptr)(ptr %local_storage, i32 %size, ptr %passed)
+        from pythonbpf.helper import emit_probe_read_kernel_call
+
+        result = emit_probe_read_kernel_call(
+            builder,
+            local_storage_i8_ptr,
+            field_size_bytes,
+            field_ptr
+        )
 
-        # Load the value
-        value = builder.load(typed_ptr)
+        # Load the value from local storage
+        value = builder.load(local_storage)
 
         # Zero-extend i32 to i64 if needed
         if needs_zext:

tests/failing_tests/vmlinux/requests.py

@@ -12,7 +12,8 @@ def example(ctx: struct_pt_regs) -> c_int64:
     d = req.__data_len
     b = ctx.r12
     c = req.timeout
-    print(f"data length {d} and {c} and {a} and {b}")
+    print(f"data length {d} and {c} and {a}")
+    print(f"ctx arg {b}")
     return c_int64(0)