diff --git a/examples/execve3.py b/examples/execve3.py index 2bcfe4b..c491145 100644 --- a/examples/execve3.py +++ b/examples/execve3.py @@ -34,9 +34,9 @@ def hello_again(ctx: c_void_p) -> c_int64: y = False if x > 0: if x < 2: - print("we prevailed") + print(f"we prevailed {x}") else: - print("we did not prevail") + print(f"we did not prevail {x}") ts = ktime() last().update(key, ts) diff --git a/pythonbpf/bpf_helper_handler.py b/pythonbpf/bpf_helper_handler.py index 165c8a6..7105d68 100644 --- a/pythonbpf/bpf_helper_handler.py +++ b/pythonbpf/bpf_helper_handler.py @@ -1,8 +1,9 @@ import ast from llvmlite import ir +from .expr_pass import eval_expr -def bpf_ktime_get_ns_emitter(call, module, builder, func): +def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, local_sym_tab=None): """ Emit LLVM IR for bpf_ktime_get_ns helper function call. """ @@ -62,10 +63,87 @@ def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No return result -def bpf_printk_emitter(call, module, builder, func): +def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None): if not hasattr(func, "_fmt_counter"): func._fmt_counter = 0 + if not call.args: + raise ValueError("print expects at least one argument") + + if isinstance(call.args[0], ast.JoinedStr): + fmt_parts = [] + exprs = [] + + for value in call.args[0].values: + if isinstance(value, ast.Constant): + if isinstance(value.value, str): + fmt_parts.append(value.value) + elif isinstance(value.value, int): + fmt_parts.append("%lld") + exprs.append(ir.Constant(ir.IntType(64), value.value)) + else: + raise NotImplementedError( + "Only string and integer constants are supported in f-string.") + elif isinstance(value, ast.FormattedValue): + # Assume int for now + fmt_parts.append("%d") + if isinstance(value.value, ast.Name): + exprs.append(value.value) + else: + raise NotImplementedError( + "Only simple variable names are supported in formatted values.") + else: + raise NotImplementedError( + "Unsupported value type in f-string.") + + fmt_str = "".join(fmt_parts) + "\n" + "\0" + fmt_name = f"{func.name}____fmt{func._fmt_counter}" + func._fmt_counter += 1 + + fmt_gvar = ir.GlobalVariable( + module, ir.ArrayType(ir.IntType(8), len(fmt_str)), name=fmt_name) + fmt_gvar.global_constant = True + fmt_gvar.initializer = ir.Constant( + ir.ArrayType(ir.IntType(8), len(fmt_str)), + bytearray(fmt_str.encode("utf8")) + ) + fmt_gvar.linkage = "internal" + fmt_gvar.align = 1 + + fmt_ptr = builder.bitcast(fmt_gvar, ir.PointerType()) + + args = [fmt_ptr, ir.Constant(ir.IntType(32), len(fmt_str))] + + # Only 3 args supported in bpf_printk + if len(exprs) > 3: + print( + "Warning: bpf_printk supports up to 3 arguments, extra arguments will be ignored.") + + for expr in exprs[:3]: + val = eval_expr(func, module, builder, expr, local_sym_tab, None) + if val: + if isinstance(val.type, ir.PointerType): + val = builder.ptrtoint(val, ir.IntType(64)) + elif isinstance(val.type, ir.IntType): + if val.type.width < 64: + val = builder.sext(val, ir.IntType(64)) + else: + print( + "Warning: Only integer and pointer types are supported in bpf_printk arguments. Others will be converted to 0.") + val = ir.Constant(ir.IntType(64), 0) + args.append(val) + else: + print( + "Warning: Failed to evaluate expression for bpf_printk argument. It will be converted to 0.") + args.append(ir.Constant(ir.IntType(64), 0)) + + 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) + return builder.call(fn_ptr, args, tail=True) + for arg in call.args: if isinstance(arg, ast.Constant) and isinstance(arg.value, str): fmt_str = arg.value + "\n" + "\0" @@ -93,6 +171,7 @@ def bpf_printk_emitter(call, module, builder, func): builder.call(fn_ptr, [fmt_ptr, ir.Constant( ir.IntType(32), len(fmt_str))], tail=True) + def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=None): """ Emit LLVM IR for bpf_map_update_elem helper function call. @@ -101,11 +180,11 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No if not call.args or len(call.args) < 2 or len(call.args) > 3: raise ValueError("Map update expects 2 or 3 arguments (key, value, flags), got " f"{len(call.args)}") - + key_arg = call.args[0] value_arg = call.args[1] flags_arg = call.args[2] if len(call.args) > 2 else None - + # Handle key if isinstance(key_arg, ast.Name): key_name = key_arg.id @@ -124,7 +203,7 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No else: raise NotImplementedError( "Only simple variable names and integer constants are supported as keys in map update.") - + # Handle value if isinstance(value_arg, ast.Name): value_name = value_arg.id @@ -143,7 +222,7 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No else: raise NotImplementedError( "Only simple variable names and integer constants are supported as values in map update.") - + # Handle flags argument (defaults to 0) if flags_arg is not None: if isinstance(flags_arg, ast.Constant) and isinstance(flags_arg.value, int): @@ -162,7 +241,7 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No "Only integer constants and simple variable names are supported as flags in map update.") else: flags_val = 0 - + if key_ptr is None or value_ptr is None: raise ValueError("Key pointer or value pointer is None.") @@ -173,20 +252,22 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No var_arg=False ) fn_ptr_type = ir.PointerType(fn_type) - + # helper id fn_addr = ir.Constant(ir.IntType(64), 2) fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) - + if isinstance(flags_val, int): flags_const = ir.Constant(ir.IntType(64), flags_val) else: flags_const = flags_val - result = builder.call(fn_ptr, [map_void_ptr, key_ptr, value_ptr, flags_const], tail=False) - + result = builder.call( + fn_ptr, [map_void_ptr, key_ptr, value_ptr, flags_const], tail=False) + return result + helper_func_list = { "lookup": bpf_map_lookup_elem_emitter, "print": bpf_printk_emitter, @@ -200,7 +281,7 @@ def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_ func_name = call.func.id if func_name in helper_func_list: # it is not a map method call - return helper_func_list[func_name](call, module, builder, func) + return helper_func_list[func_name](call, None, module, builder, func, local_sym_tab) else: raise NotImplementedError( f"Function {func_name} is not implemented as a helper function.") diff --git a/pythonbpf/expr_pass.py b/pythonbpf/expr_pass.py new file mode 100644 index 0000000..56f74a4 --- /dev/null +++ b/pythonbpf/expr_pass.py @@ -0,0 +1,49 @@ +import ast +from llvmlite import ir + + +def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab): + print(f"Evaluating expression: {expr}") + if isinstance(expr, ast.Name): + if expr.id in local_sym_tab: + var = local_sym_tab[expr.id] + val = builder.load(var) + return val + else: + print(f"Undefined variable {expr.id}") + return None + elif isinstance(expr, ast.Constant): + if isinstance(expr.value, int): + return ir.Constant(ir.IntType(64), expr.value) + elif isinstance(expr.value, bool): + return ir.Constant(ir.IntType(1), int(expr.value)) + else: + print("Unsupported constant type") + return None + elif isinstance(expr, ast.Call): + # delayed import to avoid circular dependency + from .bpf_helper_handler import helper_func_list, handle_helper_call + + if isinstance(expr.func, ast.Name): + # check for helpers first + if expr.func.id in helper_func_list: + return handle_helper_call( + expr, module, builder, func, local_sym_tab, map_sym_tab) + elif isinstance(expr.func, ast.Attribute): + if isinstance(expr.func.value, ast.Call) and isinstance(expr.func.value.func, ast.Name): + method_name = expr.func.attr + if method_name in helper_func_list: + return handle_helper_call( + expr, module, builder, func, local_sym_tab, map_sym_tab) + print("Unsupported expression evaluation") + return None + + +def handle_expr(func, module, builder, expr, local_sym_tab, map_sym_tab): + """Handle expression statements in the function body.""" + print(f"Handling expression: {ast.dump(expr)}") + call = expr.value + if isinstance(call, ast.Call): + eval_expr(func, module, builder, call, local_sym_tab, map_sym_tab) + else: + print("Unsupported expression type") diff --git a/pythonbpf/functions_pass.py b/pythonbpf/functions_pass.py index 70ee989..4c5e335 100644 --- a/pythonbpf/functions_pass.py +++ b/pythonbpf/functions_pass.py @@ -1,9 +1,11 @@ from llvmlite import ir import ast + from .bpf_helper_handler import helper_func_list, handle_helper_call from .type_deducer import ctypes_to_ir from .binary_ops import handle_binary_op +from .expr_pass import eval_expr, handle_expr def get_probe_string(func_node): @@ -22,6 +24,7 @@ def get_probe_string(func_node): return arg.value return "helper" + def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab): """Handle assignment statements in the function body.""" if len(stmt.targets) != 1: @@ -96,54 +99,12 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab): else: print("Unsupported assignment call function type") elif isinstance(rval, ast.BinOp): - handle_binary_op(rval, module, builder, func, local_sym_tab, map_sym_tab) + handle_binary_op(rval, module, builder, func, + local_sym_tab, map_sym_tab) else: print("Unsupported assignment value type") -def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab): - if isinstance(expr, ast.Name): - if expr.id in local_sym_tab: - var = local_sym_tab[expr.id] - val = builder.load(var) - return val - else: - print(f"Undefined variable {expr.id}") - return None - elif isinstance(expr, ast.Constant): - if isinstance(expr.value, int): - return ir.Constant(ir.IntType(64), expr.value) - elif isinstance(expr.value, bool): - return ir.Constant(ir.IntType(1), int(expr.value)) - else: - print("Unsupported constant type") - return None - elif isinstance(expr, ast.Call): - if isinstance(expr.func, ast.Name): - # check for helpers first - if expr.func.id in helper_func_list: - return handle_helper_call( - expr, module, builder, func, local_sym_tab, map_sym_tab) - elif isinstance(expr.func, ast.Attribute): - if isinstance(expr.func.value, ast.Call) and isinstance(expr.func.value.func, ast.Name): - method_name = expr.func.attr - if method_name in helper_func_list: - return handle_helper_call( - expr, module, builder, func, local_sym_tab, map_sym_tab) - print("Unsupported expression evaluation") - return None - - -def handle_expr(func, module, builder, expr, local_sym_tab, map_sym_tab): - """Handle expression statements in the function body.""" - print(f"Handling expression: {ast.dump(expr)}") - call = expr.value - if isinstance(call, ast.Call): - eval_expr(func, module, builder, call, local_sym_tab, map_sym_tab) - else: - print("Unsupported expression type") - - def handle_cond(func, module, builder, cond, local_sym_tab, map_sym_tab): if isinstance(cond, ast.Constant): if isinstance(cond.value, bool):