mirror of
https://github.com/varun-r-mallya/Python-BPF.git
synced 2025-12-31 21:06:25 +00:00
Compare commits
26 Commits
5a8b64f1d9
...
copilot/fi
| Author | SHA1 | Date | |
|---|---|---|---|
| 00ce3014dc | |||
| e5b3b001ce | |||
| 19b42b9a19 | |||
| 9f5ec62383 | |||
| 7af54df7c0 | |||
| 573bbb350e | |||
| 64679f8072 | |||
| 5667facf23 | |||
| 4f8af16a17 | |||
| b84884162d | |||
| e9bb90cb70 | |||
| 9d76502d5a | |||
| a10da4a277 | |||
| 29e90601b7 | |||
| 56df05a93c | |||
| a55efc6469 | |||
| 64cd2d2fc2 | |||
| cbddc0aa96 | |||
| 209df33c8f | |||
| 7a56e5d0cd | |||
| 1d7a436c9f | |||
| 5eaeb3e921 | |||
| cd52d0d91b | |||
| df981be095 | |||
| 316c21c428 | |||
| c883d95655 |
@ -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",
|
||||
|
||||
@ -7,6 +7,7 @@ from pythonbpf.helper import HelperHandlerRegistry
|
||||
from pythonbpf.vmlinux_parser.dependency_node import Field
|
||||
from .expr import VmlinuxHandlerRegistry
|
||||
from pythonbpf.type_deducer import ctypes_to_ir
|
||||
from pythonbpf.maps import BPFMapType
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -25,7 +26,9 @@ def create_targets_and_rvals(stmt):
|
||||
return stmt.targets, [stmt.value]
|
||||
|
||||
|
||||
def handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab):
|
||||
def handle_assign_allocation(
|
||||
builder, stmt, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
):
|
||||
"""Handle memory allocation for assignment statements."""
|
||||
|
||||
logger.info(f"Handling assignment for allocation: {ast.dump(stmt)}")
|
||||
@ -55,7 +58,9 @@ def handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab):
|
||||
|
||||
# Determine type and allocate based on rval
|
||||
if isinstance(rval, ast.Call):
|
||||
_allocate_for_call(builder, var_name, rval, local_sym_tab, structs_sym_tab)
|
||||
_allocate_for_call(
|
||||
builder, var_name, rval, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
)
|
||||
elif isinstance(rval, ast.Constant):
|
||||
_allocate_for_constant(builder, var_name, rval, local_sym_tab)
|
||||
elif isinstance(rval, ast.BinOp):
|
||||
@ -74,7 +79,9 @@ def handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab):
|
||||
)
|
||||
|
||||
|
||||
def _allocate_for_call(builder, var_name, rval, local_sym_tab, structs_sym_tab):
|
||||
def _allocate_for_call(
|
||||
builder, var_name, rval, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
):
|
||||
"""Allocate memory for variable assigned from a call."""
|
||||
|
||||
if isinstance(rval.func, ast.Name):
|
||||
@ -116,15 +123,74 @@ def _allocate_for_call(builder, var_name, rval, local_sym_tab, structs_sym_tab):
|
||||
|
||||
elif isinstance(rval.func, ast.Attribute):
|
||||
# Map method calls - need double allocation for ptr handling
|
||||
_allocate_for_map_method(builder, var_name, local_sym_tab)
|
||||
_allocate_for_map_method(
|
||||
builder, var_name, rval, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
)
|
||||
|
||||
else:
|
||||
logger.warning(f"Unsupported call function type for {var_name}")
|
||||
|
||||
|
||||
def _allocate_for_map_method(builder, var_name, local_sym_tab):
|
||||
def _allocate_for_map_method(
|
||||
builder, var_name, rval, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
):
|
||||
"""Allocate memory for variable assigned from map method (double alloc)."""
|
||||
|
||||
map_name = rval.func.value.id
|
||||
method_name = rval.func.attr
|
||||
|
||||
# NOTE: We will have to special case HashMap.lookup which returns a pointer to value type
|
||||
# The value type can be a struct as well, so we need to handle that properly
|
||||
# This special casing is not ideal, as over time other map methods may need similar handling
|
||||
# But for now, we will just handle lookup specifically
|
||||
if map_name not in map_sym_tab:
|
||||
logger.error(f"Map '{map_name}' not found for allocation")
|
||||
return
|
||||
|
||||
if method_name != "lookup":
|
||||
# Fallback allocation for other map methods
|
||||
_allocate_for_map_method_fallback(builder, var_name, local_sym_tab)
|
||||
return
|
||||
|
||||
map_params = map_sym_tab[map_name].params
|
||||
if map_params["type"] != BPFMapType.HASH:
|
||||
logger.warning(
|
||||
"Map method lookup used on non-hash map, using fallback allocation"
|
||||
)
|
||||
_allocate_for_map_method_fallback(builder, var_name, local_sym_tab)
|
||||
return
|
||||
|
||||
value_type = map_params["value"]
|
||||
# Determine IR type for value
|
||||
if isinstance(value_type, str) and value_type in structs_sym_tab:
|
||||
struct_info = structs_sym_tab[value_type]
|
||||
value_ir_type = struct_info.ir_type
|
||||
else:
|
||||
value_ir_type = ctypes_to_ir(value_type)
|
||||
|
||||
if value_ir_type is None:
|
||||
logger.warning(
|
||||
f"Could not determine IR type for map value '{value_type}', using fallback allocation"
|
||||
)
|
||||
_allocate_for_map_method_fallback(builder, var_name, local_sym_tab)
|
||||
return
|
||||
|
||||
# Main variable (pointer to pointer)
|
||||
ir_type = ir.PointerType(ir.IntType(64))
|
||||
var = builder.alloca(ir_type, name=var_name)
|
||||
local_sym_tab[var_name] = LocalSymbol(var, ir_type)
|
||||
# Temporary variable for computed values
|
||||
tmp_ir_type = value_ir_type
|
||||
var_tmp = builder.alloca(tmp_ir_type, name=f"{var_name}_tmp")
|
||||
local_sym_tab[f"{var_name}_tmp"] = LocalSymbol(var_tmp, tmp_ir_type)
|
||||
logger.info(
|
||||
f"Pre-allocated {var_name} and {var_name}_tmp for map method lookup of type {value_ir_type}"
|
||||
)
|
||||
|
||||
|
||||
def _allocate_for_map_method_fallback(builder, var_name, local_sym_tab):
|
||||
"""Fallback allocation for map method variable (i64* and i64**)."""
|
||||
|
||||
# Main variable (pointer to pointer)
|
||||
ir_type = ir.PointerType(ir.IntType(64))
|
||||
var = builder.alloca(ir_type, name=var_name)
|
||||
@ -135,7 +201,9 @@ def _allocate_for_map_method(builder, var_name, local_sym_tab):
|
||||
var_tmp = builder.alloca(tmp_ir_type, name=f"{var_name}_tmp")
|
||||
local_sym_tab[f"{var_name}_tmp"] = LocalSymbol(var_tmp, tmp_ir_type)
|
||||
|
||||
logger.info(f"Pre-allocated {var_name} and {var_name}_tmp for map method")
|
||||
logger.info(
|
||||
f"Pre-allocated {var_name} and {var_name}_tmp for map method (fallback)"
|
||||
)
|
||||
|
||||
|
||||
def _allocate_for_constant(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
|
||||
|
||||
@ -49,6 +49,10 @@ class DebugInfoGenerator:
|
||||
)
|
||||
return self._type_cache[key]
|
||||
|
||||
def get_uint8_type(self) -> Any:
|
||||
"""Get debug info for signed 8-bit integer"""
|
||||
return self.get_basic_type("char", 8, dc.DW_ATE_unsigned)
|
||||
|
||||
def get_int32_type(self) -> Any:
|
||||
"""Get debug info for signed 32-bit integer"""
|
||||
return self.get_basic_type("int", 32, dc.DW_ATE_signed)
|
||||
|
||||
@ -147,7 +147,9 @@ def allocate_mem(
|
||||
structs_sym_tab,
|
||||
)
|
||||
elif isinstance(stmt, ast.Assign):
|
||||
handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab)
|
||||
handle_assign_allocation(
|
||||
builder, stmt, local_sym_tab, map_sym_tab, structs_sym_tab
|
||||
)
|
||||
|
||||
allocate_temp_pool(builder, max_temps_needed, local_sym_tab)
|
||||
|
||||
|
||||
@ -12,6 +12,7 @@ from .helpers import (
|
||||
smp_processor_id,
|
||||
uid,
|
||||
skb_store_bytes,
|
||||
get_stack,
|
||||
XDP_DROP,
|
||||
XDP_PASS,
|
||||
)
|
||||
@ -83,6 +84,7 @@ __all__ = [
|
||||
"smp_processor_id",
|
||||
"uid",
|
||||
"skb_store_bytes",
|
||||
"get_stack",
|
||||
"XDP_DROP",
|
||||
"XDP_PASS",
|
||||
]
|
||||
|
||||
@ -12,11 +12,10 @@ from .helper_utils import (
|
||||
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,170 @@ 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,
|
||||
):
|
||||
"""
|
||||
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,
|
||||
@ -790,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
|
||||
|
||||
@ -287,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)
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -220,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_char_array_ptr_and_size(
|
||||
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)
|
||||
|
||||
@ -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"]
|
||||
|
||||
@ -1,22 +1,31 @@
|
||||
import logging
|
||||
from llvmlite import ir
|
||||
from pythonbpf.debuginfo import DebugInfoGenerator
|
||||
from .map_types import BPFMapType
|
||||
|
||||
logger: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
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)
|
||||
|
||||
logger.info(f"Creating debug info for map {map_name} with params {map_params}")
|
||||
uint_type = generator.get_uint32_type()
|
||||
ulong_type = generator.get_uint64_type()
|
||||
array_type = generator.create_array_type(
|
||||
uint_type, map_params.get("type", BPFMapType.UNSPEC).value
|
||||
)
|
||||
type_ptr = generator.create_pointer_type(array_type, 64)
|
||||
key_ptr = generator.create_pointer_type(
|
||||
array_type if "key_size" in map_params else ulong_type, 64
|
||||
array_type
|
||||
if "key_size" in map_params
|
||||
else _get_key_val_dbg_type(map_params.get("key"), generator, structs_sym_tab),
|
||||
64,
|
||||
)
|
||||
value_ptr = generator.create_pointer_type(
|
||||
array_type if "value_size" in map_params else ulong_type, 64
|
||||
array_type
|
||||
if "value_size" in map_params
|
||||
else _get_key_val_dbg_type(map_params.get("value"), generator, structs_sym_tab),
|
||||
64,
|
||||
)
|
||||
|
||||
elements_arr = []
|
||||
@ -64,7 +73,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)
|
||||
|
||||
@ -91,3 +106,65 @@ def create_ringbuf_debug_info(module, map_global, map_name, map_params):
|
||||
)
|
||||
map_global.set_metadata("dbg", global_var)
|
||||
return global_var
|
||||
|
||||
|
||||
def _get_key_val_dbg_type(name, generator, structs_sym_tab):
|
||||
"""Get the debug type for key/value based on type object"""
|
||||
|
||||
if not name:
|
||||
logger.warn("No name provided for key/value type, defaulting to uint64")
|
||||
return generator.get_uint64_type()
|
||||
|
||||
type_obj = structs_sym_tab.get(name)
|
||||
if type_obj:
|
||||
return _get_struct_debug_type(type_obj, generator, structs_sym_tab)
|
||||
|
||||
# Fallback to basic types
|
||||
logger.info(f"No struct named {name}, falling back to basic type")
|
||||
|
||||
# NOTE: Only handling int and long for now
|
||||
if name in ["c_int32", "c_uint32"]:
|
||||
return generator.get_uint32_type()
|
||||
|
||||
# Default fallback for now
|
||||
return generator.get_uint64_type()
|
||||
|
||||
|
||||
def _get_struct_debug_type(struct_obj, generator, structs_sym_tab):
|
||||
"""Recursively create debug type for struct"""
|
||||
elements_arr = []
|
||||
for fld in struct_obj.fields.keys():
|
||||
fld_type = struct_obj.field_type(fld)
|
||||
if isinstance(fld_type, ir.IntType):
|
||||
if fld_type.width == 32:
|
||||
fld_dbg_type = generator.get_uint32_type()
|
||||
else:
|
||||
# NOTE: Assuming 64-bit for all other int types
|
||||
fld_dbg_type = generator.get_uint64_type()
|
||||
elif isinstance(fld_type, ir.ArrayType):
|
||||
# NOTE: Array types have u8 elements only for now
|
||||
# Debug info generation should fail for other types
|
||||
elem_type = fld_type.element
|
||||
if isinstance(elem_type, ir.IntType) and elem_type.width == 8:
|
||||
char_type = generator.get_uint8_type()
|
||||
fld_dbg_type = generator.create_array_type(char_type, fld_type.count)
|
||||
else:
|
||||
logger.warning(
|
||||
f"Array element type {str(elem_type)} not supported for debug info, skipping"
|
||||
)
|
||||
continue
|
||||
else:
|
||||
# NOTE: Only handling int and char arrays for now
|
||||
logger.warning(
|
||||
f"Field type {str(fld_type)} not supported for debug info, skipping"
|
||||
)
|
||||
continue
|
||||
|
||||
member = generator.create_struct_member(
|
||||
fld, fld_dbg_type, struct_obj.field_size(fld)
|
||||
)
|
||||
elements_arr.append(member)
|
||||
struct_type = generator.create_struct_type(
|
||||
elements_arr, struct_obj.size, is_distinct=True
|
||||
)
|
||||
return struct_type
|
||||
|
||||
@ -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, params=map_params)
|
||||
|
||||
|
||||
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,48 @@ 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, structs_sym_tab
|
||||
)
|
||||
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 +121,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, structs_sym_tab)
|
||||
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 +135,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 +158,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,17 @@
|
||||
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
|
||||
params: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class MapProcessorRegistry:
|
||||
|
||||
42
tests/passing_tests/hash_map_struct.py
Normal file
42
tests/passing_tests/hash_map_struct.py
Normal file
@ -0,0 +1,42 @@
|
||||
from pythonbpf import bpf, section, struct, bpfglobal, compile, map
|
||||
from pythonbpf.maps import HashMap
|
||||
from pythonbpf.helper import pid
|
||||
from ctypes import c_void_p, c_int64
|
||||
|
||||
|
||||
@bpf
|
||||
@struct
|
||||
class val_type:
|
||||
counter: c_int64
|
||||
shizzle: c_int64
|
||||
|
||||
|
||||
@bpf
|
||||
@map
|
||||
def last() -> HashMap:
|
||||
return HashMap(key=val_type, value=c_int64, max_entries=16)
|
||||
|
||||
|
||||
@bpf
|
||||
@section("tracepoint/syscalls/sys_enter_clone")
|
||||
def hello_world(ctx: c_void_p) -> c_int64:
|
||||
obj = val_type()
|
||||
obj.counter, obj.shizzle = 42, 96
|
||||
t = last.lookup(obj)
|
||||
if t:
|
||||
print(f"Found existing entry: counter={obj.counter}, pid={t}")
|
||||
last.delete(obj)
|
||||
return 0 # type: ignore [return-value]
|
||||
val = pid()
|
||||
last.update(obj, val)
|
||||
print(f"Map updated!, {obj.counter}, {obj.shizzle}, {val}")
|
||||
return 0 # type: ignore [return-value]
|
||||
|
||||
|
||||
@bpf
|
||||
@bpfglobal
|
||||
def LICENSE() -> str:
|
||||
return "GPL"
|
||||
|
||||
|
||||
compile()
|
||||
Reference in New Issue
Block a user