Introduce HelperSignature in HelperHandlerRegistry

pythonbpf/helper/__init__.py

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

pythonbpf/helper/bpf_helper_handler.py

@@ -10,6 +10,7 @@ from .helper_utils import (
     get_buffer_ptr_and_size,
     get_char_array_ptr_and_size,
     get_ptr_from_arg,
+    get_int_value_from_arg,
 )
 from .printk_formatter import simple_string_print, handle_fstring_print
 
@@ -23,9 +24,14 @@ class BPFHelperID(Enum):
     BPF_MAP_LOOKUP_ELEM = 1
     BPF_MAP_UPDATE_ELEM = 2
     BPF_MAP_DELETE_ELEM = 3
+    BPF_PROBE_READ = 4
     BPF_KTIME_GET_NS = 5
     BPF_PRINTK = 6
+    BPF_GET_PRANDOM_U32 = 7
+    BPF_GET_SMP_PROCESSOR_ID = 8
+    BPF_SKB_STORE_BYTES = 9
     BPF_GET_CURRENT_PID_TGID = 14
+    BPF_GET_CURRENT_UID_GID = 15
     BPF_GET_CURRENT_COMM = 16
     BPF_PERF_EVENT_OUTPUT = 25
     BPF_PROBE_READ_KERNEL_STR = 115
