diff --git a/examples/execve3.py b/examples/execve3.py index 42d4270..b83a31e 100644 --- a/examples/execve3.py +++ b/examples/execve3.py @@ -24,24 +24,19 @@ def hello(ctx: c_void_p) -> c_int32: def hello_again(ctx: c_void_p) -> c_int64: print("exited") key = 0 + delta = 0 tsp = last().lookup(key) -# if tsp: -# delta = (bpf_ktime_get_ns() - tsp.value) -# if delta < 1000000000: -# print("execve called within last second") -# last().delete(key) - x = 1 - y = False - if x > 0: - if x < 2: - print(f"we prevailed {x}") - else: - print(f"we did not prevail {x}") + if True: + delta = ktime() + ddelta = deref(delta) + if ddelta < 1000000000: + print("execve called within last second") + last().delete(key) ts = ktime() last().update(key, ts) - st = "st" - last().update(key, ts) +# st = "st" +# last().update(key, ts) keena = 2 + 1 # below breaks @@ -50,11 +45,32 @@ def hello_again(ctx: c_void_p) -> c_int64: keema = 8 * 9 keesa = 10 - 11 keeda = 10 / 5 +# x = 3 +# y = False +# if x > 0: +# if x < 5: +# print(f"we prevailed {x}") +# else: +# print(f"we did not prevail {x}") +# ts = ktime() +# last().update(key, ts) +# +# st = "st" +# last().update(key, ts) +# +# keena = 2 + 1 +# # below breaks +# # keela = keena + 1 +# keema = 8 * 9 +# keesa = 10 - 11 +# keeda = 10 / 5 return c_int64(0) + @bpf @bpfglobal def LICENSE() -> str: return "GPL" + compile() diff --git a/pythonbpf/binary_ops.py b/pythonbpf/binary_ops.py index 1f0e84b..7bf2c1b 100644 --- a/pythonbpf/binary_ops.py +++ b/pythonbpf/binary_ops.py @@ -1,11 +1,15 @@ import ast from llvmlite import ir + def handle_binary_op(rval, module, builder, func, local_sym_tab, map_sym_tab): + print(module) left = rval.left right = rval.right op = rval.op + # In case of pointers, we'll deref once. + if isinstance(left, ast.Name): left = local_sym_tab[left.id] elif isinstance(left, ast.Constant): @@ -20,6 +24,8 @@ def handle_binary_op(rval, module, builder, func, local_sym_tab, map_sym_tab): else: SyntaxError("Unsupported right operand type") + print(f"left is {left}, right is {right}, op is {op}") + if isinstance(op, ast.Add): result = builder.add(left, right) elif isinstance(op, ast.Sub): diff --git a/pythonbpf/bpf_helper_handler.py b/pythonbpf/bpf_helper_handler.py index 7105d68..45b7eb4 100644 --- a/pythonbpf/bpf_helper_handler.py +++ b/pythonbpf/bpf_helper_handler.py @@ -268,11 +268,67 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No return result +def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, local_sym_tab=None): + """ + Emit LLVM IR for bpf_map_delete_elem helper function call. + Expected call signature: map.delete(key) + """ + # Check for correct number of arguments + if not call.args or len(call.args) != 1: + raise ValueError("Map delete expects exactly 1 argument (key), got " + f"{len(call.args)}") + + key_arg = call.args[0] + + # Handle key argument + if isinstance(key_arg, ast.Name): + key_name = key_arg.id + if local_sym_tab and key_name in local_sym_tab: + key_ptr = local_sym_tab[key_name] + else: + raise ValueError( + f"Key variable {key_name} not found in local symbol table.") + elif isinstance(key_arg, ast.Constant) and isinstance(key_arg.value, int): + # Handle constant integer keys + key_val = key_arg.value + key_type = ir.IntType(64) + key_ptr = builder.alloca(key_type) + key_ptr.align = key_type.width // 8 + builder.store(ir.Constant(key_type, key_val), key_ptr) + else: + raise NotImplementedError( + "Only simple variable names and integer constants are supported as keys in map delete.") + + if key_ptr is None: + raise ValueError("Key pointer is None.") + + # Cast map pointer to void* + map_void_ptr = builder.bitcast(map_ptr, ir.PointerType()) + + # Define function type for bpf_map_delete_elem + fn_type = ir.FunctionType( + ir.IntType(64), # Return type: int64 (status code) + [ir.PointerType(), ir.PointerType()], # Args: (void*, void*) + var_arg=False + ) + fn_ptr_type = ir.PointerType(fn_type) + + # Helper ID 3 is bpf_map_delete_elem + fn_addr = ir.Constant(ir.IntType(64), 3) + fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) + + # Call the helper function + result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False) + + return result + + helper_func_list = { "lookup": bpf_map_lookup_elem_emitter, "print": bpf_printk_emitter, "ktime": bpf_ktime_get_ns_emitter, "update": bpf_map_update_elem_emitter, + "delete": bpf_map_delete_elem_emitter, } diff --git a/pythonbpf/expr_pass.py b/pythonbpf/expr_pass.py index 56f74a4..46e0c2d 100644 --- a/pythonbpf/expr_pass.py +++ b/pythonbpf/expr_pass.py @@ -25,7 +25,29 @@ def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab): from .bpf_helper_handler import helper_func_list, handle_helper_call if isinstance(expr.func, ast.Name): - # check for helpers first + # check deref + if expr.func.id == "deref": + print(f"Handling deref {ast.dump(expr)}") + if len(expr.args) != 1: + print("deref takes exactly one argument") + return None + arg = expr.args[0] + if isinstance(arg, ast.Call) and isinstance(arg.func, ast.Name) and arg.func.id == "deref": + print("Multiple deref not supported") + return None + if isinstance(arg, ast.Name): + if arg.id in local_sym_tab: + arg = local_sym_tab[arg.id] + else: + print(f"Undefined variable {arg.id}") + return None + if arg is None: + print("Failed to evaluate deref argument") + return None + val = builder.load(arg) + return val + + # check for helpers if expr.func.id in helper_func_list: return handle_helper_call( expr, module, builder, func, local_sym_tab, map_sym_tab) diff --git a/pythonbpf/functions_pass.py b/pythonbpf/functions_pass.py index 4c5e335..b811d98 100644 --- a/pythonbpf/functions_pass.py +++ b/pythonbpf/functions_pass.py @@ -79,6 +79,17 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab): builder.store(val, local_sym_tab[var_name]) # local_sym_tab[var_name] = var print(f"Assigned constant {rval.func.id} to {var_name}") + elif call_type == "deref" and len(rval.args) == 1: + print(f"Handling deref assignment {ast.dump(rval)}") + val = eval_expr(func, module, builder, rval, + local_sym_tab, map_sym_tab) + if val is None: + print("Failed to evaluate deref argument") + return + print(f"Dereferenced value: {val}, storing in {var_name}") + builder.store(val, local_sym_tab[var_name]) + # local_sym_tab[var_name] = var + print(f"Dereferenced and assigned to {var_name}") else: print(f"Unsupported assignment call type: {call_type}") elif isinstance(rval.func, ast.Attribute): @@ -227,16 +238,16 @@ def process_stmt(func, module, builder, stmt, local_sym_tab, map_sym_tab, did_re return did_return -def process_func_body(module, builder, func_node, func, ret_type, map_sym_tab): - """Process the body of a bpf function""" - # TODO: A lot. We just have print -> bpf_trace_printk for now - did_return = False - - local_sym_tab = {} - - # pre-allocate dynamic variables - for stmt in func_node.body: - if isinstance(stmt, ast.Assign): +def allocate_mem(module, builder, body, func, ret_type, map_sym_tab, local_sym_tab): + for stmt in body: + if isinstance(stmt, ast.If): + if stmt.body: + local_sym_tab = allocate_mem( + module, builder, stmt.body, func, ret_type, map_sym_tab, local_sym_tab) + if stmt.orelse: + local_sym_tab = allocate_mem( + module, builder, stmt.orelse, func, ret_type, map_sym_tab, local_sym_tab) + elif isinstance(stmt, ast.Assign): if len(stmt.targets) != 1: print("Unsupported multiassignment") continue @@ -262,6 +273,13 @@ def process_func_body(module, builder, func_node, func, ret_type, map_sym_tab): var.align = ir_type.width // 8 print( f"Pre-allocated variable {var_name} for helper") + elif call_type == "deref" and len(rval.args) == 1: + # Assume return type is int64 for now + ir_type = ir.IntType(64) + var = builder.alloca(ir_type, name=var_name) + var.align = ir_type.width // 8 + print( + f"Pre-allocated variable {var_name} for deref") elif isinstance(rval.func, ast.Attribute): ir_type = ir.PointerType(ir.IntType(64)) var = builder.alloca(ir_type, name=var_name) @@ -292,6 +310,19 @@ def process_func_body(module, builder, func_node, func, ret_type, map_sym_tab): print("Unsupported assignment value type") continue local_sym_tab[var_name] = var + return local_sym_tab + + +def process_func_body(module, builder, func_node, func, ret_type, map_sym_tab): + """Process the body of a bpf function""" + # TODO: A lot. We just have print -> bpf_trace_printk for now + did_return = False + + local_sym_tab = {} + + # pre-allocate dynamic variables + local_sym_tab = allocate_mem( + module, builder, func_node.body, func, ret_type, map_sym_tab, local_sym_tab) print(f"Local symbol table: {local_sym_tab.keys()}")