Fix imports for bpf_skb_store_bytes

pythonbpf/helper/__init__.py

@@ -10,6 +10,8 @@ from .helpers import (
     random,
     probe_read,
     smp_processor_id,
+    uid,
+    skb_store_bytes,
     XDP_DROP,
     XDP_PASS,
 )
@@ -79,6 +81,8 @@ __all__ = [
     "random",
     "probe_read",
     "smp_processor_id",
+    "uid",
+    "skb_store_bytes",
     "XDP_DROP",
     "XDP_PASS",
 ]

pythonbpf/helper/bpf_helper_handler.py

@@ -29,7 +29,9 @@ class BPFHelperID(Enum):
     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
@@ -322,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)
@@ -537,6 +540,112 @@ def bpf_get_smp_processor_id_emitter(
     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 (4, 5):
+        raise ValueError(
+            f"skb_store_bytes expects 4 or 5 args (skb, offset, from, len, flags), got {len(call.args)}"
+        )
+
+    skb_ptr = get_or_create_ptr_from_arg(
+        func, module, call.args[0], builder, local_sym_tab, map_sym_tab, struct_sym_tab
+    )
+    offset_val = get_int_value_from_arg(
+        call.args[1],
+        func,
+        module,
+        builder,
+        local_sym_tab,
+        map_sym_tab,
+        struct_sym_tab,
+    )
+    from_ptr = get_or_create_ptr_from_arg(
+        func, module, call.args[2], builder, local_sym_tab, map_sym_tab, struct_sym_tab
+    )
+    len_val = get_int_value_from_arg(
+        call.args[3],
+        func,
+        module,
+        builder,
+        local_sym_tab,
+        map_sym_tab,
+        struct_sym_tab,
+    )
+    if len(call.args) == 5:
+        flags_val = get_flags_val(call.args[4], builder, local_sym_tab)
+    else:
+        flags_val = ir.Constant(ir.IntType(64), 0)
+    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_val,
+        ],
+        tail=False,
+    )
+    logger.info("Emitted bpf_skb_store_bytes call")
+    return result, ir.IntType(64)
+
+
 def handle_helper_call(
     call,
     module,

pythonbpf/helper/helpers.py

@@ -42,6 +42,16 @@ def smp_processor_id():
     return ctypes.c_int32(0)
 
 
+def uid():
+    """get current user id"""
+    return ctypes.c_int32(0)
+
+
+def skb_store_bytes(skb, 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/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()