@@ -318,6 +324,7 @@ def bpf_get_current_pid_tgid_emitter(
     result = builder.call(fn_ptr, [], tail=False)
 
     # Extract the lower 32 bits (PID) using bitwise AND with 0xFFFFFFFF
+    # TODO: return both PID and TGID if we end up needing TGID somewhere
     mask = ir.Constant(ir.IntType(64), 0xFFFFFFFF)
     pid = builder.and_(result, mask)
     return pid, ir.IntType(64)
@@ -433,6 +440,211 @@ def bpf_probe_read_kernel_str_emitter(
     return result, ir.IntType(64)
 
 
+@HelperHandlerRegistry.register("random")
+def bpf_get_prandom_u32_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_prandom_u32 helper function call.
+    """
+    helper_id = ir.Constant(ir.IntType(64), BPFHelperID.BPF_GET_PRANDOM_U32.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)
+    return result, ir.IntType(32)
+
+
+@HelperHandlerRegistry.register("probe_read")
+def bpf_probe_read_emitter(
+    call,
+    map_ptr,
+    module,
+    builder,
+    func,
+    local_sym_tab=None,
+    struct_sym_tab=None,
+    map_sym_tab=None,
+):
+    """
+    Emit LLVM IR for bpf_probe_read helper function
+    """
+
+    if len(call.args) != 3:
+        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
+    )
+    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_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),
+        [ir.PointerType(), ir.IntType(32), ir.PointerType()],
+        var_arg=False,
+    )
+    fn_ptr = builder.inttoptr(
+        ir.Constant(ir.IntType(64), BPFHelperID.BPF_PROBE_READ.value),
+        ir.PointerType(fn_type),
+    )
+    result = builder.call(
+        fn_ptr,
+        [
+            builder.bitcast(dst_ptr, ir.PointerType()),
+            builder.trunc(size_val, ir.IntType(32)),
+            builder.bitcast(src_ptr, ir.PointerType()),
+        ],
+        tail=False,
+    )
+    logger.info(f"Emitted bpf_probe_read (size={size_val})")
+    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)
+
+
+@HelperHandlerRegistry.register("uid")
+def bpf_get_current_uid_gid_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_current_uid_gid helper function call.
+    """
+    helper_id = ir.Constant(ir.IntType(64), BPFHelperID.BPF_GET_CURRENT_UID_GID.value)
+    fn_type = ir.FunctionType(ir.IntType(64), [], 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)
+
+    # Extract the lower 32 bits (UID) using bitwise AND with 0xFFFFFFFF
+    # TODO: return both UID and GID if we end up needing GID somewhere
+    mask = ir.Constant(ir.IntType(64), 0xFFFFFFFF)
+    pid = builder.and_(result, mask)
+    return pid, ir.IntType(64)
+
+
+@HelperHandlerRegistry.register("skb_store_bytes")
+def bpf_skb_store_bytes_emitter(
+    call,
+    map_ptr,
+    module,
+    builder,
+    func,
+    local_sym_tab=None,
+    struct_sym_tab=None,
+    map_sym_tab=None,
+):
+    """
+    Emit LLVM IR for bpf_skb_store_bytes helper function call.
+    Expected call signature: skb_store_bytes(skb, offset, from, len, flags)
+    """
+
+    if len(call.args) not in (3, 4):
+        raise ValueError(
+            f"skb_store_bytes expects 3 or 4 args (offset, from, len, flags), got {len(call.args)}"
+        )
+
+    skb_ptr = func.args[0]  # First argument to the function is skb
+    offset_val = get_int_value_from_arg(
+        call.args[0],
+        func,
+        module,
+        builder,
+        local_sym_tab,
+        map_sym_tab,
+        struct_sym_tab,
+    )
+    from_ptr = get_or_create_ptr_from_arg(
+        func, module, call.args[1], builder, local_sym_tab, map_sym_tab, struct_sym_tab
+    )
+    len_val = get_int_value_from_arg(
+        call.args[2],
+        func,
+        module,
+        builder,
+        local_sym_tab,
+        map_sym_tab,
+        struct_sym_tab,
+    )
+    if len(call.args) == 4:
+        flags_val = get_flags_val(call.args[3], builder, local_sym_tab)
+    else:
+        flags_val = 0
+    flags = ir.Constant(ir.IntType(64), flags_val)
+    fn_type = ir.FunctionType(
+        ir.IntType(64),
+        [
+            ir.PointerType(),  # skb
+            ir.IntType(32),  # offset
+            ir.PointerType(),  # from
+            ir.IntType(32),  # len
+            ir.IntType(64),  # flags
+        ],
+        var_arg=False,
+    )
+    fn_ptr = builder.inttoptr(
+        ir.Constant(ir.IntType(64), BPFHelperID.BPF_SKB_STORE_BYTES.value),
+        ir.PointerType(fn_type),
+    )
+    result = builder.call(
+        fn_ptr,
+        [
+            builder.bitcast(skb_ptr, ir.PointerType()),
+            builder.trunc(offset_val, ir.IntType(32)),
+            builder.bitcast(from_ptr, ir.PointerType()),
+            builder.trunc(len_val, ir.IntType(32)),
+            flags,
+        ],
+        tail=False,
+    )
+    logger.info("Emitted bpf_skb_store_bytes call")
+    return result, ir.IntType(64)
+
+
 def handle_helper_call(
     call,
     module,

pythonbpf/helper/helper_registry.py

@@ -1,17 +1,31 @@
+from dataclasses import dataclass
+from llvmlite import ir
 from typing import Callable
 
 
+@dataclass
+class HelperSignature:
+    """Signature of a BPF helper function"""
+
+    arg_types: list[ir.Type]
+    return_type: ir.Type
+    func: Callable
+
+
 class HelperHandlerRegistry:
     """Registry for BPF helpers"""
 
-    _handlers: dict[str, Callable] = {}
+    _handlers: dict[str, HelperSignature] = {}
 
     @classmethod
-    def register(cls, helper_name):
+    def register(cls, helper_name, param_types=None, return_type=None):
         """Decorator to register a handler function for a helper"""
 
         def decorator(func):
-            cls._handlers[helper_name] = func
+            helper_sig = HelperSignature(
+                arg_types=param_types, return_type=return_type, func=func
+            )
+            cls._handlers[helper_name] = helper_sig
             return func
 
         return decorator
@@ -19,7 +33,8 @@ class HelperHandlerRegistry:
     @classmethod
     def get_handler(cls, helper_name):
         """Get the handler function for a helper"""
-        return cls._handlers.get(helper_name)
+        handler = cls._handlers.get(helper_name)
+        return handler.func if handler else None
 
     @classmethod
     def has_handler(cls, helper_name):

pythonbpf/helper/helper_utils.py

@@ -274,3 +274,23 @@ def get_ptr_from_arg(
         raise ValueError(f"Expected pointer type, got {val_type}")
 
     return val, val_type
+
+
+def get_int_value_from_arg(
+    arg, func, module, builder, local_sym_tab, map_sym_tab, struct_sym_tab
+):
+    """Evaluate argument and return integer value"""
+
+    result = eval_expr(
+        func, module, builder, arg, local_sym_tab, map_sym_tab, struct_sym_tab
+    )
+
+    if not result:
+        raise ValueError("Failed to evaluate argument")
+
+    val, val_type = result
+
+    if not isinstance(val_type, ir.IntType):
+        raise ValueError(f"Expected integer type, got {val_type}")
+
+    return val

pythonbpf/helper/helpers.py

@@ -27,6 +27,31 @@ def probe_read_str(dst, src):
     return ctypes.c_int64(0)
 
 
+def random():
+    """get a pseudorandom u32 number"""
+    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)
+
+
+def uid():
+    """get current user id"""
+    return ctypes.c_int32(0)
+
+
+def skb_store_bytes(offset, from_buf, size, flags=0):
+    """store bytes into a socket buffer"""
+    return ctypes.c_int64(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/prandom.py

@@ -0,0 +1,25 @@
+from pythonbpf import bpf, bpfglobal, section, BPF, trace_pipe
+from ctypes import c_void_p, c_int64
+from pythonbpf.helper import random
+
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_clone")
+def hello_world(ctx: c_void_p) -> c_int64:
+    r = random()
+    print(f"Hello, World!, {r}")
+    return 0  # type: ignore [return-value]
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+# Compile and load
+b = BPF()
+b.load()
+b.attach_all()
+
+trace_pipe()

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()

tests/passing_tests/helpers/uid_gid.py

@@ -0,0 +1,31 @@
+from pythonbpf import bpf, section, bpfglobal, compile
+from ctypes import c_void_p, c_int64
+from pythonbpf.helper import uid, pid
+
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_execve")
+def filter_by_user(ctx: c_void_p) -> c_int64:
+    """Filter events by specific user ID"""
+
+    current_uid = uid()
+
+    # Only trace root user (UID 0)
+    if current_uid == 0:
+        process_id = pid()
+        print(f"Root process {process_id} executed")
+
+    # Or trace specific user (e.g., UID 1000)
+    if current_uid == 1002:
+        print("User 1002 executed something")
+
+    return 0
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+compile()