6 Commits

5 changed files with 130 additions and 17 deletions

View File

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

View File

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

View File

@ -32,6 +32,16 @@ def random():
return ctypes.c_int32(0) 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_ABORTED = ctypes.c_int64(0)
XDP_DROP = ctypes.c_int64(1) XDP_DROP = ctypes.c_int64(1)
XDP_PASS = ctypes.c_int64(2) XDP_PASS = ctypes.c_int64(2)

View File

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

View File

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