mirror of
https://github.com/varun-r-mallya/Python-BPF.git
synced 2025-12-31 21:06:25 +00:00
133 lines
4.9 KiB
Python
133 lines
4.9 KiB
Python
import ast
|
|
from llvmlite import ir
|
|
|
|
|
|
def bpf_ktime_get_ns_emitter(call, module, builder, func):
|
|
"""
|
|
Emit LLVM IR for bpf_ktime_get_ns helper function call.
|
|
"""
|
|
# func is an arg to just have a uniform signature with other emitters
|
|
helper_id = ir.Constant(ir.IntType(64), 5)
|
|
fn_type = ir.FunctionType(ir.IntType(64), [], var_arg=False)
|
|
fn_ptr_type = ir.PointerType(fn_type)
|
|
fn_ptr = builder.inttoptr(helper_id, fn_ptr_type)
|
|
result = builder.call(fn_ptr, [], tail=False)
|
|
return result
|
|
|
|
|
|
def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, local_sym_tab=None):
|
|
"""
|
|
Emit LLVM IR for bpf_map_lookup_elem helper function call.
|
|
"""
|
|
|
|
if call.args and len(call.args) != 1:
|
|
raise ValueError("Map lookup expects exactly one argument, got "
|
|
f"{len(call.args)}")
|
|
key_arg = call.args[0]
|
|
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 // 8
|
|
builder.store(ir.Constant(key_type, key_val), key_ptr)
|
|
else:
|
|
raise NotImplementedError(
|
|
"Only simple variable names are supported as keys in map lookup.")
|
|
|
|
if key_ptr is None:
|
|
raise ValueError("Key pointer is None.")
|
|
|
|
map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
|
|
|
|
fn_type = ir.FunctionType(
|
|
ir.PointerType(), # Return type: void*
|
|
[ir.PointerType(), ir.PointerType()], # Args: (void*, void*)
|
|
var_arg=False
|
|
)
|
|
fn_ptr_type = ir.PointerType(fn_type)
|
|
|
|
# Helper ID 1 is bpf_map_lookup_elem
|
|
fn_addr = ir.Constant(ir.IntType(64), 1)
|
|
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
|
|
|
|
result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False)
|
|
|
|
return result
|
|
|
|
|
|
def bpf_printk_emitter(call, module, builder, func):
|
|
if not hasattr(func, "_fmt_counter"):
|
|
func._fmt_counter = 0
|
|
|
|
for arg in call.args:
|
|
if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
|
|
fmt_str = arg.value + "\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( # type: ignore
|
|
ir.ArrayType(ir.IntType(8), len(fmt_str)),
|
|
bytearray(fmt_str.encode("utf8"))
|
|
)
|
|
fmt_gvar.linkage = "internal"
|
|
fmt_gvar.align = 1 # type: ignore
|
|
|
|
fmt_ptr = builder.bitcast(fmt_gvar, ir.PointerType())
|
|
|
|
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)
|
|
|
|
builder.call(fn_ptr, [fmt_ptr, ir.Constant(
|
|
ir.IntType(32), len(fmt_str))], tail=True)
|
|
|
|
|
|
helper_func_list = {
|
|
"lookup": bpf_map_lookup_elem_emitter,
|
|
"print": bpf_printk_emitter,
|
|
"ktime": bpf_ktime_get_ns_emitter,
|
|
}
|
|
|
|
|
|
def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_tab=None):
|
|
if isinstance(call.func, ast.Name):
|
|
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)
|
|
else:
|
|
raise NotImplementedError(
|
|
f"Function {func_name} is not implemented as a helper function.")
|
|
elif isinstance(call.func, ast.Attribute):
|
|
# likely a map method call
|
|
if isinstance(call.func.value, ast.Call) and isinstance(call.func.value.func, ast.Name):
|
|
map_name = call.func.value.func.id
|
|
method_name = call.func.attr
|
|
if map_sym_tab and map_name in map_sym_tab:
|
|
map_ptr = map_sym_tab[map_name]
|
|
if method_name in helper_func_list:
|
|
return helper_func_list[method_name](
|
|
call, map_ptr, module, builder, local_sym_tab)
|
|
else:
|
|
raise NotImplementedError(
|
|
f"Map method {method_name} is not implemented as a helper function.")
|
|
else:
|
|
raise ValueError(
|
|
f"Map variable {map_name} not found in symbol tables.")
|
|
else:
|
|
raise NotImplementedError(
|
|
"Attribute not supported for map method calls.")
|