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

Use HelperHandleRegitry

Browse Source
This commit is contained in:
Pragyansh Chaturvedi
2025-10-01 03:53:11 +05:30
parent 6cd07498fe
commit 61f6743f0a
5 changed files with 36 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
pythonbpf/expr_pass.py
View File

@ -23,7 +23,7 @@ def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab, structs_s
return None return None
elif isinstance(expr, ast.Call): elif isinstance(expr, ast.Call):
# delayed import to avoid circular dependency # delayed import to avoid circular dependency
from .helper.bpf_helper_handler import helper_func_list, handle_helper_call from pythonbpf.helper import HelperHandlerRegistry, handle_helper_call
if isinstance(expr.func, ast.Name): if isinstance(expr.func, ast.Name):
# check deref # check deref
@ -50,21 +50,21 @@ def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab, structs_s
return val, local_sym_tab[expr.args[0].id][1] return val, local_sym_tab[expr.args[0].id][1]
# check for helpers # check for helpers
if expr.func.id in helper_func_list: if expr.func.id in HelperHandlerRegistry._handlers:
return handle_helper_call( return handle_helper_call(
expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata) expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
elif isinstance(expr.func, ast.Attribute): elif isinstance(expr.func, ast.Attribute):
print(f"Handling method call: {ast.dump(expr.func)}") print(f"Handling method call: {ast.dump(expr.func)}")
if isinstance(expr.func.value, ast.Call) and isinstance(expr.func.value.func, ast.Name): if isinstance(expr.func.value, ast.Call) and isinstance(expr.func.value.func, ast.Name):
method_name = expr.func.attr method_name = expr.func.attr
if method_name in helper_func_list: if method_name in HelperHandlerRegistry._handlers:
return handle_helper_call( return handle_helper_call(
expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata) expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
elif isinstance(expr.func.value, ast.Name): elif isinstance(expr.func.value, ast.Name):
obj_name = expr.func.value.id obj_name = expr.func.value.id
method_name = expr.func.attr method_name = expr.func.attr
if obj_name in map_sym_tab: if obj_name in map_sym_tab:
if method_name in helper_func_list: if method_name in HelperHandlerRegistry._handlers:
return handle_helper_call( return handle_helper_call(
expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata) expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
elif isinstance(expr, ast.Attribute): elif isinstance(expr, ast.Attribute):

8
pythonbpf/functions_pass.py
View File

@ -2,7 +2,7 @@ from llvmlite import ir
import ast import ast
from .helper.bpf_helper_handler import helper_func_list, handle_helper_call from .helper import HelperHandlerRegistry, handle_helper_call
from .type_deducer import ctypes_to_ir from .type_deducer import ctypes_to_ir
from .binary_ops import handle_binary_op from .binary_ops import handle_binary_op
from .expr_pass import eval_expr, handle_expr from .expr_pass import eval_expr, handle_expr
@ -113,7 +113,7 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc
print(f"Assigned {call_type} constant " print(f"Assigned {call_type} constant "
f"{rval.args[0].value} to {var_name}") f"{rval.args[0].value} to {var_name}")
# local_sym_tab[var_name] = var # local_sym_tab[var_name] = var
elif call_type in helper_func_list: elif call_type in HelperHandlerRegistry._handlers:
# var = builder.alloca(ir.IntType(64), name=var_name) # var = builder.alloca(ir.IntType(64), name=var_name)
# var.align = 8 # var.align = 8
val = handle_helper_call( val = handle_helper_call(
@ -154,7 +154,7 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc
method_name = rval.func.attr method_name = rval.func.attr
if map_name in map_sym_tab: if map_name in map_sym_tab:
map_ptr = map_sym_tab[map_name] map_ptr = map_sym_tab[map_name]
if method_name in helper_func_list: if method_name in HelperHandlerRegistry._handlers:
val = handle_helper_call( val = handle_helper_call(
rval, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata) rval, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
# var = builder.alloca(ir.IntType(64), name=var_name) # var = builder.alloca(ir.IntType(64), name=var_name)
@ -344,7 +344,7 @@ def allocate_mem(module, builder, body, func, ret_type, map_sym_tab, local_sym_t
var.align = ir_type.width // 8 var.align = ir_type.width // 8
print( print(
f"Pre-allocated variable {var_name} of type {call_type}") f"Pre-allocated variable {var_name} of type {call_type}")
elif call_type in helper_func_list: elif call_type in HelperHandlerRegistry._handlers:
# Assume return type is int64 for now # Assume return type is int64 for now
ir_type = ir.IntType(64) ir_type = ir.IntType(64)
var = builder.alloca(ir_type, name=var_name) var = builder.alloca(ir_type, name=var_name)

2
pythonbpf/helper/__init__.py Normal file
View File

@ -0,0 +1,2 @@
from .helper_utils import HelperHandlerRegistry
from .bpf_helper_handler import handle_helper_call

48
pythonbpf/helper/bpf_helper_handler.py
View File

@ -2,6 +2,7 @@ import ast
from llvmlite import ir from llvmlite import ir
from pythonbpf.expr_pass import eval_expr from pythonbpf.expr_pass import eval_expr
from enum import Enum from enum import Enum
from .helper_utils import HelperHandlerRegistry
class BPFHelperID(Enum): class BPFHelperID(Enum):
@ -14,6 +15,7 @@ class BPFHelperID(Enum):
BPF_PERF_EVENT_OUTPUT = 25 BPF_PERF_EVENT_OUTPUT = 25
@HelperHandlerRegistry.register("ktime")
def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None): def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
""" """
Emit LLVM IR for bpf_ktime_get_ns helper function call. Emit LLVM IR for bpf_ktime_get_ns helper function call.
@ -27,6 +29,7 @@ def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, local_sym_tab
return result, ir.IntType(64) return result, ir.IntType(64)
@HelperHandlerRegistry.register("lookup")
def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None): def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
""" """
Emit LLVM IR for bpf_map_lookup_elem helper function call. Emit LLVM IR for bpf_map_lookup_elem helper function call.
@ -66,7 +69,8 @@ def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func, local_sym_
fn_ptr_type = ir.PointerType(fn_type) fn_ptr_type = ir.PointerType(fn_type)
# Helper ID 1 is bpf_map_lookup_elem # Helper ID 1 is bpf_map_lookup_elem
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_MAP_LOOKUP_ELEM.value) fn_addr = ir.Constant(ir.IntType(
64), BPFHelperID.BPF_MAP_LOOKUP_ELEM.value)
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False) result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False)
@ -74,6 +78,7 @@ def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func, local_sym_
return result, ir.PointerType() return result, ir.PointerType()
@HelperHandlerRegistry.register("print")
def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None): def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
if not hasattr(func, "_fmt_counter"): if not hasattr(func, "_fmt_counter"):
func._fmt_counter = 0 func._fmt_counter = 0
@ -233,6 +238,7 @@ def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None,
return None return None
@HelperHandlerRegistry.register("update")
def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None): def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
""" """
Emit LLVM IR for bpf_map_update_elem helper function call. Emit LLVM IR for bpf_map_update_elem helper function call.
@ -315,7 +321,8 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func, local_sym_
fn_ptr_type = ir.PointerType(fn_type) fn_ptr_type = ir.PointerType(fn_type)
# helper id # helper id
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_MAP_UPDATE_ELEM.value) fn_addr = ir.Constant(ir.IntType(
64), BPFHelperID.BPF_MAP_UPDATE_ELEM.value)
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
if isinstance(flags_val, int): if isinstance(flags_val, int):
@ -329,6 +336,7 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func, local_sym_
return result, None return result, None
@HelperHandlerRegistry.register("delete")
def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None): def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
""" """
Emit LLVM IR for bpf_map_delete_elem helper function call. Emit LLVM IR for bpf_map_delete_elem helper function call.
@ -375,7 +383,8 @@ def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func, local_sym_
fn_ptr_type = ir.PointerType(fn_type) fn_ptr_type = ir.PointerType(fn_type)
# Helper ID 3 is bpf_map_delete_elem # Helper ID 3 is bpf_map_delete_elem
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_MAP_DELETE_ELEM.value) fn_addr = ir.Constant(ir.IntType(
64), BPFHelperID.BPF_MAP_DELETE_ELEM.value)
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
# Call the helper function # Call the helper function
@ -384,12 +393,14 @@ def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func, local_sym_
return result, None return result, None
@HelperHandlerRegistry.register("pid")
def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None): def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
""" """
Emit LLVM IR for bpf_get_current_pid_tgid helper function call. Emit LLVM IR for bpf_get_current_pid_tgid helper function call.
""" """
# func is an arg to just have a uniform signature with other emitters # func is an arg to just have a uniform signature with other emitters
helper_id = ir.Constant(ir.IntType(64), BPFHelperID.BPF_GET_CURRENT_PID_TGID.value) helper_id = ir.Constant(ir.IntType(
64), BPFHelperID.BPF_GET_CURRENT_PID_TGID.value)
fn_type = ir.FunctionType(ir.IntType(64), [], var_arg=False) fn_type = ir.FunctionType(ir.IntType(64), [], var_arg=False)
fn_ptr_type = ir.PointerType(fn_type) fn_ptr_type = ir.PointerType(fn_type)
fn_ptr = builder.inttoptr(helper_id, fn_ptr_type) fn_ptr = builder.inttoptr(helper_id, fn_ptr_type)
@ -442,7 +453,8 @@ def bpf_perf_event_output_handler(call, map_ptr, module, builder, func, local_sy
fn_ptr_type = ir.PointerType(fn_type) fn_ptr_type = ir.PointerType(fn_type)
# helper id # helper id
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_PERF_EVENT_OUTPUT.value) fn_addr = ir.Constant(ir.IntType(
64), BPFHelperID.BPF_PERF_EVENT_OUTPUT.value)
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
result = builder.call( result = builder.call(
@ -453,24 +465,14 @@ def bpf_perf_event_output_handler(call, map_ptr, module, builder, func, local_sy
"Only simple object names are supported as data in perf event output.") "Only simple object names are supported as data in perf event output.")
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,
"pid": bpf_get_current_pid_tgid_emitter,
"output": bpf_perf_event_output_handler,
}
def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_tab=None, struct_sym_tab=None, local_var_metadata=None): def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
print(local_var_metadata) print(local_var_metadata)
if isinstance(call.func, ast.Name): if isinstance(call.func, ast.Name):
func_name = call.func.id func_name = call.func.id
if func_name in helper_func_list: hdl_func = HelperHandlerRegistry.get_handler(func_name)
if hdl_func:
# it is not a map method call # it is not a map method call
return helper_func_list[func_name](call, None, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata) return hdl_func(call, None, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata)
else: else:
raise NotImplementedError( raise NotImplementedError(
f"Function {func_name} is not implemented as a helper function.") f"Function {func_name} is not implemented as a helper function.")
@ -481,9 +483,10 @@ def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_
method_name = call.func.attr method_name = call.func.attr
if map_sym_tab and map_name in map_sym_tab: if map_sym_tab and map_name in map_sym_tab:
map_ptr = map_sym_tab[map_name] map_ptr = map_sym_tab[map_name]
if method_name in helper_func_list: hdl_func = HelperHandlerRegistry.get_handler(method_name)
if hdl_func:
print(local_var_metadata) print(local_var_metadata)
return helper_func_list[method_name]( return hdl_func(
call, map_ptr, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata) call, map_ptr, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata)
else: else:
raise NotImplementedError( raise NotImplementedError(
@ -496,8 +499,9 @@ def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_
method_name = call.func.attr method_name = call.func.attr
if map_sym_tab and obj_name in map_sym_tab: if map_sym_tab and obj_name in map_sym_tab:
map_ptr = map_sym_tab[obj_name] map_ptr = map_sym_tab[obj_name]
if method_name in helper_func_list: hdl_func = HelperHandlerRegistry.get_handler(method_name)
return helper_func_list[method_name]( if hdl_func:
return hdl_func(
call, map_ptr, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata) call, map_ptr, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata)
else: else:
raise NotImplementedError( raise NotImplementedError(

0
pythonbpf/helper/helpers.py → pythonbpf/helpers.py
View File