mirror of
https://github.com/varun-r-mallya/Python-BPF.git
synced 2025-12-31 21:06:25 +00:00
281 lines
11 KiB
Python
281 lines
11 KiB
Python
import ast
|
|
from llvmlite import ir
|
|
from pythonbpf.expr_pass import eval_expr
|
|
from enum import Enum
|
|
from .helper_utils import HelperHandlerRegistry, get_or_create_ptr_from_arg, get_flags_val, _handle_fstring_print, _simple_string_print, _get_data_ptr_and_size
|
|
|
|
|
|
class BPFHelperID(Enum):
|
|
BPF_MAP_LOOKUP_ELEM = 1
|
|
BPF_MAP_UPDATE_ELEM = 2
|
|
BPF_MAP_DELETE_ELEM = 3
|
|
BPF_KTIME_GET_NS = 5
|
|
BPF_PRINTK = 6
|
|
BPF_GET_CURRENT_PID_TGID = 14
|
|
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):
|
|
"""
|
|
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), BPFHelperID.BPF_KTIME_GET_NS.value)
|
|
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, 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):
|
|
"""
|
|
Emit LLVM IR for bpf_map_lookup_elem helper function call.
|
|
"""
|
|
if not call.args or len(call.args) != 1:
|
|
raise ValueError("Map lookup expects exactly one argument (key), got "
|
|
f"{len(call.args)}")
|
|
key_ptr = get_or_create_ptr_from_arg(call.args[0], builder, local_sym_tab)
|
|
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)
|
|
|
|
fn_addr = ir.Constant(ir.IntType(
|
|
64), BPFHelperID.BPF_MAP_LOOKUP_ELEM.value)
|
|
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
|
|
|
|
result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False)
|
|
|
|
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):
|
|
"""Emit LLVM IR for bpf_printk helper function call."""
|
|
if not hasattr(func, "_fmt_counter"):
|
|
func._fmt_counter = 0
|
|
|
|
if not call.args:
|
|
raise ValueError(
|
|
"bpf_printk expects at least one argument (format string)")
|
|
|
|
args = []
|
|
if isinstance(call.args[0], ast.JoinedStr):
|
|
args = _handle_fstring_print(call.args[0], module, builder, func,
|
|
local_sym_tab, struct_sym_tab,
|
|
local_var_metadata)
|
|
elif isinstance(call.args[0], ast.Constant) and isinstance(call.args[0].value, str):
|
|
# TODO: We are onbly supporting single arguments for now.
|
|
# In case of multiple args, the first one will be taken.
|
|
args = _simple_string_print(call.args[0], module, builder, func)
|
|
else:
|
|
raise NotImplementedError(
|
|
"Only simple string literals or f-strings are supported in bpf_printk.")
|
|
|
|
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), BPFHelperID.BPF_PRINTK.value)
|
|
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
|
|
|
|
builder.call(fn_ptr, args, tail=True)
|
|
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):
|
|
"""
|
|
Emit LLVM IR for bpf_map_update_elem helper function call.
|
|
Expected call signature: map.update(key, value, flags=0)
|
|
"""
|
|
if (not call.args or
|
|
len(call.args) < 2 or
|
|
len(call.args) > 3):
|
|
raise ValueError("Map update expects 2 or 3 arguments (key, value, flags), got "
|
|
f"{len(call.args)}")
|
|
|
|
key_arg = call.args[0]
|
|
value_arg = call.args[1]
|
|
flags_arg = call.args[2] if len(call.args) > 2 else None
|
|
|
|
key_ptr = get_or_create_ptr_from_arg(key_arg, builder, local_sym_tab)
|
|
value_ptr = get_or_create_ptr_from_arg(value_arg, builder, local_sym_tab)
|
|
flags_val = get_flags_val(flags_arg, builder, local_sym_tab)
|
|
|
|
map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
|
|
fn_type = ir.FunctionType(
|
|
ir.IntType(64),
|
|
[ir.PointerType(), ir.PointerType(), ir.PointerType(), ir.IntType(64)],
|
|
var_arg=False
|
|
)
|
|
fn_ptr_type = ir.PointerType(fn_type)
|
|
|
|
fn_addr = ir.Constant(ir.IntType(
|
|
64), BPFHelperID.BPF_MAP_UPDATE_ELEM.value)
|
|
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
|
|
|
|
if isinstance(flags_val, int):
|
|
flags_const = ir.Constant(ir.IntType(64), flags_val)
|
|
else:
|
|
flags_const = flags_val
|
|
|
|
result = builder.call(
|
|
fn_ptr, [map_void_ptr, key_ptr, value_ptr, flags_const], tail=False)
|
|
|
|
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):
|
|
"""
|
|
Emit LLVM IR for bpf_map_delete_elem helper function call.
|
|
Expected call signature: map.delete(key)
|
|
"""
|
|
if not call.args or len(call.args) != 1:
|
|
raise ValueError("Map delete expects exactly one argument (key), got "
|
|
f"{len(call.args)}")
|
|
key_ptr = get_or_create_ptr_from_arg(call.args[0], builder, local_sym_tab)
|
|
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)
|
|
|
|
fn_addr = ir.Constant(ir.IntType(
|
|
64), BPFHelperID.BPF_MAP_DELETE_ELEM.value)
|
|
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
|
|
|
|
result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False)
|
|
|
|
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):
|
|
"""
|
|
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
|
|
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_ptr_type = ir.PointerType(fn_type)
|
|
fn_ptr = builder.inttoptr(helper_id, fn_ptr_type)
|
|
result = builder.call(fn_ptr, [], tail=False)
|
|
|
|
# Extract the lower 32 bits (PID) using bitwise AND with 0xFFFFFFFF
|
|
mask = ir.Constant(ir.IntType(64), 0xFFFFFFFF)
|
|
pid = builder.and_(result, mask)
|
|
return pid, ir.IntType(64)
|
|
|
|
|
|
def bpf_perf_event_output_handler(call, map_ptr, module, builder, func,
|
|
local_sym_tab=None, struct_sym_tab=None,
|
|
local_var_metadata=None):
|
|
if len(call.args) != 1:
|
|
raise ValueError("Perf event output expects exactly one argument (data), got "
|
|
f"{len(call.args)}")
|
|
data_arg = call.args[0]
|
|
ctx_ptr = func.args[0] # First argument to the function is ctx
|
|
|
|
data_ptr, size_val = _get_data_ptr_and_size(data_arg, local_sym_tab,
|
|
struct_sym_tab,
|
|
local_var_metadata)
|
|
|
|
# BPF_F_CURRENT_CPU is -1 in 32 bit
|
|
flags_val = ir.Constant(ir.IntType(64), 0xFFFFFFFF)
|
|
|
|
map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
|
|
data_void_ptr = builder.bitcast(data_ptr, ir.PointerType())
|
|
fn_type = ir.FunctionType(
|
|
ir.IntType(64),
|
|
[ir.PointerType(ir.IntType(8)), ir.PointerType(), ir.IntType(64),
|
|
ir.PointerType(), ir.IntType(64)],
|
|
var_arg=False
|
|
)
|
|
fn_ptr_type = ir.PointerType(fn_type)
|
|
|
|
# helper id
|
|
fn_addr = ir.Constant(ir.IntType(64),
|
|
BPFHelperID.BPF_PERF_EVENT_OUTPUT.value)
|
|
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
|
|
|
|
result = builder.call(
|
|
fn_ptr, [ctx_ptr, map_void_ptr, flags_val, data_void_ptr, size_val], tail=False)
|
|
return result, 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)
|
|
if isinstance(call.func, ast.Name):
|
|
func_name = call.func.id
|
|
hdl_func = HelperHandlerRegistry.get_handler(func_name)
|
|
if hdl_func:
|
|
# it is not a map method call
|
|
return hdl_func(call, None, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata)
|
|
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]
|
|
hdl_func = HelperHandlerRegistry.get_handler(method_name)
|
|
if hdl_func:
|
|
print(local_var_metadata)
|
|
return hdl_func(
|
|
call, map_ptr, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata)
|
|
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.")
|
|
elif isinstance(call.func.value, ast.Name):
|
|
obj_name = call.func.value.id
|
|
method_name = call.func.attr
|
|
if map_sym_tab and obj_name in map_sym_tab:
|
|
map_ptr = map_sym_tab[obj_name]
|
|
hdl_func = HelperHandlerRegistry.get_handler(method_name)
|
|
if hdl_func:
|
|
return hdl_func(
|
|
call, map_ptr, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata)
|
|
else:
|
|
raise NotImplementedError(
|
|
f"Map method {method_name} is not implemented as a helper function.")
|
|
else:
|
|
raise ValueError(
|
|
f"Map variable {obj_name} not found in symbol tables.")
|
|
else:
|
|
raise NotImplementedError(
|
|
"Attribute not supported for map method calls.")
|
|
return None
|