diff --git a/pythonbpf/bpf_helper_handler.py b/pythonbpf/bpf_helper_handler.py new file mode 100644 index 0000000..2f80468 --- /dev/null +++ b/pythonbpf/bpf_helper_handler.py @@ -0,0 +1,33 @@ +import ast +from llvmlite import ir + +def bpf_printk_emitter(call, module, builder, func): + # Handle print statement + for arg in call.args: + if isinstance(arg, ast.Constant) and isinstance(arg.value, str): + fmt_str = arg.value + "\n" + "\0" + # Create a global variable for the format string + fmt_gvar = ir.GlobalVariable( + module, ir.ArrayType(ir.IntType(8), len(fmt_str)), name=f"{func.name}____fmt") + fmt_gvar.global_constant = True + fmt_gvar.initializer = ir.Constant( # type: ignore + ir.ArrayType(ir.IntType(8), len(fmt_str)), + bytearray(fmt_str.encode("utf8")) + ) + fmt_gvar.linkage = "internal" + fmt_gvar.align = 1 # type: ignore + + # Cast the global variable to i8* + fmt_ptr = builder.bitcast( + fmt_gvar, ir.PointerType()) + + # Call bpf_trace_printk (assumed to be at address 6) + fn_type = ir.FunctionType(ir.IntType( + 64), [ir.PointerType(), ir.IntType(32)], var_arg=True) + fn_ptr_type = ir.PointerType(fn_type) + fn_addr = ir.Constant(ir.IntType(64), 6) + fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) + + # Call the function + builder.call(fn_ptr, [fmt_ptr, ir.Constant( + ir.IntType(32), len(fmt_str))], tail=True) \ No newline at end of file diff --git a/pythonbpf/functions_pass.py b/pythonbpf/functions_pass.py index c88be8c..9c579f2 100644 --- a/pythonbpf/functions_pass.py +++ b/pythonbpf/functions_pass.py @@ -1,5 +1,7 @@ from llvmlite import ir import ast + +from .bpf_helper_handler import bpf_printk_emitter from .type_deducer import ctypes_to_ir def get_probe_string(func_node): @@ -26,35 +28,7 @@ def process_func_body(module, builder, func_node, func, ret_type): if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Call): call = stmt.value if isinstance(call.func, ast.Name) and call.func.id == "print": - # Handle print statement - for arg in call.args: - if isinstance(arg, ast.Constant) and isinstance(arg.value, str): - fmt_str = arg.value + "\n" + "\0" - # Create a global variable for the format string - fmt_gvar = ir.GlobalVariable( - module, ir.ArrayType(ir.IntType(8), len(fmt_str)), name=f"{func.name}____fmt") - fmt_gvar.global_constant = True - fmt_gvar.initializer = ir.Constant( # type: ignore - ir.ArrayType(ir.IntType(8), len(fmt_str)), - bytearray(fmt_str.encode("utf8")) - ) - fmt_gvar.linkage = "internal" - fmt_gvar.align = 1 # type: ignore - - # Cast the global variable to i8* - fmt_ptr = builder.bitcast( - fmt_gvar, ir.PointerType()) - - # Call bpf_trace_printk (assumed to be at address 6) - fn_type = ir.FunctionType(ir.IntType( - 64), [ir.PointerType(), ir.IntType(32)], var_arg=True) - fn_ptr_type = ir.PointerType(fn_type) - fn_addr = ir.Constant(ir.IntType(64), 6) - fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) - - # Call the function - builder.call(fn_ptr, [fmt_ptr, ir.Constant( - ir.IntType(32), len(fmt_str))], tail=True) + bpf_printk_emitter(call, module, builder, func) elif isinstance(stmt, ast.Return): if stmt.value is None: builder.ret(ir.Constant(ir.IntType(32), 0))