varun-r-mallya/python-bpf
1
0
Fork 0
mirror of https://github.com/varun-r-mallya/Python-BPF.git synced 2025-12-31 21:06:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add deref(), add delete helper, refactor pre-alloc

Browse Source
This commit is contained in:
Pragyansh Chaturvedi
2025-09-12 04:24:40 +05:30
parent 0950d0550c
commit a09e4e1bb6
5 changed files with 156 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
examples/execve3.py
View File

@ -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()

6
pythonbpf/binary_ops.py
View File

@ -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):

56
pythonbpf/bpf_helper_handler.py
View File

@ -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,
}

24
pythonbpf/expr_pass.py
View File

@ -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)

51
pythonbpf/functions_pass.py
View File

@ -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()}")