Fix call to get_or_create_ptr_from_arg for probe_read_str

pythonbpf/functions/functions_pass.py

@@ -62,7 +62,9 @@ def count_temps_in_call(call_node, local_sym_tab):
         ):
             continue
         param_type = HelperHandlerRegistry.get_param_type(func_name, arg_idx)
-        count[param_type] = count.get(param_type, 0) + 1
+        if isinstance(param_type, ir.PointerType):
+            pointee_type = param_type.pointee
+            count[pointee_type] = count.get(pointee_type, 0) + 1
 
     return count
 

pythonbpf/helper/bpf_helper_handler.py

@@ -8,7 +8,6 @@ from .helper_utils import (
     get_flags_val,
     get_data_ptr_and_size,
     get_buffer_ptr_and_size,
-    get_char_array_ptr_and_size,
     get_ptr_from_arg,
     get_int_value_from_arg,
 )
@@ -37,7 +36,11 @@ class BPFHelperID(Enum):
     BPF_PROBE_READ_KERNEL_STR = 115
 
 
-@HelperHandlerRegistry.register("ktime")
+@HelperHandlerRegistry.register(
+    "ktime",
+    param_types=[],
+    return_type=ir.IntType(64),
+)
 def bpf_ktime_get_ns_emitter(
     call,
     map_ptr,
@@ -60,7 +63,11 @@ def bpf_ktime_get_ns_emitter(
     return result, ir.IntType(64)
 
 
-@HelperHandlerRegistry.register("lookup")
+@HelperHandlerRegistry.register(
+    "lookup",
+    param_types=[ir.PointerType(ir.IntType(64))],
+    return_type=ir.PointerType(ir.IntType(64)),
+)
 def bpf_map_lookup_elem_emitter(
     call,
     map_ptr,
@@ -102,6 +109,7 @@ def bpf_map_lookup_elem_emitter(
     return result, ir.PointerType()
 
 
+# NOTE: This has special handling so we won't reflect the signature here.
 @HelperHandlerRegistry.register("print")
 def bpf_printk_emitter(
     call,
@@ -150,7 +158,15 @@ def bpf_printk_emitter(
     return True
 
 
-@HelperHandlerRegistry.register("update")
+@HelperHandlerRegistry.register(
+    "update",
+    param_types=[
+        ir.PointerType(ir.IntType(64)),
+        ir.PointerType(ir.IntType(64)),
+        ir.IntType(64),
+    ],
+    return_type=ir.PointerType(ir.IntType(64)),
+)
 def bpf_map_update_elem_emitter(
     call,
     map_ptr,
@@ -205,7 +221,11 @@ def bpf_map_update_elem_emitter(
     return result, None
 
 
-@HelperHandlerRegistry.register("delete")
+@HelperHandlerRegistry.register(
+    "delete",
+    param_types=[ir.PointerType(ir.IntType(64))],
+    return_type=ir.PointerType(ir.IntType(64)),
+)
 def bpf_map_delete_elem_emitter(
     call,
     map_ptr,
@@ -245,7 +265,11 @@ def bpf_map_delete_elem_emitter(
     return result, None
 
 
-@HelperHandlerRegistry.register("comm")
+@HelperHandlerRegistry.register(
+    "comm",
+    param_types=[ir.PointerType(ir.IntType(8))],
+    return_type=ir.IntType(64),
+)
 def bpf_get_current_comm_emitter(
     call,
     map_ptr,
@@ -302,7 +326,11 @@ def bpf_get_current_comm_emitter(
     return result, None
 
 
-@HelperHandlerRegistry.register("pid")
+@HelperHandlerRegistry.register(
+    "pid",
+    param_types=[],
+    return_type=ir.IntType(64),
+)
 def bpf_get_current_pid_tgid_emitter(
     call,
     map_ptr,
@@ -330,7 +358,11 @@ def bpf_get_current_pid_tgid_emitter(
     return pid, ir.IntType(64)
 
 
-@HelperHandlerRegistry.register("output")
+@HelperHandlerRegistry.register(
+    "output",
+    param_types=[ir.PointerType(ir.IntType(8))],
+    return_type=ir.IntType(64),
+)
 def bpf_perf_event_output_handler(
     call,
     map_ptr,
@@ -405,7 +437,14 @@ def emit_probe_read_kernel_str_call(builder, dst_ptr, dst_size, src_ptr):
     return result
 
 
-@HelperHandlerRegistry.register("probe_read_str")
+@HelperHandlerRegistry.register(
+    "probe_read_str",
+    param_types=[
+        ir.PointerType(ir.IntType(8)),
+        ir.PointerType(ir.IntType(8)),
+    ],
+    return_type=ir.IntType(64),
+)
 def bpf_probe_read_kernel_str_emitter(
     call,
     map_ptr,
@@ -424,8 +463,8 @@ def bpf_probe_read_kernel_str_emitter(
         )
 
     # Get destination buffer (char array -> i8*)
-    dst_ptr, dst_size = get_char_array_ptr_and_size(
-        call.args[0], builder, local_sym_tab, struct_sym_tab
+    dst_ptr, dst_size = get_or_create_ptr_from_arg(
+        func, module, call.args[0], builder, local_sym_tab, map_sym_tab, struct_sym_tab
     )
 
     # Get source pointer (evaluate expression)
@@ -440,7 +479,11 @@ def bpf_probe_read_kernel_str_emitter(
     return result, ir.IntType(64)
 
 
-@HelperHandlerRegistry.register("random")
+@HelperHandlerRegistry.register(
+    "random",
+    param_types=[],
+    return_type=ir.IntType(32),
+)
 def bpf_get_prandom_u32_emitter(
     call,
     map_ptr,
@@ -462,7 +505,15 @@ def bpf_get_prandom_u32_emitter(
     return result, ir.IntType(32)
 
 
-@HelperHandlerRegistry.register("probe_read")
+@HelperHandlerRegistry.register(
+    "probe_read",
+    param_types=[
+        ir.PointerType(ir.IntType(8)),
+        ir.IntType(32),
+        ir.PointerType(ir.IntType(8)),
+    ],
+    return_type=ir.IntType(64),
+)
 def bpf_probe_read_emitter(
     call,
     map_ptr,
@@ -481,7 +532,14 @@ def bpf_probe_read_emitter(
         logger.warn("Expected 3 args for probe_read helper")
         return
     dst_ptr = get_or_create_ptr_from_arg(
-        func, module, call.args[0], builder, local_sym_tab, map_sym_tab, struct_sym_tab
+        func,
+        module,
+        call.args[0],
+        builder,
+        local_sym_tab,
+        map_sym_tab,
+        struct_sym_tab,
+        ir.IntType(8),
     )
     size_val = get_int_value_from_arg(
         call.args[1],
@@ -493,7 +551,14 @@ def bpf_probe_read_emitter(
         struct_sym_tab,
     )
     src_ptr = get_or_create_ptr_from_arg(
-        func, module, call.args[2], builder, local_sym_tab, map_sym_tab, struct_sym_tab
+        func,
+        module,
+        call.args[2],
+        builder,
+        local_sym_tab,
+        map_sym_tab,
+        struct_sym_tab,
+        ir.IntType(8),
     )
     fn_type = ir.FunctionType(
         ir.IntType(64),
@@ -517,7 +582,11 @@ def bpf_probe_read_emitter(
     return result, ir.IntType(64)
 
 
-@HelperHandlerRegistry.register("smp_processor_id")
+@HelperHandlerRegistry.register(
+    "smp_processor_id",
+    param_types=[],
+    return_type=ir.IntType(32),
+)
 def bpf_get_smp_processor_id_emitter(
     call,
     map_ptr,
@@ -540,7 +609,11 @@ def bpf_get_smp_processor_id_emitter(
     return result, ir.IntType(32)
 
 
-@HelperHandlerRegistry.register("uid")
+@HelperHandlerRegistry.register(
+    "uid",
+    param_types=[],
+    return_type=ir.IntType(64),
+)
 def bpf_get_current_uid_gid_emitter(
     call,
     map_ptr,
@@ -569,7 +642,12 @@ def bpf_get_current_uid_gid_emitter(
 
 @HelperHandlerRegistry.register(
     "skb_store_bytes",
-    param_types=[ir.IntType(32), ir.PointerType(), ir.IntType(32), ir.IntType(64)],
+    param_types=[
+        ir.IntType(32),
+        ir.PointerType(ir.IntType(8)),
+        ir.IntType(32),
+        ir.IntType(64),
+    ],
     return_type=ir.IntType(64),
 )
 def bpf_skb_store_bytes_emitter(
@@ -608,7 +686,6 @@ def bpf_skb_store_bytes_emitter(
         builder,
         local_sym_tab,
         map_sym_tab,
-        args_signature[1],
         struct_sym_tab,
     )
     from_ptr = get_or_create_ptr_from_arg(
@@ -618,8 +695,8 @@ def bpf_skb_store_bytes_emitter(
         builder,
         local_sym_tab,
         map_sym_tab,
-        args_signature[2],
         struct_sym_tab,
+        args_signature[2],
     )
     len_val = get_int_value_from_arg(
         call.args[2],
@@ -628,7 +705,6 @@ def bpf_skb_store_bytes_emitter(
         builder,
         local_sym_tab,
         map_sym_tab,
-        args_signature[3],
         struct_sym_tab,
     )
     if len(call.args) == 4:

pythonbpf/helper/helper_utils.py

@@ -18,7 +18,7 @@ class ScratchPoolManager:
 
     @property
     def counter(self):
-        return sum(self._counter.values())
+        return sum(self._counters.values())
 
     def reset(self):
         self._counters.clear()
@@ -92,18 +92,56 @@ def get_or_create_ptr_from_arg(
     builder,
     local_sym_tab,
     map_sym_tab,
-    expected_type=None,
     struct_sym_tab=None,
+    expected_type=None,
 ):
     """Extract or create pointer from the call arguments."""
 
+    logger.info(f"Getting pointer from arg: {ast.dump(arg)}")
+    sz = None
     if isinstance(arg, ast.Name):
         # Stack space is already allocated
         ptr = get_var_ptr_from_name(arg.id, local_sym_tab)
     elif isinstance(arg, ast.Constant) and isinstance(arg.value, int):
+        int_width = 64  # Default to i64
         if expected_type and isinstance(expected_type, ir.IntType):
             int_width = expected_type.width
         ptr = create_int_constant_ptr(arg.value, builder, local_sym_tab, int_width)
+    elif isinstance(arg, ast.Attribute):
+        # A struct field
+        struct_name = arg.value.id
+        field_name = arg.attr
+
+        if not local_sym_tab or struct_name not in local_sym_tab:
+            raise ValueError(f"Struct '{struct_name}' not found")
+
+        struct_type = local_sym_tab[struct_name].metadata
+        if not struct_sym_tab or struct_type not in struct_sym_tab:
+            raise ValueError(f"Struct type '{struct_type}' not found")
+
+        struct_info = struct_sym_tab[struct_type]
+        if field_name not in struct_info.fields:
+            raise ValueError(
+                f"Field '{field_name}' not found in struct '{struct_name}'"
+            )
+
+        field_type = struct_info.field_type(field_name)
+        struct_ptr = local_sym_tab[struct_name].var
+
+        # Special handling for char arrays
+        if (
+            isinstance(field_type, ir.ArrayType)
+            and isinstance(field_type.element, ir.IntType)
+            and field_type.element.width == 8
+        ):
+            ptr, sz = get_char_array_ptr_and_size(
+                arg, builder, local_sym_tab, struct_sym_tab
+            )
+            if not ptr:
+                raise ValueError("Failed to get char array pointer from struct field")
+        else:
+            ptr = struct_info.gep(builder, struct_ptr, field_name)
+
     else:
         # NOTE: For any integer expression reaching this branch, it is probably a struct field or a binop
         # Evaluate the expression and store the result in a temp variable
@@ -123,6 +161,10 @@ def get_or_create_ptr_from_arg(
             val = builder.trunc(val, expected_type)
         builder.store(val, ptr)
 
+    # NOTE: For char arrays, also return size
+    if sz:
+        return ptr, sz
+
     return ptr
 
 

pythonbpf/helper/printk_formatter.py

@@ -4,6 +4,7 @@ import logging
 from llvmlite import ir
 from pythonbpf.expr import eval_expr, get_base_type_and_depth, deref_to_depth
 from pythonbpf.expr.vmlinux_registry import VmlinuxHandlerRegistry
+from pythonbpf.helper.helper_utils import get_char_array_ptr_and_size
 
 logger = logging.getLogger(__name__)
 
@@ -219,7 +220,7 @@ def _prepare_expr_args(expr, func, module, builder, local_sym_tab, struct_sym_ta
     """Evaluate and prepare an expression to use as an arg for bpf_printk."""
 
     # Special case: struct field char array needs pointer to first element
-    char_array_ptr = _get_struct_char_array_ptr(
+    char_array_ptr, _ = get_char_array_ptr_and_size(
         expr, builder, local_sym_tab, struct_sym_tab
     )
     if char_array_ptr:
@@ -242,52 +243,6 @@ def _prepare_expr_args(expr, func, module, builder, local_sym_tab, struct_sym_ta
         return ir.Constant(ir.IntType(64), 0)
 
 
-def _get_struct_char_array_ptr(expr, builder, local_sym_tab, struct_sym_tab):
-    """Get pointer to first element of char array in struct field, or None."""
-    if not (isinstance(expr, ast.Attribute) and isinstance(expr.value, ast.Name)):
-        return None
-
-    var_name = expr.value.id
-    field_name = expr.attr
-
-    # Check if it's a valid struct field
-    if not (
-        local_sym_tab
-        and var_name in local_sym_tab
-        and struct_sym_tab
-        and local_sym_tab[var_name].metadata in struct_sym_tab
-    ):
-        return None
-
-    struct_type = local_sym_tab[var_name].metadata
-    struct_info = struct_sym_tab[struct_type]
-
-    if field_name not in struct_info.fields:
-        return None
-
-    field_type = struct_info.field_type(field_name)
-
-    # Check if it's a char array
-    is_char_array = (
-        isinstance(field_type, ir.ArrayType)
-        and isinstance(field_type.element, ir.IntType)
-        and field_type.element.width == 8
-    )
-
-    if not is_char_array:
-        return None
-
-    # Get field pointer and GEP to first element: [N x i8]* -> i8*
-    struct_ptr = local_sym_tab[var_name].var
-    field_ptr = struct_info.gep(builder, struct_ptr, field_name)
-
-    return builder.gep(
-        field_ptr,
-        [ir.Constant(ir.IntType(32), 0), ir.Constant(ir.IntType(32), 0)],
-        inbounds=True,
-    )
-
-
 def _handle_pointer_arg(val, func, builder):
     """Convert pointer type for bpf_printk."""
     target, depth = get_base_type_and_depth(val.type)