diff --git a/pyproject.toml b/pyproject.toml index 3d0f11df..9ab259d2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ classifiers = [ ] readme = "README.md" license = {text = "Apache-2.0"} -requires-python = ">=3.8" +requires-python = ">=3.10" dependencies = [ "llvmlite", diff --git a/pythonbpf/codegen.py b/pythonbpf/codegen.py index 60628e63..c9b75d3a 100644 --- a/pythonbpf/codegen.py +++ b/pythonbpf/codegen.py @@ -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 diff --git a/pythonbpf/helper/__init__.py b/pythonbpf/helper/__init__.py index 4c1d2831..6d38e791 100644 --- a/pythonbpf/helper/__init__.py +++ b/pythonbpf/helper/__init__.py @@ -12,6 +12,7 @@ smp_processor_id, uid, skb_store_bytes, + get_stack, XDP_DROP, XDP_PASS, ) @@ -83,6 +84,7 @@ def helper_call_handler( "smp_processor_id", "uid", "skb_store_bytes", + "get_stack", "XDP_DROP", "XDP_PASS", ] diff --git a/pythonbpf/helper/bpf_helper_handler.py b/pythonbpf/helper/bpf_helper_handler.py index 17ff7978..d54fda3c 100644 --- a/pythonbpf/helper/bpf_helper_handler.py +++ b/pythonbpf/helper/bpf_helper_handler.py @@ -12,11 +12,10 @@ 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): @@ -33,7 +32,12 @@ class BPFHelperID(Enum): 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( @@ -358,11 +362,6 @@ def bpf_get_current_pid_tgid_emitter( return pid, ir.IntType(64) -@HelperHandlerRegistry.register( - "output", - param_types=[ir.PointerType(ir.IntType(8))], - return_type=ir.IntType(64), -) def bpf_perf_event_output_handler( call, map_ptr, @@ -373,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)}" @@ -410,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""" @@ -711,7 +806,10 @@ def bpf_skb_store_bytes_emitter( flags_val = get_flags_val(call.args[3], builder, local_sym_tab) else: flags_val = 0 - flags = ir.Constant(ir.IntType(64), flags_val) + 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, @@ -736,6 +834,167 @@ def bpf_skb_store_bytes_emitter( 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, +): + 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, @@ -790,6 +1049,6 @@ def invoke_helper(method_name, map_ptr=None): 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 diff --git a/pythonbpf/helper/helpers.py b/pythonbpf/helper/helpers.py index 302b5263..c80d57d5 100644 --- a/pythonbpf/helper/helpers.py +++ b/pythonbpf/helper/helpers.py @@ -52,6 +52,11 @@ def skb_store_bytes(offset, from_buf, size, flags=0): 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) diff --git a/pythonbpf/maps/__init__.py b/pythonbpf/maps/__init__.py index 48fc9ff9..eb2007da 100644 --- a/pythonbpf/maps/__init__.py +++ b/pythonbpf/maps/__init__.py @@ -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"] diff --git a/pythonbpf/maps/maps.py b/pythonbpf/maps/maps.py index a2d7c219..583e9570 100644 --- a/pythonbpf/maps/maps.py +++ b/pythonbpf/maps/maps.py @@ -36,11 +36,14 @@ def output(self, data): 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 @@ def reserve(self, size: int, flags=0): 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 diff --git a/pythonbpf/maps/maps_pass.py b/pythonbpf/maps/maps_pass.py index 85837d72..87b9c03a 100644 --- a/pythonbpf/maps/maps_pass.py +++ b/pythonbpf/maps/maps_pass.py @@ -3,7 +3,7 @@ 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 @@ -46,7 +46,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): @@ -79,17 +79,28 @@ def _parse_map_params(rval, expected_args=None): return params -@MapProcessorRegistry.register("RingBuf") +@MapProcessorRegistry.register("RingBuffer") def process_ringbuf_map(map_name, rval, module): """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 @@ -103,7 +114,7 @@ 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 @@ -117,7 +128,7 @@ 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 diff --git a/pythonbpf/maps/maps_utils.py b/pythonbpf/maps/maps_utils.py index ee3ad083..eaa43b22 100644 --- a/pythonbpf/maps/maps_utils.py +++ b/pythonbpf/maps/maps_utils.py @@ -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: