Add passing test smp_processor_id.py for helpers

pythonbpf/helper/__init__.py

@@ -1,7 +1,18 @@
 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 .helpers import ktime, pid, deref, comm, probe_read_str, random, XDP_DROP, XDP_PASS
+from .helpers import (
+    ktime,
+    pid,
+    deref,
+    comm,
+    probe_read_str,
+    random,
+    probe_read,
+    smp_processor_id,
+    XDP_DROP,
+    XDP_PASS,
+)
 
 
 # Register the helper handler with expr module
@@ -66,6 +77,8 @@ __all__ = [
     "comm",
     "probe_read_str",
     "random",
+    "probe_read",
+    "smp_processor_id",
     "XDP_DROP",
     "XDP_PASS",
 ]

pythonbpf/helper/bpf_helper_handler.py

@@ -28,6 +28,7 @@ class BPFHelperID(Enum):
     BPF_KTIME_GET_NS = 5
     BPF_PRINTK = 6
     BPF_GET_PRANDOM_U32 = 7
+    BPF_GET_SMP_PROCESSOR_ID = 8
     BPF_GET_CURRENT_PID_TGID = 14
     BPF_GET_CURRENT_COMM = 16
     BPF_PERF_EVENT_OUTPUT = 25
@@ -476,23 +477,20 @@ def bpf_probe_read_emitter(
     if len(call.args) != 3:
         logger.warn("Expected 3 args for probe_read helper")
         return
-    dst_ptr, _ = get_ptr_from_arg(
-        call.args[0], func, module, builder, local_sym_tab, map_sym_tab, struct_sym_tab
+    dst_ptr = get_or_create_ptr_from_arg(
+        func, module, call.args[0], builder, local_sym_tab, map_sym_tab, struct_sym_tab
     )
-    size_val = (
-        get_int_value_from_arg(
-            call.args[1],
-            func,
-            module,
-            builder,
-            local_sym_tab,
-            map_sym_tab,
-            struct_sym_tab,
-        )
-        & 0xFFFFFFFF
+    size_val = get_int_value_from_arg(
+        call.args[1],
+        func,
+        module,
+        builder,
+        local_sym_tab,
+        map_sym_tab,
+        struct_sym_tab,
     )
-    src_ptr, _ = get_ptr_from_arg(
-        call.args[2], func, module, builder, local_sym_tab, map_sym_tab, 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
     )
     fn_type = ir.FunctionType(
         ir.IntType(64),
@@ -507,7 +505,7 @@ def bpf_probe_read_emitter(
         fn_ptr,
         [
             builder.bitcast(dst_ptr, ir.PointerType()),
-            ir.Constant(ir.IntType(32), size_val),
+            builder.trunc(size_val, ir.IntType(32)),
             builder.bitcast(src_ptr, ir.PointerType()),
         ],
         tail=False,
@@ -516,6 +514,29 @@ def bpf_probe_read_emitter(
     return result, ir.IntType(64)
 
 
+@HelperHandlerRegistry.register("smp_processor_id")
+def bpf_get_smp_processor_id_emitter(
+    call,
+    map_ptr,
+    module,
+    builder,
+    func,
+    local_sym_tab=None,
+    struct_sym_tab=None,
+    map_sym_tab=None,
+):
+    """
+    Emit LLVM IR for bpf_get_smp_processor_id helper function call.
+    """
+    helper_id = ir.Constant(ir.IntType(64), BPFHelperID.BPF_GET_SMP_PROCESSOR_ID.value)
+    fn_type = ir.FunctionType(ir.IntType(32), [], var_arg=False)
+    fn_ptr_type = ir.PointerType(fn_type)
+    fn_ptr = builder.inttoptr(helper_id, fn_ptr_type)
+    result = builder.call(fn_ptr, [], tail=False)
+    logger.info("Emitted bpf_get_smp_processor_id call")
+    return result, ir.IntType(32)
+
+
 def handle_helper_call(
     call,
     module,

pythonbpf/helper/helpers.py

@@ -32,6 +32,16 @@ def random():
     return ctypes.c_int32(0)
 
 
+def probe_read(dst, size, src):
+    """Safely read data from kernel memory"""
+    return ctypes.c_int64(0)
+
+
+def smp_processor_id():
+    """get the current CPU id"""
+    return ctypes.c_int32(0)
+
+
 XDP_ABORTED = ctypes.c_int64(0)
 XDP_DROP = ctypes.c_int64(1)
 XDP_PASS = ctypes.c_int64(2)

tests/passing_tests/helpers/bpf_probe_read.py

@@ -0,0 +1,29 @@
+from pythonbpf import bpf, section, bpfglobal, compile, struct
+from ctypes import c_void_p, c_int64, c_uint64, c_uint32
+from pythonbpf.helper import probe_read
+
+
+@bpf
+@struct
+class data_t:
+    pid: c_uint32
+    value: c_uint64
+
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_execve")
+def test_probe_read(ctx: c_void_p) -> c_int64:
+    """Test bpf_probe_read helper function"""
+    data = data_t()
+    probe_read(data.value, 8, ctx)
+    probe_read(data.pid, 4, ctx)
+    return 0
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+compile()

tests/passing_tests/helpers/smp_processor_id.py

@@ -0,0 +1,40 @@
+from pythonbpf import bpf, section, bpfglobal, compile, struct
+from ctypes import c_void_p, c_int64, c_uint32, c_uint64
+from pythonbpf.helper import smp_processor_id, ktime
+
+
+@bpf
+@struct
+class cpu_event_t:
+    cpu_id: c_uint32
+    timestamp: c_uint64
+
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_execve")
+def trace_with_cpu(ctx: c_void_p) -> c_int64:
+    """Test bpf_get_smp_processor_id helper function"""
+
+    # Get the current CPU ID
+    cpu = smp_processor_id()
+
+    # Print it
+    print(f"Running on CPU {cpu}")
+
+    # Use it in a struct
+    event = cpu_event_t()
+    event.cpu_id = smp_processor_id()
+    event.timestamp = ktime()
+
+    print(f"Event on CPU {event.cpu_id} at time {event.timestamp}")
+
+    return 0
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+compile()