mirror of
https://github.com/varun-r-mallya/Python-BPF.git
synced 2026-03-25 22:51:28 +00:00
790 lines
28 KiB
Python
790 lines
28 KiB
Python
"""
|
|
BPF function processing and LLVM IR generation.
|
|
|
|
This module handles the core function processing, converting Python function
|
|
definitions into LLVM IR for BPF programs. It manages local variables,
|
|
control flow, and statement processing.
|
|
"""
|
|
|
|
from llvmlite import ir
|
|
import ast
|
|
import logging
|
|
from typing import Any
|
|
from dataclasses import dataclass
|
|
|
|
from pythonbpf.helper import HelperHandlerRegistry, handle_helper_call
|
|
from pythonbpf.type_deducer import ctypes_to_ir
|
|
from pythonbpf.binary_ops import handle_binary_op
|
|
from pythonbpf.expr import eval_expr, handle_expr, convert_to_bool
|
|
|
|
from .return_utils import _handle_none_return, _handle_xdp_return, _is_xdp_name
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class LocalSymbol:
|
|
"""
|
|
Represents a local variable in a BPF function.
|
|
|
|
Attributes:
|
|
var: LLVM IR alloca instruction for the variable
|
|
ir_type: LLVM IR type of the variable
|
|
metadata: Optional metadata (e.g., struct type name)
|
|
"""
|
|
var: ir.AllocaInstr
|
|
ir_type: ir.Type
|
|
metadata: Any = None
|
|
|
|
def __iter__(self):
|
|
"""Support tuple unpacking of LocalSymbol."""
|
|
yield self.var
|
|
yield self.ir_type
|
|
yield self.metadata
|
|
|
|
|
|
def get_probe_string(func_node):
|
|
"""Extract the probe string from the decorator of the function node."""
|
|
# TODO: right now we have the whole string in the section decorator
|
|
# But later we can implement typed tuples for tracepoints and kprobes
|
|
# For helper functions, we return "helper"
|
|
|
|
for decorator in func_node.decorator_list:
|
|
if isinstance(decorator, ast.Name) and decorator.id == "bpfglobal":
|
|
return None
|
|
if isinstance(decorator, ast.Call) and isinstance(decorator.func, ast.Name):
|
|
if decorator.func.id == "section" and len(decorator.args) == 1:
|
|
arg = decorator.args[0]
|
|
if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
|
|
return arg.value
|
|
return "helper"
|
|
|
|
|
|
def handle_assign(
|
|
func, module, builder, stmt, map_sym_tab, local_sym_tab, structs_sym_tab
|
|
):
|
|
"""Handle assignment statements in the function body."""
|
|
if len(stmt.targets) != 1:
|
|
logger.info("Unsupported multiassignment")
|
|
return
|
|
|
|
num_types = ("c_int32", "c_int64", "c_uint32", "c_uint64")
|
|
|
|
target = stmt.targets[0]
|
|
logger.info(f"Handling assignment to {ast.dump(target)}")
|
|
if not isinstance(target, ast.Name) and not isinstance(target, ast.Attribute):
|
|
logger.info("Unsupported assignment target")
|
|
return
|
|
var_name = target.id if isinstance(target, ast.Name) else target.value.id
|
|
rval = stmt.value
|
|
if isinstance(target, ast.Attribute):
|
|
# struct field assignment
|
|
field_name = target.attr
|
|
if var_name in local_sym_tab:
|
|
struct_type = local_sym_tab[var_name].metadata
|
|
struct_info = structs_sym_tab[struct_type]
|
|
if field_name in struct_info.fields:
|
|
field_ptr = struct_info.gep(
|
|
builder, local_sym_tab[var_name].var, field_name
|
|
)
|
|
val = eval_expr(
|
|
func,
|
|
module,
|
|
builder,
|
|
rval,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
)
|
|
if isinstance(struct_info.field_type(field_name), ir.ArrayType) and val[
|
|
1
|
|
] == ir.PointerType(ir.IntType(8)):
|
|
# TODO: Figure it out, not a priority rn
|
|
# Special case for string assignment to char array
|
|
# str_len = struct_info["field_types"][field_idx].count
|
|
# assign_string_to_array(builder, field_ptr, val[0], str_len)
|
|
# print(f"Assigned to struct field {var_name}.{field_name}")
|
|
pass
|
|
if val is None:
|
|
logger.info("Failed to evaluate struct field assignment")
|
|
return
|
|
logger.info(field_ptr)
|
|
builder.store(val[0], field_ptr)
|
|
logger.info(f"Assigned to struct field {var_name}.{field_name}")
|
|
return
|
|
elif isinstance(rval, ast.Constant):
|
|
if isinstance(rval.value, bool):
|
|
if rval.value:
|
|
builder.store(
|
|
ir.Constant(ir.IntType(1), 1), local_sym_tab[var_name].var
|
|
)
|
|
else:
|
|
builder.store(
|
|
ir.Constant(ir.IntType(1), 0), local_sym_tab[var_name].var
|
|
)
|
|
logger.info(f"Assigned constant {rval.value} to {var_name}")
|
|
elif isinstance(rval.value, int):
|
|
# Assume c_int64 for now
|
|
# var = builder.alloca(ir.IntType(64), name=var_name)
|
|
# var.align = 8
|
|
builder.store(
|
|
ir.Constant(ir.IntType(64), rval.value), local_sym_tab[var_name].var
|
|
)
|
|
logger.info(f"Assigned constant {rval.value} to {var_name}")
|
|
elif isinstance(rval.value, str):
|
|
str_val = rval.value.encode("utf-8") + b"\x00"
|
|
str_const = ir.Constant(
|
|
ir.ArrayType(ir.IntType(8), len(str_val)), bytearray(str_val)
|
|
)
|
|
global_str = ir.GlobalVariable(
|
|
module, str_const.type, name=f"{var_name}_str"
|
|
)
|
|
global_str.linkage = "internal"
|
|
global_str.global_constant = True
|
|
global_str.initializer = str_const
|
|
str_ptr = builder.bitcast(global_str, ir.PointerType(ir.IntType(8)))
|
|
builder.store(str_ptr, local_sym_tab[var_name].var)
|
|
logger.info(f"Assigned string constant '{rval.value}' to {var_name}")
|
|
else:
|
|
logger.info("Unsupported constant type")
|
|
elif isinstance(rval, ast.Call):
|
|
if isinstance(rval.func, ast.Name):
|
|
call_type = rval.func.id
|
|
logger.info(f"Assignment call type: {call_type}")
|
|
if (
|
|
call_type in num_types
|
|
and len(rval.args) == 1
|
|
and isinstance(rval.args[0], ast.Constant)
|
|
and isinstance(rval.args[0].value, int)
|
|
):
|
|
ir_type = ctypes_to_ir(call_type)
|
|
# var = builder.alloca(ir_type, name=var_name)
|
|
# var.align = ir_type.width // 8
|
|
builder.store(
|
|
ir.Constant(ir_type, rval.args[0].value),
|
|
local_sym_tab[var_name].var,
|
|
)
|
|
logger.info(
|
|
f"Assigned {call_type} constant {rval.args[0].value} to {var_name}"
|
|
)
|
|
elif HelperHandlerRegistry.has_handler(call_type):
|
|
# var = builder.alloca(ir.IntType(64), name=var_name)
|
|
# var.align = 8
|
|
val = handle_helper_call(
|
|
rval,
|
|
module,
|
|
builder,
|
|
func,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
)
|
|
builder.store(val[0], local_sym_tab[var_name].var)
|
|
logger.info(f"Assigned constant {rval.func.id} to {var_name}")
|
|
elif call_type == "deref" and len(rval.args) == 1:
|
|
logger.info(f"Handling deref assignment {ast.dump(rval)}")
|
|
val = eval_expr(
|
|
func,
|
|
module,
|
|
builder,
|
|
rval,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
)
|
|
if val is None:
|
|
logger.info("Failed to evaluate deref argument")
|
|
return
|
|
logger.info(f"Dereferenced value: {val}, storing in {var_name}")
|
|
builder.store(val[0], local_sym_tab[var_name].var)
|
|
logger.info(f"Dereferenced and assigned to {var_name}")
|
|
elif call_type in structs_sym_tab and len(rval.args) == 0:
|
|
struct_info = structs_sym_tab[call_type]
|
|
ir_type = struct_info.ir_type
|
|
# var = builder.alloca(ir_type, name=var_name)
|
|
# Null init
|
|
builder.store(ir.Constant(ir_type, None), local_sym_tab[var_name].var)
|
|
logger.info(f"Assigned struct {call_type} to {var_name}")
|
|
else:
|
|
logger.info(f"Unsupported assignment call type: {call_type}")
|
|
elif isinstance(rval.func, ast.Attribute):
|
|
logger.info(f"Assignment call attribute: {ast.dump(rval.func)}")
|
|
if isinstance(rval.func.value, ast.Name):
|
|
if rval.func.value.id in map_sym_tab:
|
|
map_name = rval.func.value.id
|
|
method_name = rval.func.attr
|
|
if HelperHandlerRegistry.has_handler(method_name):
|
|
val = handle_helper_call(
|
|
rval,
|
|
module,
|
|
builder,
|
|
func,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
)
|
|
builder.store(val[0], local_sym_tab[var_name].var)
|
|
else:
|
|
# TODO: probably a struct access
|
|
logger.info(f"TODO STRUCT ACCESS {ast.dump(rval)}")
|
|
elif isinstance(rval.func.value, ast.Call) and isinstance(
|
|
rval.func.value.func, ast.Name
|
|
):
|
|
map_name = rval.func.value.func.id
|
|
method_name = rval.func.attr
|
|
if map_name in map_sym_tab:
|
|
if HelperHandlerRegistry.has_handler(method_name):
|
|
val = handle_helper_call(
|
|
rval,
|
|
module,
|
|
builder,
|
|
func,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
)
|
|
# var = builder.alloca(ir.IntType(64), name=var_name)
|
|
# var.align = 8
|
|
builder.store(val[0], local_sym_tab[var_name].var)
|
|
else:
|
|
logger.info("Unsupported assignment call structure")
|
|
else:
|
|
logger.info("Unsupported assignment call function type")
|
|
elif isinstance(rval, ast.BinOp):
|
|
handle_binary_op(rval, builder, var_name, local_sym_tab)
|
|
else:
|
|
logger.info("Unsupported assignment value type")
|
|
|
|
|
|
def handle_cond(
|
|
func, module, builder, cond, local_sym_tab, map_sym_tab, structs_sym_tab=None
|
|
):
|
|
"""
|
|
Evaluate a condition expression and convert it to a boolean value.
|
|
|
|
Args:
|
|
func: The LLVM IR function being built
|
|
module: The LLVM IR module
|
|
builder: LLVM IR builder
|
|
cond: The AST condition node to evaluate
|
|
local_sym_tab: Local symbol table
|
|
map_sym_tab: Map symbol table
|
|
structs_sym_tab: Struct symbol table
|
|
|
|
Returns:
|
|
LLVM IR boolean value representing the condition result
|
|
"""
|
|
val = eval_expr(
|
|
func, module, builder, cond, local_sym_tab, map_sym_tab, structs_sym_tab
|
|
)[0]
|
|
return convert_to_bool(builder, val)
|
|
|
|
|
|
def handle_if(
|
|
func, module, builder, stmt, map_sym_tab, local_sym_tab, structs_sym_tab=None
|
|
):
|
|
"""Handle if statements in the function body."""
|
|
logger.info("Handling if statement")
|
|
# start = builder.block.parent
|
|
then_block = func.append_basic_block(name="if.then")
|
|
merge_block = func.append_basic_block(name="if.end")
|
|
if stmt.orelse:
|
|
else_block = func.append_basic_block(name="if.else")
|
|
else:
|
|
else_block = None
|
|
|
|
cond = handle_cond(
|
|
func, module, builder, stmt.test, local_sym_tab, map_sym_tab, structs_sym_tab
|
|
)
|
|
if else_block:
|
|
builder.cbranch(cond, then_block, else_block)
|
|
else:
|
|
builder.cbranch(cond, then_block, merge_block)
|
|
|
|
builder.position_at_end(then_block)
|
|
for s in stmt.body:
|
|
process_stmt(
|
|
func, module, builder, s, local_sym_tab, map_sym_tab, structs_sym_tab, False
|
|
)
|
|
if not builder.block.is_terminated:
|
|
builder.branch(merge_block)
|
|
|
|
if else_block:
|
|
builder.position_at_end(else_block)
|
|
for s in stmt.orelse:
|
|
process_stmt(
|
|
func,
|
|
module,
|
|
builder,
|
|
s,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
False,
|
|
)
|
|
if not builder.block.is_terminated:
|
|
builder.branch(merge_block)
|
|
|
|
builder.position_at_end(merge_block)
|
|
|
|
|
|
def handle_return(builder, stmt, local_sym_tab, ret_type):
|
|
"""
|
|
Handle return statements in BPF functions.
|
|
|
|
Args:
|
|
builder: LLVM IR builder
|
|
stmt: The AST Return node
|
|
local_sym_tab: Local symbol table
|
|
ret_type: Expected return type
|
|
|
|
Returns:
|
|
True if a return was emitted, False otherwise
|
|
"""
|
|
logger.info(f"Handling return statement: {ast.dump(stmt)}")
|
|
if stmt.value is None:
|
|
return _handle_none_return(builder)
|
|
elif isinstance(stmt.value, ast.Name) and _is_xdp_name(stmt.value.id):
|
|
return _handle_xdp_return(stmt, builder, ret_type)
|
|
else:
|
|
val = eval_expr(
|
|
func=None,
|
|
module=None,
|
|
builder=builder,
|
|
expr=stmt.value,
|
|
local_sym_tab=local_sym_tab,
|
|
map_sym_tab={},
|
|
structs_sym_tab={},
|
|
)
|
|
logger.info(f"Evaluated return expression to {val}")
|
|
builder.ret(val[0])
|
|
return True
|
|
|
|
|
|
def process_stmt(
|
|
func,
|
|
module,
|
|
builder,
|
|
stmt,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
did_return,
|
|
ret_type=ir.IntType(64),
|
|
):
|
|
"""
|
|
Process a single statement in a BPF function.
|
|
|
|
Args:
|
|
func: The LLVM IR function being built
|
|
module: The LLVM IR module
|
|
builder: LLVM IR builder
|
|
stmt: The AST statement node to process
|
|
local_sym_tab: Local symbol table
|
|
map_sym_tab: Map symbol table
|
|
structs_sym_tab: Struct symbol table
|
|
did_return: Whether a return has been emitted
|
|
ret_type: Expected return type
|
|
|
|
Returns:
|
|
True if a return was emitted, False otherwise
|
|
"""
|
|
logger.info(f"Processing statement: {ast.dump(stmt)}")
|
|
if isinstance(stmt, ast.Expr):
|
|
handle_expr(
|
|
func,
|
|
module,
|
|
builder,
|
|
stmt,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
)
|
|
elif isinstance(stmt, ast.Assign):
|
|
handle_assign(
|
|
func, module, builder, stmt, map_sym_tab, local_sym_tab, structs_sym_tab
|
|
)
|
|
elif isinstance(stmt, ast.AugAssign):
|
|
raise SyntaxError("Augmented assignment not supported")
|
|
elif isinstance(stmt, ast.If):
|
|
handle_if(
|
|
func, module, builder, stmt, map_sym_tab, local_sym_tab, structs_sym_tab
|
|
)
|
|
elif isinstance(stmt, ast.Return):
|
|
did_return = handle_return(
|
|
builder,
|
|
stmt,
|
|
local_sym_tab,
|
|
ret_type,
|
|
)
|
|
return did_return
|
|
|
|
|
|
def allocate_mem(
|
|
module, builder, body, func, ret_type, map_sym_tab, local_sym_tab, structs_sym_tab
|
|
):
|
|
"""
|
|
Pre-allocate stack memory for local variables in a BPF function.
|
|
|
|
This function scans the function body and creates alloca instructions
|
|
for all local variables before processing the function statements.
|
|
|
|
Args:
|
|
module: The LLVM IR module
|
|
builder: LLVM IR builder
|
|
body: List of AST statements in the function body
|
|
func: The LLVM IR function being built
|
|
ret_type: Expected return type
|
|
map_sym_tab: Map symbol table
|
|
local_sym_tab: Local symbol table to populate
|
|
structs_sym_tab: Struct symbol table
|
|
|
|
Returns:
|
|
Updated local symbol table
|
|
"""
|
|
for stmt in body:
|
|
has_metadata = False
|
|
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,
|
|
structs_sym_tab,
|
|
)
|
|
if stmt.orelse:
|
|
local_sym_tab = allocate_mem(
|
|
module,
|
|
builder,
|
|
stmt.orelse,
|
|
func,
|
|
ret_type,
|
|
map_sym_tab,
|
|
local_sym_tab,
|
|
structs_sym_tab,
|
|
)
|
|
elif isinstance(stmt, ast.Assign):
|
|
if len(stmt.targets) != 1:
|
|
logger.info("Unsupported multiassignment")
|
|
continue
|
|
target = stmt.targets[0]
|
|
if not isinstance(target, ast.Name):
|
|
logger.info("Unsupported assignment target")
|
|
continue
|
|
var_name = target.id
|
|
rval = stmt.value
|
|
if var_name in local_sym_tab:
|
|
logger.info(f"Variable {var_name} already allocated")
|
|
continue
|
|
if isinstance(rval, ast.Call):
|
|
if isinstance(rval.func, ast.Name):
|
|
call_type = rval.func.id
|
|
if call_type in ("c_int32", "c_int64", "c_uint32", "c_uint64"):
|
|
ir_type = ctypes_to_ir(call_type)
|
|
var = builder.alloca(ir_type, name=var_name)
|
|
var.align = ir_type.width // 8
|
|
logger.info(
|
|
f"Pre-allocated variable {var_name} of type {call_type}"
|
|
)
|
|
elif HelperHandlerRegistry.has_handler(call_type):
|
|
# 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
|
|
logger.info(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
|
|
logger.info(f"Pre-allocated variable {var_name} for deref")
|
|
elif call_type in structs_sym_tab:
|
|
struct_info = structs_sym_tab[call_type]
|
|
ir_type = struct_info.ir_type
|
|
var = builder.alloca(ir_type, name=var_name)
|
|
has_metadata = True
|
|
logger.info(
|
|
f"Pre-allocated variable {var_name} for struct {call_type}"
|
|
)
|
|
elif isinstance(rval.func, ast.Attribute):
|
|
ir_type = ir.PointerType(ir.IntType(64))
|
|
var = builder.alloca(ir_type, name=var_name)
|
|
# var.align = ir_type.width // 8
|
|
logger.info(f"Pre-allocated variable {var_name} for map")
|
|
else:
|
|
logger.info("Unsupported assignment call function type")
|
|
continue
|
|
elif isinstance(rval, ast.Constant):
|
|
if isinstance(rval.value, bool):
|
|
ir_type = ir.IntType(1)
|
|
var = builder.alloca(ir_type, name=var_name)
|
|
var.align = 1
|
|
logger.info(f"Pre-allocated variable {var_name} of type c_bool")
|
|
elif isinstance(rval.value, int):
|
|
# Assume c_int64 for now
|
|
ir_type = ir.IntType(64)
|
|
var = builder.alloca(ir_type, name=var_name)
|
|
var.align = ir_type.width // 8
|
|
logger.info(f"Pre-allocated variable {var_name} of type c_int64")
|
|
elif isinstance(rval.value, str):
|
|
ir_type = ir.PointerType(ir.IntType(8))
|
|
var = builder.alloca(ir_type, name=var_name)
|
|
var.align = 8
|
|
logger.info(f"Pre-allocated variable {var_name} of type string")
|
|
else:
|
|
logger.info("Unsupported constant type")
|
|
continue
|
|
elif isinstance(rval, ast.BinOp):
|
|
# Assume c_int64 for now
|
|
ir_type = ir.IntType(64)
|
|
var = builder.alloca(ir_type, name=var_name)
|
|
var.align = ir_type.width // 8
|
|
logger.info(f"Pre-allocated variable {var_name} of type c_int64")
|
|
else:
|
|
logger.info("Unsupported assignment value type")
|
|
continue
|
|
|
|
if has_metadata:
|
|
local_sym_tab[var_name] = LocalSymbol(var, ir_type, call_type)
|
|
else:
|
|
local_sym_tab[var_name] = LocalSymbol(var, ir_type)
|
|
return local_sym_tab
|
|
|
|
|
|
def process_func_body(
|
|
module, builder, func_node, func, ret_type, map_sym_tab, structs_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,
|
|
structs_sym_tab,
|
|
)
|
|
|
|
logger.info(f"Local symbol table: {local_sym_tab.keys()}")
|
|
|
|
for stmt in func_node.body:
|
|
did_return = process_stmt(
|
|
func,
|
|
module,
|
|
builder,
|
|
stmt,
|
|
local_sym_tab,
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
did_return,
|
|
ret_type,
|
|
)
|
|
|
|
if not did_return:
|
|
builder.ret(ir.Constant(ir.IntType(64), 0))
|
|
|
|
|
|
def process_bpf_chunk(func_node, module, return_type, map_sym_tab, structs_sym_tab):
|
|
"""Process a single BPF chunk (function) and emit corresponding LLVM IR."""
|
|
|
|
func_name = func_node.name
|
|
|
|
ret_type = return_type
|
|
|
|
# TODO: parse parameters
|
|
param_types = []
|
|
if func_node.args.args:
|
|
# Assume first arg to be ctx
|
|
param_types.append(ir.PointerType())
|
|
|
|
func_ty = ir.FunctionType(ret_type, param_types)
|
|
func = ir.Function(module, func_ty, func_name)
|
|
|
|
func.linkage = "dso_local"
|
|
func.attributes.add("nounwind")
|
|
func.attributes.add("noinline")
|
|
func.attributes.add("optnone")
|
|
|
|
if func_node.args.args:
|
|
# Only look at the first argument for now
|
|
param = func.args[0]
|
|
param.add_attribute("nocapture")
|
|
|
|
probe_string = get_probe_string(func_node)
|
|
if probe_string is not None:
|
|
func.section = probe_string
|
|
|
|
block = func.append_basic_block(name="entry")
|
|
builder = ir.IRBuilder(block)
|
|
|
|
process_func_body(
|
|
module, builder, func_node, func, ret_type, map_sym_tab, structs_sym_tab
|
|
)
|
|
return func
|
|
|
|
|
|
def func_proc(tree, module, chunks, map_sym_tab, structs_sym_tab):
|
|
"""
|
|
Process all BPF function chunks and generate LLVM IR.
|
|
|
|
Args:
|
|
tree: The Python AST (not used in current implementation)
|
|
module: The LLVM IR module to add functions to
|
|
chunks: List of AST function nodes decorated with @bpf
|
|
map_sym_tab: Map symbol table
|
|
structs_sym_tab: Struct symbol table
|
|
"""
|
|
for func_node in chunks:
|
|
is_global = False
|
|
for decorator in func_node.decorator_list:
|
|
if isinstance(decorator, ast.Name) and decorator.id in (
|
|
"map",
|
|
"bpfglobal",
|
|
"struct",
|
|
):
|
|
is_global = True
|
|
break
|
|
if is_global:
|
|
continue
|
|
func_type = get_probe_string(func_node)
|
|
logger.info(f"Found probe_string of {func_node.name}: {func_type}")
|
|
|
|
process_bpf_chunk(
|
|
func_node,
|
|
module,
|
|
ctypes_to_ir(infer_return_type(func_node)),
|
|
map_sym_tab,
|
|
structs_sym_tab,
|
|
)
|
|
|
|
|
|
def infer_return_type(func_node: ast.FunctionDef):
|
|
"""
|
|
Infer the return type of a BPF function from annotations or return statements.
|
|
|
|
Args:
|
|
func_node: The AST function node
|
|
|
|
Returns:
|
|
String representation of the return type (e.g., 'c_int64')
|
|
|
|
Raises:
|
|
TypeError: If func_node is not a FunctionDef
|
|
"""
|
|
if not isinstance(func_node, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
|
raise TypeError("Expected ast.FunctionDef")
|
|
if func_node.returns is not None:
|
|
try:
|
|
return ast.unparse(func_node.returns)
|
|
except Exception:
|
|
node = func_node.returns
|
|
if isinstance(node, ast.Name):
|
|
return node.id
|
|
if isinstance(node, ast.Attribute):
|
|
return getattr(node, "attr", type(node).__name__)
|
|
try:
|
|
return str(node)
|
|
except Exception:
|
|
return type(node).__name__
|
|
found_type = None
|
|
|
|
def _expr_type(e):
|
|
"""Helper function to extract type from an expression."""
|
|
if e is None:
|
|
return "None"
|
|
if isinstance(e, ast.Constant):
|
|
return type(e.value).__name__
|
|
if isinstance(e, ast.Name):
|
|
return e.id
|
|
if isinstance(e, ast.Call):
|
|
f = e.func
|
|
if isinstance(f, ast.Name):
|
|
return f.id
|
|
if isinstance(f, ast.Attribute):
|
|
try:
|
|
return ast.unparse(f)
|
|
except Exception:
|
|
return getattr(f, "attr", type(f).__name__)
|
|
try:
|
|
return ast.unparse(f)
|
|
except Exception:
|
|
return type(f).__name__
|
|
if isinstance(e, ast.Attribute):
|
|
try:
|
|
return ast.unparse(e)
|
|
except Exception:
|
|
return getattr(e, "attr", type(e).__name__)
|
|
try:
|
|
return ast.unparse(e)
|
|
except Exception:
|
|
return type(e).__name__
|
|
|
|
for walked_node in ast.walk(func_node):
|
|
if isinstance(walked_node, ast.Return):
|
|
t = _expr_type(walked_node.value)
|
|
if found_type is None:
|
|
found_type = t
|
|
elif found_type != t:
|
|
raise ValueError(f"Conflicting return types: {found_type} vs {t}")
|
|
return found_type or "None"
|
|
|
|
|
|
# For string assignment to fixed-size arrays
|
|
|
|
|
|
def assign_string_to_array(builder, target_array_ptr, source_string_ptr, array_length):
|
|
"""
|
|
Copy a string (i8*) to a fixed-size array ([N x i8]*)
|
|
"""
|
|
# Create a loop to copy characters one by one
|
|
# entry_block = builder.block
|
|
copy_block = builder.append_basic_block("copy_char")
|
|
end_block = builder.append_basic_block("copy_end")
|
|
|
|
# Create loop counter
|
|
i = builder.alloca(ir.IntType(32))
|
|
builder.store(ir.Constant(ir.IntType(32), 0), i)
|
|
|
|
# Start the loop
|
|
builder.branch(copy_block)
|
|
|
|
# Copy loop
|
|
builder.position_at_end(copy_block)
|
|
idx = builder.load(i)
|
|
in_bounds = builder.icmp_unsigned(
|
|
"<", idx, ir.Constant(ir.IntType(32), array_length)
|
|
)
|
|
builder.cbranch(in_bounds, copy_block, end_block)
|
|
|
|
with builder.if_then(in_bounds):
|
|
# Load character from source
|
|
src_ptr = builder.gep(source_string_ptr, [idx])
|
|
char = builder.load(src_ptr)
|
|
|
|
# Store character in target
|
|
dst_ptr = builder.gep(target_array_ptr, [ir.Constant(ir.IntType(32), 0), idx])
|
|
builder.store(char, dst_ptr)
|
|
|
|
# Increment counter
|
|
next_idx = builder.add(idx, ir.Constant(ir.IntType(32), 1))
|
|
builder.store(next_idx, i)
|
|
|
|
builder.position_at_end(end_block)
|
|
|
|
# Ensure null termination
|
|
last_idx = ir.Constant(ir.IntType(32), array_length - 1)
|
|
null_ptr = builder.gep(target_array_ptr, [ir.Constant(ir.IntType(32), 0), last_idx])
|
|
builder.store(ir.Constant(ir.IntType(8), 0), null_ptr)
|