mirror of
https://github.com/varun-r-mallya/Python-BPF.git
synced 2025-12-31 21:06:25 +00:00
Compare commits
56 Commits
32int_supp
...
symex
| Author | SHA1 | Date | |
|---|---|---|---|
| da45daa972 | |||
| 39a0746db4 | |||
| e9bb90cb70 | |||
| 9d76502d5a | |||
| a10da4a277 | |||
| 29e90601b7 | |||
| 56df05a93c | |||
| a55efc6469 | |||
| 64cd2d2fc2 | |||
| cbddc0aa96 | |||
| 209df33c8f | |||
| 7a56e5d0cd | |||
| 1d7a436c9f | |||
| 5eaeb3e921 | |||
| cd52d0d91b | |||
| df981be095 | |||
| 316c21c428 | |||
| c883d95655 | |||
| 5a8b64f1d9 | |||
| cf99b3bb9a | |||
| 6c85b248ce | |||
| b5a3494cc6 | |||
| be62972974 | |||
| 2f4a7d2f90 | |||
| 3ccd3f767e | |||
| 2e37726922 | |||
| 5b36726b7d | |||
| faad3555dc | |||
| 3e6cea2b67 | |||
| 338d4994d8 | |||
| 3078d4224d | |||
| 7d29790f00 | |||
| 963e2a8171 | |||
| 123a92af1d | |||
| 752f564d3f | |||
| d8cddb9799 | |||
| 33e18f6d6d | |||
| 5e371787eb | |||
| 67c9d9b932 | |||
| f757a32a63 | |||
| c5de92b9d0 | |||
| 4efd3223cd | |||
| 4884ed7577 | |||
| 5b7769dd38 | |||
| b7c1e92f05 | |||
| 8b28a927c3 | |||
| f9ee43e7ef | |||
| dabb8bf0df | |||
| 19dedede53 | |||
| 82cac8f8ef | |||
| 70a04f54d1 | |||
| ec2ea835e5 | |||
| 2257c175ed | |||
| 5bf60d69b8 | |||
| 0006e26b08 | |||
| 5cbd9a531e |
@ -68,8 +68,6 @@ def callback(cpu, event):
|
||||
|
||||
perf = b["events"].open_perf_buffer(callback, struct_name="data_t")
|
||||
print("Starting to poll... (Ctrl+C to stop)")
|
||||
print("Try running: fork() or clone() system calls to trigger events")
|
||||
|
||||
try:
|
||||
while True:
|
||||
b["events"].poll(1000)
|
||||
|
||||
@ -26,7 +26,7 @@ classifiers = [
|
||||
]
|
||||
readme = "README.md"
|
||||
license = {text = "Apache-2.0"}
|
||||
requires-python = ">=3.8"
|
||||
requires-python = ">=3.10"
|
||||
|
||||
dependencies = [
|
||||
"llvmlite",
|
||||
|
||||
@ -177,17 +177,33 @@ def _allocate_for_binop(builder, var_name, local_sym_tab):
|
||||
logger.info(f"Pre-allocated {var_name} for binop result")
|
||||
|
||||
|
||||
def _get_type_name(ir_type):
|
||||
"""Get a string representation of an IR type."""
|
||||
if isinstance(ir_type, ir.IntType):
|
||||
return f"i{ir_type.width}"
|
||||
elif isinstance(ir_type, ir.PointerType):
|
||||
return "ptr"
|
||||
elif isinstance(ir_type, ir.ArrayType):
|
||||
return f"[{ir_type.count}x{_get_type_name(ir_type.element)}]"
|
||||
else:
|
||||
return str(ir_type).replace(" ", "")
|
||||
|
||||
|
||||
def allocate_temp_pool(builder, max_temps, local_sym_tab):
|
||||
"""Allocate the temporary scratch space pool for helper arguments."""
|
||||
if max_temps == 0:
|
||||
if not max_temps:
|
||||
logger.info("No temp pool allocation needed")
|
||||
return
|
||||
|
||||
logger.info(f"Allocating temp pool of {max_temps} variables")
|
||||
for i in range(max_temps):
|
||||
temp_name = f"__helper_temp_{i}"
|
||||
temp_var = builder.alloca(ir.IntType(64), name=temp_name)
|
||||
temp_var.align = 8
|
||||
local_sym_tab[temp_name] = LocalSymbol(temp_var, ir.IntType(64))
|
||||
for tmp_type, cnt in max_temps.items():
|
||||
type_name = _get_type_name(tmp_type)
|
||||
logger.info(f"Allocating temp pool of {cnt} variables of type {type_name}")
|
||||
for i in range(cnt):
|
||||
temp_name = f"__helper_temp_{type_name}_{i}"
|
||||
temp_var = builder.alloca(tmp_type, name=temp_name)
|
||||
temp_var.align = _get_alignment(tmp_type)
|
||||
local_sym_tab[temp_name] = LocalSymbol(temp_var, tmp_type)
|
||||
logger.debug(f"Allocated temp variable: {temp_name}")
|
||||
|
||||
|
||||
def _allocate_for_name(builder, var_name, rval, local_sym_tab):
|
||||
|
||||
@ -86,7 +86,7 @@ def processor(source_code, filename, module):
|
||||
license_processing(tree, module)
|
||||
globals_processing(tree, module)
|
||||
structs_sym_tab = structs_proc(tree, module, bpf_chunks)
|
||||
map_sym_tab = maps_proc(tree, module, bpf_chunks)
|
||||
map_sym_tab = maps_proc(tree, module, bpf_chunks, structs_sym_tab)
|
||||
func_proc(tree, module, bpf_chunks, map_sym_tab, structs_sym_tab)
|
||||
|
||||
globals_list_creation(tree, module)
|
||||
@ -218,13 +218,11 @@ def compile(loglevel=logging.WARNING) -> bool:
|
||||
def BPF(loglevel=logging.WARNING) -> BpfObject:
|
||||
caller_frame = inspect.stack()[1]
|
||||
src = inspect.getsource(caller_frame.frame)
|
||||
with tempfile.NamedTemporaryFile(
|
||||
mode="w+", delete=True, suffix=".py"
|
||||
) as f, tempfile.NamedTemporaryFile(
|
||||
mode="w+", delete=True, suffix=".ll"
|
||||
) as inter, tempfile.NamedTemporaryFile(
|
||||
mode="w+", delete=False, suffix=".o"
|
||||
) as obj_file:
|
||||
with (
|
||||
tempfile.NamedTemporaryFile(mode="w+", delete=True, suffix=".py") as f,
|
||||
tempfile.NamedTemporaryFile(mode="w+", delete=True, suffix=".ll") as inter,
|
||||
tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".o") as obj_file,
|
||||
):
|
||||
f.write(src)
|
||||
f.flush()
|
||||
source = f.name
|
||||
|
||||
@ -39,7 +39,7 @@ logger = logging.getLogger(__name__)
|
||||
def count_temps_in_call(call_node, local_sym_tab):
|
||||
"""Count the number of temporary variables needed for a function call."""
|
||||
|
||||
count = 0
|
||||
count = {}
|
||||
is_helper = False
|
||||
|
||||
# NOTE: We exclude print calls for now
|
||||
@ -49,21 +49,28 @@ def count_temps_in_call(call_node, local_sym_tab):
|
||||
and call_node.func.id != "print"
|
||||
):
|
||||
is_helper = True
|
||||
func_name = call_node.func.id
|
||||
elif isinstance(call_node.func, ast.Attribute):
|
||||
if HelperHandlerRegistry.has_handler(call_node.func.attr):
|
||||
is_helper = True
|
||||
func_name = call_node.func.attr
|
||||
|
||||
if not is_helper:
|
||||
return 0
|
||||
return {} # No temps needed
|
||||
|
||||
for arg in call_node.args:
|
||||
for arg_idx in range(len(call_node.args)):
|
||||
# NOTE: Count all non-name arguments
|
||||
# For struct fields, if it is being passed as an argument,
|
||||
# The struct object should already exist in the local_sym_tab
|
||||
if not isinstance(arg, ast.Name) and not (
|
||||
arg = call_node.args[arg_idx]
|
||||
if isinstance(arg, ast.Name) or (
|
||||
isinstance(arg, ast.Attribute) and arg.value.id in local_sym_tab
|
||||
):
|
||||
count += 1
|
||||
continue
|
||||
param_type = HelperHandlerRegistry.get_param_type(func_name, arg_idx)
|
||||
if isinstance(param_type, ir.PointerType):
|
||||
pointee_type = param_type.pointee
|
||||
count[pointee_type] = count.get(pointee_type, 0) + 1
|
||||
|
||||
return count
|
||||
|
||||
@ -99,11 +106,15 @@ def handle_if_allocation(
|
||||
def allocate_mem(
|
||||
module, builder, body, func, ret_type, map_sym_tab, local_sym_tab, structs_sym_tab
|
||||
):
|
||||
max_temps_needed = 0
|
||||
max_temps_needed = {}
|
||||
|
||||
def merge_type_counts(count_dict):
|
||||
nonlocal max_temps_needed
|
||||
for typ, cnt in count_dict.items():
|
||||
max_temps_needed[typ] = max(max_temps_needed.get(typ, 0), cnt)
|
||||
|
||||
def update_max_temps_for_stmt(stmt):
|
||||
nonlocal max_temps_needed
|
||||
temps_needed = 0
|
||||
|
||||
if isinstance(stmt, ast.If):
|
||||
for s in stmt.body:
|
||||
@ -112,10 +123,13 @@ def allocate_mem(
|
||||
update_max_temps_for_stmt(s)
|
||||
return
|
||||
|
||||
stmt_temps = {}
|
||||
for node in ast.walk(stmt):
|
||||
if isinstance(node, ast.Call):
|
||||
temps_needed += count_temps_in_call(node, local_sym_tab)
|
||||
max_temps_needed = max(max_temps_needed, temps_needed)
|
||||
call_temps = count_temps_in_call(node, local_sym_tab)
|
||||
for typ, cnt in call_temps.items():
|
||||
stmt_temps[typ] = stmt_temps.get(typ, 0) + cnt
|
||||
merge_type_counts(stmt_temps)
|
||||
|
||||
for stmt in body:
|
||||
update_max_temps_for_stmt(stmt)
|
||||
|
||||
@ -1,7 +1,21 @@
|
||||
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,
|
||||
get_stack,
|
||||
XDP_DROP,
|
||||
XDP_PASS,
|
||||
)
|
||||
|
||||
|
||||
# Register the helper handler with expr module
|
||||
@ -65,6 +79,12 @@ __all__ = [
|
||||
"deref",
|
||||
"comm",
|
||||
"probe_read_str",
|
||||
"random",
|
||||
"probe_read",
|
||||
"smp_processor_id",
|
||||
"uid",
|
||||
"skb_store_bytes",
|
||||
"get_stack",
|
||||
"XDP_DROP",
|
||||
"XDP_PASS",
|
||||
]
|
||||
|
||||
@ -8,30 +8,43 @@ 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,
|
||||
)
|
||||
from .printk_formatter import simple_string_print, handle_fstring_print
|
||||
|
||||
from logging import Logger
|
||||
from pythonbpf.maps import BPFMapType
|
||||
import logging
|
||||
|
||||
logger: Logger = logging.getLogger(__name__)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
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_GET_STACK = 67
|
||||
BPF_PROBE_READ_KERNEL_STR = 115
|
||||
BPF_RINGBUF_OUTPUT = 130
|
||||
BPF_RINGBUF_RESERVE = 131
|
||||
BPF_RINGBUF_SUBMIT = 132
|
||||
BPF_RINGBUF_DISCARD = 133
|
||||
|
||||
|
||||
@HelperHandlerRegistry.register("ktime")
|
||||
@HelperHandlerRegistry.register(
|
||||
"ktime",
|
||||
param_types=[],
|
||||
return_type=ir.IntType(64),
|
||||
)
|
||||
def bpf_ktime_get_ns_emitter(
|
||||
call,
|
||||
map_ptr,
|
||||
@ -54,7 +67,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,
|
||||
@ -96,6 +113,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,
|
||||
@ -144,7 +162,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,
|
||||
@ -199,7 +225,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,
|
||||
@ -239,7 +269,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,
|
||||
@ -296,7 +330,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,
|
||||
@ -318,12 +356,12 @@ 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)
|
||||
|
||||
|
||||
@HelperHandlerRegistry.register("output")
|
||||
def bpf_perf_event_output_handler(
|
||||
call,
|
||||
map_ptr,
|
||||
@ -334,6 +372,10 @@ def bpf_perf_event_output_handler(
|
||||
struct_sym_tab=None,
|
||||
map_sym_tab=None,
|
||||
):
|
||||
"""
|
||||
Emit LLVM IR for bpf_perf_event_output helper function call.
|
||||
"""
|
||||
|
||||
if len(call.args) != 1:
|
||||
raise ValueError(
|
||||
f"Perf event output expects exactly one argument, got {len(call.args)}"
|
||||
@ -371,6 +413,98 @@ def bpf_perf_event_output_handler(
|
||||
return result, None
|
||||
|
||||
|
||||
def bpf_ringbuf_output_emitter(
|
||||
call,
|
||||
map_ptr,
|
||||
module,
|
||||
builder,
|
||||
func,
|
||||
local_sym_tab=None,
|
||||
struct_sym_tab=None,
|
||||
map_sym_tab=None,
|
||||
):
|
||||
"""
|
||||
Emit LLVM IR for bpf_ringbuf_output helper function call.
|
||||
"""
|
||||
|
||||
if len(call.args) != 1:
|
||||
raise ValueError(
|
||||
f"Ringbuf output expects exactly one argument, got {len(call.args)}"
|
||||
)
|
||||
data_arg = call.args[0]
|
||||
data_ptr, size_val = get_data_ptr_and_size(data_arg, local_sym_tab, struct_sym_tab)
|
||||
flags_val = ir.Constant(ir.IntType(64), 0)
|
||||
|
||||
map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
|
||||
data_void_ptr = builder.bitcast(data_ptr, ir.PointerType())
|
||||
fn_type = ir.FunctionType(
|
||||
ir.IntType(64),
|
||||
[
|
||||
ir.PointerType(),
|
||||
ir.PointerType(),
|
||||
ir.IntType(64),
|
||||
ir.IntType(64),
|
||||
],
|
||||
var_arg=False,
|
||||
)
|
||||
fn_ptr_type = ir.PointerType(fn_type)
|
||||
|
||||
# helper id
|
||||
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_RINGBUF_OUTPUT.value)
|
||||
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
|
||||
|
||||
result = builder.call(
|
||||
fn_ptr, [map_void_ptr, data_void_ptr, size_val, flags_val], tail=False
|
||||
)
|
||||
return result, None
|
||||
|
||||
|
||||
@HelperHandlerRegistry.register(
|
||||
"output",
|
||||
param_types=[ir.PointerType(ir.IntType(8))],
|
||||
return_type=ir.IntType(64),
|
||||
)
|
||||
def handle_output_helper(
|
||||
call,
|
||||
map_ptr,
|
||||
module,
|
||||
builder,
|
||||
func,
|
||||
local_sym_tab=None,
|
||||
struct_sym_tab=None,
|
||||
map_sym_tab=None,
|
||||
):
|
||||
"""
|
||||
Route output helper to the appropriate emitter based on map type.
|
||||
"""
|
||||
match map_sym_tab[map_ptr.name].type:
|
||||
case BPFMapType.PERF_EVENT_ARRAY:
|
||||
return bpf_perf_event_output_handler(
|
||||
call,
|
||||
map_ptr,
|
||||
module,
|
||||
builder,
|
||||
func,
|
||||
local_sym_tab,
|
||||
struct_sym_tab,
|
||||
map_sym_tab,
|
||||
)
|
||||
case BPFMapType.RINGBUF:
|
||||
return bpf_ringbuf_output_emitter(
|
||||
call,
|
||||
map_ptr,
|
||||
module,
|
||||
builder,
|
||||
func,
|
||||
local_sym_tab,
|
||||
struct_sym_tab,
|
||||
map_sym_tab,
|
||||
)
|
||||
case _:
|
||||
logger.error("Unsupported map type for output helper.")
|
||||
raise NotImplementedError("Output helper for this map type is not implemented.")
|
||||
|
||||
|
||||
def emit_probe_read_kernel_str_call(builder, dst_ptr, dst_size, src_ptr):
|
||||
"""Emit LLVM IR call to bpf_probe_read_kernel_str"""
|
||||
|
||||
@ -398,7 +532,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,
|
||||
@ -417,8 +558,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)
|
||||
@ -433,6 +574,430 @@ def bpf_probe_read_kernel_str_emitter(
|
||||
return result, ir.IntType(64)
|
||||
|
||||
|
||||
@HelperHandlerRegistry.register(
|
||||
"random",
|
||||
param_types=[],
|
||||
return_type=ir.IntType(32),
|
||||
)
|
||||
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",
|
||||
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,
|
||||
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,
|
||||
ir.IntType(8),
|
||||
)
|
||||
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,
|
||||
ir.IntType(8),
|
||||
)
|
||||
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",
|
||||
param_types=[],
|
||||
return_type=ir.IntType(32),
|
||||
)
|
||||
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",
|
||||
param_types=[],
|
||||
return_type=ir.IntType(64),
|
||||
)
|
||||
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",
|
||||
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(
|
||||
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)
|
||||
"""
|
||||
|
||||
args_signature = [
|
||||
ir.PointerType(), # skb pointer
|
||||
ir.IntType(32), # offset
|
||||
ir.PointerType(), # from
|
||||
ir.IntType(32), # len
|
||||
ir.IntType(64), # 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,
|
||||
args_signature[2],
|
||||
)
|
||||
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
|
||||
if isinstance(flags_val, int):
|
||||
flags = ir.Constant(ir.IntType(64), flags_val)
|
||||
else:
|
||||
flags = flags_val
|
||||
fn_type = ir.FunctionType(
|
||||
ir.IntType(64),
|
||||
args_signature,
|
||||
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)
|
||||
|
||||
|
||||
@HelperHandlerRegistry.register(
|
||||
"reserve",
|
||||
param_types=[ir.IntType(64)],
|
||||
return_type=ir.PointerType(ir.IntType(8)),
|
||||
)
|
||||
def bpf_ringbuf_reserve_emitter(
|
||||
call,
|
||||
map_ptr,
|
||||
module,
|
||||
builder,
|
||||
func,
|
||||
local_sym_tab=None,
|
||||
struct_sym_tab=None,
|
||||
map_sym_tab=None,
|
||||
):
|
||||
"""
|
||||
Emit LLVM IR for bpf_ringbuf_reserve helper function call.
|
||||
Expected call signature: ringbuf.reserve(size)
|
||||
"""
|
||||
|
||||
if len(call.args) != 1:
|
||||
raise ValueError(
|
||||
f"ringbuf.reserve expects exactly one argument (size), got {len(call.args)}"
|
||||
)
|
||||
|
||||
size_val = get_int_value_from_arg(
|
||||
call.args[0],
|
||||
func,
|
||||
module,
|
||||
builder,
|
||||
local_sym_tab,
|
||||
map_sym_tab,
|
||||
struct_sym_tab,
|
||||
)
|
||||
|
||||
map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
|
||||
fn_type = ir.FunctionType(
|
||||
ir.PointerType(ir.IntType(8)),
|
||||
[ir.PointerType(), ir.IntType(64)],
|
||||
var_arg=False,
|
||||
)
|
||||
fn_ptr_type = ir.PointerType(fn_type)
|
||||
|
||||
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_RINGBUF_RESERVE.value)
|
||||
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
|
||||
|
||||
result = builder.call(fn_ptr, [map_void_ptr, size_val], tail=False)
|
||||
|
||||
return result, ir.PointerType(ir.IntType(8))
|
||||
|
||||
|
||||
@HelperHandlerRegistry.register(
|
||||
"submit",
|
||||
param_types=[ir.PointerType(ir.IntType(8)), ir.IntType(64)],
|
||||
return_type=ir.VoidType(),
|
||||
)
|
||||
def bpf_ringbuf_submit_emitter(
|
||||
call,
|
||||
map_ptr,
|
||||
module,
|
||||
builder,
|
||||
func,
|
||||
local_sym_tab=None,
|
||||
struct_sym_tab=None,
|
||||
map_sym_tab=None,
|
||||
):
|
||||
"""
|
||||
Emit LLVM IR for bpf_ringbuf_submit helper function call.
|
||||
Expected call signature: ringbuf.submit(data, flags=0)
|
||||
"""
|
||||
|
||||
if len(call.args) not in (1, 2):
|
||||
raise ValueError(
|
||||
f"ringbuf.submit expects 1 or 2 args (data, flags), got {len(call.args)}"
|
||||
)
|
||||
|
||||
data_arg = call.args[0]
|
||||
flags_arg = call.args[1] if len(call.args) == 2 else None
|
||||
|
||||
data_ptr = get_or_create_ptr_from_arg(
|
||||
func,
|
||||
module,
|
||||
data_arg,
|
||||
builder,
|
||||
local_sym_tab,
|
||||
map_sym_tab,
|
||||
struct_sym_tab,
|
||||
ir.PointerType(ir.IntType(8)),
|
||||
)
|
||||
|
||||
flags_const = get_flags_val(flags_arg, builder, local_sym_tab)
|
||||
if isinstance(flags_const, int):
|
||||
flags_const = ir.Constant(ir.IntType(64), flags_const)
|
||||
|
||||
map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
|
||||
fn_type = ir.FunctionType(
|
||||
ir.VoidType(),
|
||||
[ir.PointerType(), ir.PointerType(), ir.IntType(64)],
|
||||
var_arg=False,
|
||||
)
|
||||
fn_ptr_type = ir.PointerType(fn_type)
|
||||
|
||||
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_RINGBUF_SUBMIT.value)
|
||||
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
|
||||
|
||||
result = builder.call(fn_ptr, [map_void_ptr, data_ptr, flags_const], tail=False)
|
||||
|
||||
return result, None
|
||||
|
||||
|
||||
@HelperHandlerRegistry.register(
|
||||
"get_stack",
|
||||
param_types=[ir.PointerType(ir.IntType(8)), ir.IntType(64)],
|
||||
return_type=ir.IntType(64),
|
||||
)
|
||||
def bpf_get_stack_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_stack helper function call.
|
||||
"""
|
||||
if len(call.args) not in (1, 2):
|
||||
raise ValueError(
|
||||
f"get_stack expects atmost two arguments (buf, flags), got {len(call.args)}"
|
||||
)
|
||||
ctx_ptr = func.args[0] # First argument to the function is ctx
|
||||
buf_arg = call.args[0]
|
||||
flags_arg = call.args[1] if len(call.args) == 2 else None
|
||||
buf_ptr, buf_size = get_buffer_ptr_and_size(
|
||||
buf_arg, builder, local_sym_tab, struct_sym_tab
|
||||
)
|
||||
flags_val = get_flags_val(flags_arg, builder, local_sym_tab)
|
||||
if isinstance(flags_val, int):
|
||||
flags_val = ir.Constant(ir.IntType(64), flags_val)
|
||||
|
||||
buf_void_ptr = builder.bitcast(buf_ptr, ir.PointerType())
|
||||
fn_type = ir.FunctionType(
|
||||
ir.IntType(64),
|
||||
[
|
||||
ir.PointerType(ir.IntType(8)),
|
||||
ir.PointerType(),
|
||||
ir.IntType(64),
|
||||
ir.IntType(64),
|
||||
],
|
||||
var_arg=False,
|
||||
)
|
||||
fn_ptr_type = ir.PointerType(fn_type)
|
||||
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_GET_STACK.value)
|
||||
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
|
||||
result = builder.call(
|
||||
fn_ptr,
|
||||
[ctx_ptr, buf_void_ptr, ir.Constant(ir.IntType(64), buf_size), flags_val],
|
||||
tail=False,
|
||||
)
|
||||
return result, ir.IntType(64)
|
||||
|
||||
|
||||
def handle_helper_call(
|
||||
call,
|
||||
module,
|
||||
@ -487,6 +1052,6 @@ def handle_helper_call(
|
||||
if not map_sym_tab or map_name not in map_sym_tab:
|
||||
raise ValueError(f"Map '{map_name}' not found in symbol table")
|
||||
|
||||
return invoke_helper(method_name, map_sym_tab[map_name])
|
||||
return invoke_helper(method_name, map_sym_tab[map_name].sym)
|
||||
|
||||
return None
|
||||
|
||||
@ -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,9 +33,29 @@ 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):
|
||||
"""Check if a handler function is registered for a helper"""
|
||||
return helper_name in cls._handlers
|
||||
|
||||
@classmethod
|
||||
def get_signature(cls, helper_name):
|
||||
"""Get the signature of a helper function"""
|
||||
return cls._handlers.get(helper_name)
|
||||
|
||||
@classmethod
|
||||
def get_param_type(cls, helper_name, index):
|
||||
"""Get the type of a parameter of a helper function by the index"""
|
||||
signature = cls.get_signature(helper_name)
|
||||
if signature and signature.arg_types and 0 <= index < len(signature.arg_types):
|
||||
return signature.arg_types[index]
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def get_return_type(cls, helper_name):
|
||||
"""Get the return type of a helper function"""
|
||||
signature = cls.get_signature(helper_name)
|
||||
return signature.return_type if signature else None
|
||||
|
||||
@ -14,26 +14,43 @@ class ScratchPoolManager:
|
||||
"""Manage the temporary helper variables in local_sym_tab"""
|
||||
|
||||
def __init__(self):
|
||||
self._counter = 0
|
||||
self._counters = {}
|
||||
|
||||
@property
|
||||
def counter(self):
|
||||
return self._counter
|
||||
return sum(self._counters.values())
|
||||
|
||||
def reset(self):
|
||||
self._counter = 0
|
||||
self._counters.clear()
|
||||
logger.debug("Scratch pool counter reset to 0")
|
||||
|
||||
def get_next_temp(self, local_sym_tab):
|
||||
temp_name = f"__helper_temp_{self._counter}"
|
||||
self._counter += 1
|
||||
def _get_type_name(self, ir_type):
|
||||
if isinstance(ir_type, ir.PointerType):
|
||||
return "ptr"
|
||||
elif isinstance(ir_type, ir.IntType):
|
||||
return f"i{ir_type.width}"
|
||||
elif isinstance(ir_type, ir.ArrayType):
|
||||
return f"[{ir_type.count}x{self._get_type_name(ir_type.element)}]"
|
||||
else:
|
||||
return str(ir_type).replace(" ", "")
|
||||
|
||||
def get_next_temp(self, local_sym_tab, expected_type=None):
|
||||
# Default to i64 if no expected type provided
|
||||
type_name = self._get_type_name(expected_type) if expected_type else "i64"
|
||||
if type_name not in self._counters:
|
||||
self._counters[type_name] = 0
|
||||
|
||||
counter = self._counters[type_name]
|
||||
temp_name = f"__helper_temp_{type_name}_{counter}"
|
||||
self._counters[type_name] += 1
|
||||
|
||||
if temp_name not in local_sym_tab:
|
||||
raise ValueError(
|
||||
f"Scratch pool exhausted or inadequate: {temp_name}. "
|
||||
f"Current counter: {self._counter}"
|
||||
f"Type: {type_name} Counter: {counter}"
|
||||
)
|
||||
|
||||
logger.debug(f"Using {temp_name} for type {type_name}")
|
||||
return local_sym_tab[temp_name].var, temp_name
|
||||
|
||||
|
||||
@ -60,24 +77,73 @@ def get_var_ptr_from_name(var_name, local_sym_tab):
|
||||
def create_int_constant_ptr(value, builder, local_sym_tab, int_width=64):
|
||||
"""Create a pointer to an integer constant."""
|
||||
|
||||
# Default to 64-bit integer
|
||||
ptr, temp_name = _temp_pool_manager.get_next_temp(local_sym_tab)
|
||||
int_type = ir.IntType(int_width)
|
||||
ptr, temp_name = _temp_pool_manager.get_next_temp(local_sym_tab, int_type)
|
||||
logger.info(f"Using temp variable '{temp_name}' for int constant {value}")
|
||||
const_val = ir.Constant(ir.IntType(int_width), value)
|
||||
const_val = ir.Constant(int_type, value)
|
||||
builder.store(const_val, ptr)
|
||||
return ptr
|
||||
|
||||
|
||||
def get_or_create_ptr_from_arg(
|
||||
func, module, arg, builder, local_sym_tab, map_sym_tab, struct_sym_tab=None
|
||||
func,
|
||||
module,
|
||||
arg,
|
||||
builder,
|
||||
local_sym_tab,
|
||||
map_sym_tab,
|
||||
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):
|
||||
ptr = create_int_constant_ptr(arg.value, builder, local_sym_tab)
|
||||
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
|
||||
val = get_operand_value(
|
||||
func, module, arg, builder, local_sym_tab, map_sym_tab, struct_sym_tab
|
||||
@ -85,13 +151,20 @@ def get_or_create_ptr_from_arg(
|
||||
if val is None:
|
||||
raise ValueError("Failed to evaluate expression for helper arg.")
|
||||
|
||||
# NOTE: We assume the result is an int64 for now
|
||||
# if isinstance(arg, ast.Attribute):
|
||||
# return val
|
||||
ptr, temp_name = _temp_pool_manager.get_next_temp(local_sym_tab)
|
||||
ptr, temp_name = _temp_pool_manager.get_next_temp(local_sym_tab, expected_type)
|
||||
logger.info(f"Using temp variable '{temp_name}' for expression result")
|
||||
if (
|
||||
isinstance(val.type, ir.IntType)
|
||||
and expected_type
|
||||
and val.type.width > expected_type.width
|
||||
):
|
||||
val = builder.trunc(val, expected_type)
|
||||
builder.store(val, ptr)
|
||||
|
||||
# NOTE: For char arrays, also return size
|
||||
if sz:
|
||||
return ptr, sz
|
||||
|
||||
return ptr
|
||||
|
||||
|
||||
@ -214,7 +287,10 @@ def get_char_array_ptr_and_size(buf_arg, builder, local_sym_tab, struct_sym_tab)
|
||||
|
||||
field_type = struct_info.field_type(field_name)
|
||||
if not _is_char_array(field_type):
|
||||
raise ValueError("Expected char array field")
|
||||
logger.info(
|
||||
"Field is not a char array, falling back to int or ptr detection"
|
||||
)
|
||||
return None, 0
|
||||
|
||||
struct_ptr = local_sym_tab[var_name].var
|
||||
field_ptr = struct_info.gep(builder, struct_ptr, field_name)
|
||||
@ -274,3 +350,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
|
||||
|
||||
@ -27,6 +27,36 @@ 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)
|
||||
|
||||
|
||||
def get_stack(buf, flags=0):
|
||||
"""get the current stack trace"""
|
||||
return ctypes.c_int64(0)
|
||||
|
||||
|
||||
XDP_ABORTED = ctypes.c_int64(0)
|
||||
XDP_DROP = ctypes.c_int64(1)
|
||||
XDP_PASS = ctypes.c_int64(2)
|
||||
|
||||
@ -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,11 +220,12 @@ 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(
|
||||
expr, builder, local_sym_tab, struct_sym_tab
|
||||
)
|
||||
if char_array_ptr:
|
||||
return char_array_ptr
|
||||
if isinstance(expr, ast.Attribute):
|
||||
char_array_ptr, _ = get_char_array_ptr_and_size(
|
||||
expr, builder, local_sym_tab, struct_sym_tab
|
||||
)
|
||||
if char_array_ptr:
|
||||
return char_array_ptr
|
||||
|
||||
# Regular expression evaluation
|
||||
val, _ = eval_expr(func, module, builder, expr, local_sym_tab, None, struct_sym_tab)
|
||||
@ -242,52 +244,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)
|
||||
|
||||
@ -1,4 +1,5 @@
|
||||
from .maps import HashMap, PerfEventArray, RingBuf
|
||||
from .maps import HashMap, PerfEventArray, RingBuffer
|
||||
from .maps_pass import maps_proc
|
||||
from .map_types import BPFMapType
|
||||
|
||||
__all__ = ["HashMap", "PerfEventArray", "maps_proc", "RingBuf"]
|
||||
__all__ = ["HashMap", "PerfEventArray", "maps_proc", "RingBuffer", "BPFMapType"]
|
||||
|
||||
@ -2,7 +2,7 @@ from pythonbpf.debuginfo import DebugInfoGenerator
|
||||
from .map_types import BPFMapType
|
||||
|
||||
|
||||
def create_map_debug_info(module, map_global, map_name, map_params):
|
||||
def create_map_debug_info(module, map_global, map_name, map_params, structs_sym_tab):
|
||||
"""Generate debug info metadata for BPF maps HASH and PERF_EVENT_ARRAY"""
|
||||
generator = DebugInfoGenerator(module)
|
||||
|
||||
@ -64,7 +64,13 @@ def create_map_debug_info(module, map_global, map_name, map_params):
|
||||
return global_var
|
||||
|
||||
|
||||
def create_ringbuf_debug_info(module, map_global, map_name, map_params):
|
||||
# TODO: This should not be exposed outside of the module.
|
||||
# Ideally we should expose a single create_map_debug_info function that handles all map types.
|
||||
# We can probably use a registry pattern to register different map types and their debug info generators.
|
||||
# map_params["type"] will be used to determine which generator to use.
|
||||
def create_ringbuf_debug_info(
|
||||
module, map_global, map_name, map_params, structs_sym_tab
|
||||
):
|
||||
"""Generate debug information metadata for BPF RINGBUF map"""
|
||||
generator = DebugInfoGenerator(module)
|
||||
|
||||
|
||||
@ -36,11 +36,14 @@ class PerfEventArray:
|
||||
pass # Placeholder for output method
|
||||
|
||||
|
||||
class RingBuf:
|
||||
class RingBuffer:
|
||||
def __init__(self, max_entries):
|
||||
self.max_entries = max_entries
|
||||
|
||||
def reserve(self, size: int, flags=0):
|
||||
def output(self, data, flags=0):
|
||||
pass
|
||||
|
||||
def reserve(self, size: int):
|
||||
if size > self.max_entries:
|
||||
raise ValueError("size cannot be greater than set maximum entries")
|
||||
return 0
|
||||
@ -48,4 +51,7 @@ class RingBuf:
|
||||
def submit(self, data, flags=0):
|
||||
pass
|
||||
|
||||
def discard(self, data, flags=0):
|
||||
pass
|
||||
|
||||
# add discard, output and also give names to flags and stuff
|
||||
|
||||
@ -3,7 +3,7 @@ import logging
|
||||
from logging import Logger
|
||||
from llvmlite import ir
|
||||
|
||||
from .maps_utils import MapProcessorRegistry
|
||||
from .maps_utils import MapProcessorRegistry, MapSymbol
|
||||
from .map_types import BPFMapType
|
||||
from .map_debug_info import create_map_debug_info, create_ringbuf_debug_info
|
||||
from pythonbpf.expr.vmlinux_registry import VmlinuxHandlerRegistry
|
||||
@ -12,13 +12,15 @@ from pythonbpf.expr.vmlinux_registry import VmlinuxHandlerRegistry
|
||||
logger: Logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def maps_proc(tree, module, chunks):
|
||||
def maps_proc(tree, module, chunks, structs_sym_tab):
|
||||
"""Process all functions decorated with @map to find BPF maps"""
|
||||
map_sym_tab = {}
|
||||
for func_node in chunks:
|
||||
if is_map(func_node):
|
||||
logger.info(f"Found BPF map: {func_node.name}")
|
||||
map_sym_tab[func_node.name] = process_bpf_map(func_node, module)
|
||||
map_sym_tab[func_node.name] = process_bpf_map(
|
||||
func_node, module, structs_sym_tab
|
||||
)
|
||||
return map_sym_tab
|
||||
|
||||
|
||||
@ -46,7 +48,7 @@ def create_bpf_map(module, map_name, map_params):
|
||||
map_global.align = 8
|
||||
|
||||
logger.info(f"Created BPF map: {map_name} with params {map_params}")
|
||||
return map_global
|
||||
return MapSymbol(type=map_params["type"], sym=map_global)
|
||||
|
||||
|
||||
def _parse_map_params(rval, expected_args=None):
|
||||
@ -60,7 +62,8 @@ def _parse_map_params(rval, expected_args=None):
|
||||
if i < len(rval.args):
|
||||
arg = rval.args[i]
|
||||
if isinstance(arg, ast.Name):
|
||||
params[arg_name] = arg.id
|
||||
result = _get_vmlinux_enum(handler, arg.id)
|
||||
params[arg_name] = result if result is not None else arg.id
|
||||
elif isinstance(arg, ast.Constant):
|
||||
params[arg_name] = arg.value
|
||||
|
||||
@ -68,33 +71,46 @@ def _parse_map_params(rval, expected_args=None):
|
||||
for keyword in rval.keywords:
|
||||
if isinstance(keyword.value, ast.Name):
|
||||
name = keyword.value.id
|
||||
if handler and handler.is_vmlinux_enum(name):
|
||||
result = handler.get_vmlinux_enum_value(name)
|
||||
params[keyword.arg] = result if result is not None else name
|
||||
else:
|
||||
params[keyword.arg] = name
|
||||
result = _get_vmlinux_enum(handler, name)
|
||||
params[keyword.arg] = result if result is not None else name
|
||||
elif isinstance(keyword.value, ast.Constant):
|
||||
params[keyword.arg] = keyword.value.value
|
||||
|
||||
return params
|
||||
|
||||
|
||||
@MapProcessorRegistry.register("RingBuf")
|
||||
def process_ringbuf_map(map_name, rval, module):
|
||||
def _get_vmlinux_enum(handler, name):
|
||||
if handler and handler.is_vmlinux_enum(name):
|
||||
return handler.get_vmlinux_enum_value(name)
|
||||
|
||||
|
||||
@MapProcessorRegistry.register("RingBuffer")
|
||||
def process_ringbuf_map(map_name, rval, module, structs_sym_tab):
|
||||
"""Process a BPF_RINGBUF map declaration"""
|
||||
logger.info(f"Processing Ringbuf: {map_name}")
|
||||
map_params = _parse_map_params(rval, expected_args=["max_entries"])
|
||||
map_params["type"] = BPFMapType.RINGBUF
|
||||
|
||||
# NOTE: constraints borrowed from https://docs.ebpf.io/linux/map-type/BPF_MAP_TYPE_RINGBUF/
|
||||
max_entries = map_params.get("max_entries")
|
||||
if (
|
||||
not isinstance(max_entries, int)
|
||||
or max_entries < 4096
|
||||
or (max_entries & (max_entries - 1)) != 0
|
||||
):
|
||||
raise ValueError(
|
||||
"Ringbuf max_entries must be a power of two greater than or equal to the page size (4096)"
|
||||
)
|
||||
|
||||
logger.info(f"Ringbuf map parameters: {map_params}")
|
||||
|
||||
map_global = create_bpf_map(module, map_name, map_params)
|
||||
create_ringbuf_debug_info(module, map_global, map_name, map_params)
|
||||
create_ringbuf_debug_info(module, map_global.sym, map_name, map_params)
|
||||
return map_global
|
||||
|
||||
|
||||
@MapProcessorRegistry.register("HashMap")
|
||||
def process_hash_map(map_name, rval, module):
|
||||
def process_hash_map(map_name, rval, module, structs_sym_tab):
|
||||
"""Process a BPF_HASH map declaration"""
|
||||
logger.info(f"Processing HashMap: {map_name}")
|
||||
map_params = _parse_map_params(rval, expected_args=["key", "value", "max_entries"])
|
||||
@ -103,12 +119,12 @@ def process_hash_map(map_name, rval, module):
|
||||
logger.info(f"Map parameters: {map_params}")
|
||||
map_global = create_bpf_map(module, map_name, map_params)
|
||||
# Generate debug info for BTF
|
||||
create_map_debug_info(module, map_global, map_name, map_params)
|
||||
create_map_debug_info(module, map_global.sym, map_name, map_params)
|
||||
return map_global
|
||||
|
||||
|
||||
@MapProcessorRegistry.register("PerfEventArray")
|
||||
def process_perf_event_map(map_name, rval, module):
|
||||
def process_perf_event_map(map_name, rval, module, structs_sym_tab):
|
||||
"""Process a BPF_PERF_EVENT_ARRAY map declaration"""
|
||||
logger.info(f"Processing PerfEventArray: {map_name}")
|
||||
map_params = _parse_map_params(rval, expected_args=["key_size", "value_size"])
|
||||
@ -117,11 +133,11 @@ def process_perf_event_map(map_name, rval, module):
|
||||
logger.info(f"Map parameters: {map_params}")
|
||||
map_global = create_bpf_map(module, map_name, map_params)
|
||||
# Generate debug info for BTF
|
||||
create_map_debug_info(module, map_global, map_name, map_params)
|
||||
create_map_debug_info(module, map_global.sym, map_name, map_params)
|
||||
return map_global
|
||||
|
||||
|
||||
def process_bpf_map(func_node, module):
|
||||
def process_bpf_map(func_node, module, structs_sym_tab):
|
||||
"""Process a BPF map (a function decorated with @map)"""
|
||||
map_name = func_node.name
|
||||
logger.info(f"Processing BPF map: {map_name}")
|
||||
@ -140,7 +156,7 @@ def process_bpf_map(func_node, module):
|
||||
if isinstance(rval, ast.Call) and isinstance(rval.func, ast.Name):
|
||||
handler = MapProcessorRegistry.get_processor(rval.func.id)
|
||||
if handler:
|
||||
return handler(map_name, rval, module)
|
||||
return handler(map_name, rval, module, structs_sym_tab)
|
||||
else:
|
||||
logger.warning(f"Unknown map type {rval.func.id}, defaulting to HashMap")
|
||||
return process_hash_map(map_name, rval, module)
|
||||
|
||||
@ -1,5 +1,16 @@
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from llvmlite import ir
|
||||
from typing import Any
|
||||
from .map_types import BPFMapType
|
||||
|
||||
|
||||
@dataclass
|
||||
class MapSymbol:
|
||||
"""Class representing a symbol on the map"""
|
||||
|
||||
type: BPFMapType
|
||||
sym: ir.GlobalVariable
|
||||
|
||||
|
||||
class MapProcessorRegistry:
|
||||
|
||||
29
tests/passing_tests/helpers/bpf_probe_read.py
Normal file
29
tests/passing_tests/helpers/bpf_probe_read.py
Normal 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()
|
||||
25
tests/passing_tests/helpers/prandom.py
Normal file
25
tests/passing_tests/helpers/prandom.py
Normal file
@ -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()
|
||||
40
tests/passing_tests/helpers/smp_processor_id.py
Normal file
40
tests/passing_tests/helpers/smp_processor_id.py
Normal 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()
|
||||
31
tests/passing_tests/helpers/uid_gid.py
Normal file
31
tests/passing_tests/helpers/uid_gid.py
Normal file
@ -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()
|
||||
Reference in New Issue
Block a user