mirror of
https://github.com/varun-r-mallya/Python-BPF.git
synced 2025-12-31 21:06:25 +00:00
Add deref(), add delete helper, refactor pre-alloc
This commit is contained in:
@ -24,24 +24,19 @@ def hello(ctx: c_void_p) -> c_int32:
|
|||||||
def hello_again(ctx: c_void_p) -> c_int64:
|
def hello_again(ctx: c_void_p) -> c_int64:
|
||||||
print("exited")
|
print("exited")
|
||||||
key = 0
|
key = 0
|
||||||
|
delta = 0
|
||||||
tsp = last().lookup(key)
|
tsp = last().lookup(key)
|
||||||
# if tsp:
|
if True:
|
||||||
# delta = (bpf_ktime_get_ns() - tsp.value)
|
delta = ktime()
|
||||||
# if delta < 1000000000:
|
ddelta = deref(delta)
|
||||||
# print("execve called within last second")
|
if ddelta < 1000000000:
|
||||||
# last().delete(key)
|
print("execve called within last second")
|
||||||
x = 1
|
last().delete(key)
|
||||||
y = False
|
|
||||||
if x > 0:
|
|
||||||
if x < 2:
|
|
||||||
print(f"we prevailed {x}")
|
|
||||||
else:
|
|
||||||
print(f"we did not prevail {x}")
|
|
||||||
ts = ktime()
|
ts = ktime()
|
||||||
last().update(key, ts)
|
last().update(key, ts)
|
||||||
|
|
||||||
st = "st"
|
# st = "st"
|
||||||
last().update(key, ts)
|
# last().update(key, ts)
|
||||||
|
|
||||||
keena = 2 + 1
|
keena = 2 + 1
|
||||||
# below breaks
|
# below breaks
|
||||||
@ -50,11 +45,32 @@ def hello_again(ctx: c_void_p) -> c_int64:
|
|||||||
keema = 8 * 9
|
keema = 8 * 9
|
||||||
keesa = 10 - 11
|
keesa = 10 - 11
|
||||||
keeda = 10 / 5
|
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)
|
return c_int64(0)
|
||||||
|
|
||||||
|
|
||||||
@bpf
|
@bpf
|
||||||
@bpfglobal
|
@bpfglobal
|
||||||
def LICENSE() -> str:
|
def LICENSE() -> str:
|
||||||
return "GPL"
|
return "GPL"
|
||||||
|
|
||||||
|
|
||||||
compile()
|
compile()
|
||||||
|
|||||||
@ -1,11 +1,15 @@
|
|||||||
import ast
|
import ast
|
||||||
from llvmlite import ir
|
from llvmlite import ir
|
||||||
|
|
||||||
|
|
||||||
def handle_binary_op(rval, module, builder, func, local_sym_tab, map_sym_tab):
|
def handle_binary_op(rval, module, builder, func, local_sym_tab, map_sym_tab):
|
||||||
|
print(module)
|
||||||
left = rval.left
|
left = rval.left
|
||||||
right = rval.right
|
right = rval.right
|
||||||
op = rval.op
|
op = rval.op
|
||||||
|
|
||||||
|
# In case of pointers, we'll deref once.
|
||||||
|
|
||||||
if isinstance(left, ast.Name):
|
if isinstance(left, ast.Name):
|
||||||
left = local_sym_tab[left.id]
|
left = local_sym_tab[left.id]
|
||||||
elif isinstance(left, ast.Constant):
|
elif isinstance(left, ast.Constant):
|
||||||
@ -20,6 +24,8 @@ def handle_binary_op(rval, module, builder, func, local_sym_tab, map_sym_tab):
|
|||||||
else:
|
else:
|
||||||
SyntaxError("Unsupported right operand type")
|
SyntaxError("Unsupported right operand type")
|
||||||
|
|
||||||
|
print(f"left is {left}, right is {right}, op is {op}")
|
||||||
|
|
||||||
if isinstance(op, ast.Add):
|
if isinstance(op, ast.Add):
|
||||||
result = builder.add(left, right)
|
result = builder.add(left, right)
|
||||||
elif isinstance(op, ast.Sub):
|
elif isinstance(op, ast.Sub):
|
||||||
|
|||||||
@ -268,11 +268,67 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No
|
|||||||
return result
|
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 = {
|
helper_func_list = {
|
||||||
"lookup": bpf_map_lookup_elem_emitter,
|
"lookup": bpf_map_lookup_elem_emitter,
|
||||||
"print": bpf_printk_emitter,
|
"print": bpf_printk_emitter,
|
||||||
"ktime": bpf_ktime_get_ns_emitter,
|
"ktime": bpf_ktime_get_ns_emitter,
|
||||||
"update": bpf_map_update_elem_emitter,
|
"update": bpf_map_update_elem_emitter,
|
||||||
|
"delete": bpf_map_delete_elem_emitter,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@ -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
|
from .bpf_helper_handler import helper_func_list, handle_helper_call
|
||||||
|
|
||||||
if isinstance(expr.func, ast.Name):
|
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:
|
if expr.func.id in helper_func_list:
|
||||||
return handle_helper_call(
|
return handle_helper_call(
|
||||||
expr, module, builder, func, local_sym_tab, map_sym_tab)
|
expr, module, builder, func, local_sym_tab, map_sym_tab)
|
||||||
|
|||||||
@ -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])
|
builder.store(val, local_sym_tab[var_name])
|
||||||
# local_sym_tab[var_name] = var
|
# local_sym_tab[var_name] = var
|
||||||
print(f"Assigned constant {rval.func.id} to {var_name}")
|
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:
|
else:
|
||||||
print(f"Unsupported assignment call type: {call_type}")
|
print(f"Unsupported assignment call type: {call_type}")
|
||||||
elif isinstance(rval.func, ast.Attribute):
|
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
|
return did_return
|
||||||
|
|
||||||
|
|
||||||
def process_func_body(module, builder, func_node, func, ret_type, map_sym_tab):
|
def allocate_mem(module, builder, body, func, ret_type, map_sym_tab, local_sym_tab):
|
||||||
"""Process the body of a bpf function"""
|
for stmt in body:
|
||||||
# TODO: A lot. We just have print -> bpf_trace_printk for now
|
if isinstance(stmt, ast.If):
|
||||||
did_return = False
|
if stmt.body:
|
||||||
|
local_sym_tab = allocate_mem(
|
||||||
local_sym_tab = {}
|
module, builder, stmt.body, func, ret_type, map_sym_tab, local_sym_tab)
|
||||||
|
if stmt.orelse:
|
||||||
# pre-allocate dynamic variables
|
local_sym_tab = allocate_mem(
|
||||||
for stmt in func_node.body:
|
module, builder, stmt.orelse, func, ret_type, map_sym_tab, local_sym_tab)
|
||||||
if isinstance(stmt, ast.Assign):
|
elif isinstance(stmt, ast.Assign):
|
||||||
if len(stmt.targets) != 1:
|
if len(stmt.targets) != 1:
|
||||||
print("Unsupported multiassignment")
|
print("Unsupported multiassignment")
|
||||||
continue
|
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
|
var.align = ir_type.width // 8
|
||||||
print(
|
print(
|
||||||
f"Pre-allocated variable {var_name} for helper")
|
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):
|
elif isinstance(rval.func, ast.Attribute):
|
||||||
ir_type = ir.PointerType(ir.IntType(64))
|
ir_type = ir.PointerType(ir.IntType(64))
|
||||||
var = builder.alloca(ir_type, name=var_name)
|
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")
|
print("Unsupported assignment value type")
|
||||||
continue
|
continue
|
||||||
local_sym_tab[var_name] = var
|
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()}")
|
print(f"Local symbol table: {local_sym_tab.keys()}")
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user