fix ruff errors

This commit is contained in:
2025-10-02 01:23:55 +05:30
parent 6831f11179
commit 59b3d6514b
5 changed files with 268 additions and 160 deletions

View File

@ -27,10 +27,7 @@ def hello(ctx: c_void_p) -> c_int32:
dataobj.pid = pid() dataobj.pid = pid()
dataobj.ts = ktime() dataobj.ts = ktime()
# dataobj.comm = strobj # dataobj.comm = strobj
print( print(f"clone called at {dataobj.ts} by pid" f"{dataobj.pid}, comm {strobj}")
f"clone called at {dataobj.ts} by pid"
f"{dataobj.pid}, comm {strobj}"
)
events.output(dataobj) events.output(dataobj)
return c_int32(0) return c_int32(0)

View File

@ -83,19 +83,16 @@ def handle_assign(
elif isinstance(rval, ast.Constant): elif isinstance(rval, ast.Constant):
if isinstance(rval.value, bool): if isinstance(rval.value, bool):
if rval.value: if rval.value:
builder.store(ir.Constant(ir.IntType(1), 1), builder.store(ir.Constant(ir.IntType(1), 1), local_sym_tab[var_name][0])
local_sym_tab[var_name][0])
else: else:
builder.store(ir.Constant(ir.IntType(1), 0), builder.store(ir.Constant(ir.IntType(1), 0), local_sym_tab[var_name][0])
local_sym_tab[var_name][0])
print(f"Assigned constant {rval.value} to {var_name}") print(f"Assigned constant {rval.value} to {var_name}")
elif isinstance(rval.value, int): elif isinstance(rval.value, int):
# Assume c_int64 for now # Assume c_int64 for now
# var = builder.alloca(ir.IntType(64), name=var_name) # var = builder.alloca(ir.IntType(64), name=var_name)
# var.align = 8 # var.align = 8
builder.store( builder.store(
ir.Constant(ir.IntType(64), ir.Constant(ir.IntType(64), rval.value), local_sym_tab[var_name][0]
rval.value), local_sym_tab[var_name][0]
) )
# local_sym_tab[var_name] = var # local_sym_tab[var_name] = var
print(f"Assigned constant {rval.value} to {var_name}") print(f"Assigned constant {rval.value} to {var_name}")
@ -110,8 +107,7 @@ def handle_assign(
global_str.linkage = "internal" global_str.linkage = "internal"
global_str.global_constant = True global_str.global_constant = True
global_str.initializer = str_const global_str.initializer = str_const
str_ptr = builder.bitcast( str_ptr = builder.bitcast(global_str, ir.PointerType(ir.IntType(8)))
global_str, ir.PointerType(ir.IntType(8)))
builder.store(str_ptr, local_sym_tab[var_name][0]) builder.store(str_ptr, local_sym_tab[var_name][0])
print(f"Assigned string constant '{rval.value}' to {var_name}") print(f"Assigned string constant '{rval.value}' to {var_name}")
else: else:
@ -130,8 +126,7 @@ def handle_assign(
# var = builder.alloca(ir_type, name=var_name) # var = builder.alloca(ir_type, name=var_name)
# var.align = ir_type.width // 8 # var.align = ir_type.width // 8
builder.store( builder.store(
ir.Constant( ir.Constant(ir_type, rval.args[0].value), local_sym_tab[var_name][0]
ir_type, rval.args[0].value), local_sym_tab[var_name][0]
) )
print( print(
f"Assigned {call_type} constant " f"Assigned {call_type} constant "
@ -177,8 +172,7 @@ def handle_assign(
ir_type = struct_info.ir_type ir_type = struct_info.ir_type
# var = builder.alloca(ir_type, name=var_name) # var = builder.alloca(ir_type, name=var_name)
# Null init # Null init
builder.store(ir.Constant(ir_type, None), builder.store(ir.Constant(ir_type, None), local_sym_tab[var_name][0])
local_sym_tab[var_name][0])
local_var_metadata[var_name] = call_type local_var_metadata[var_name] = call_type
print(f"Assigned struct {call_type} to {var_name}") print(f"Assigned struct {call_type} to {var_name}")
# local_sym_tab[var_name] = var # local_sym_tab[var_name] = var
@ -249,8 +243,7 @@ def handle_cond(func, module, builder, cond, local_sym_tab, map_sym_tab):
print(f"Undefined variable {cond.id} in condition") print(f"Undefined variable {cond.id} in condition")
return None return None
elif isinstance(cond, ast.Compare): elif isinstance(cond, ast.Compare):
lhs = eval_expr(func, module, builder, cond.left, lhs = eval_expr(func, module, builder, cond.left, local_sym_tab, map_sym_tab)[0]
local_sym_tab, map_sym_tab)[0]
if len(cond.ops) != 1 or len(cond.comparators) != 1: if len(cond.ops) != 1 or len(cond.comparators) != 1:
print("Unsupported complex comparison") print("Unsupported complex comparison")
return None return None
@ -303,8 +296,7 @@ def handle_if(
else: else:
else_block = None else_block = None
cond = handle_cond(func, module, builder, stmt.test, cond = handle_cond(func, module, builder, stmt.test, local_sym_tab, map_sym_tab)
local_sym_tab, map_sym_tab)
if else_block: if else_block:
builder.cbranch(cond, then_block, else_block) builder.cbranch(cond, then_block, else_block)
else: else:
@ -449,8 +441,7 @@ def allocate_mem(
ir_type = ctypes_to_ir(call_type) ir_type = ctypes_to_ir(call_type)
var = builder.alloca(ir_type, name=var_name) var = builder.alloca(ir_type, name=var_name)
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 HelperHandlerRegistry.has_handler(call_type): elif HelperHandlerRegistry.has_handler(call_type):
# Assume return type is int64 for now # Assume return type is int64 for now
ir_type = ir.IntType(64) ir_type = ir.IntType(64)
@ -671,8 +662,7 @@ def infer_return_type(func_node: ast.FunctionDef):
if found_type is None: if found_type is None:
found_type = t found_type = t
elif found_type != t: elif found_type != t:
raise ValueError("Conflicting return types:" raise ValueError("Conflicting return types:" f"{found_type} vs {t}")
f"{found_type} vs {t}")
return found_type or "None" return found_type or "None"
@ -709,8 +699,7 @@ def assign_string_to_array(builder, target_array_ptr, source_string_ptr, array_l
char = builder.load(src_ptr) char = builder.load(src_ptr)
# Store character in target # Store character in target
dst_ptr = builder.gep( dst_ptr = builder.gep(target_array_ptr, [ir.Constant(ir.IntType(32), 0), idx])
target_array_ptr, [ir.Constant(ir.IntType(32), 0), idx])
builder.store(char, dst_ptr) builder.store(char, dst_ptr)
# Increment counter # Increment counter
@ -721,6 +710,5 @@ def assign_string_to_array(builder, target_array_ptr, source_string_ptr, array_l
# Ensure null termination # Ensure null termination
last_idx = ir.Constant(ir.IntType(32), array_length - 1) last_idx = ir.Constant(ir.IntType(32), array_length - 1)
null_ptr = builder.gep( null_ptr = builder.gep(target_array_ptr, [ir.Constant(ir.IntType(32), 0), last_idx])
target_array_ptr, [ir.Constant(ir.IntType(32), 0), last_idx])
builder.store(ir.Constant(ir.IntType(8), 0), null_ptr) builder.store(ir.Constant(ir.IntType(8), 0), null_ptr)

View File

@ -1,2 +1,4 @@
from .helper_utils import HelperHandlerRegistry from .helper_utils import HelperHandlerRegistry
from .bpf_helper_handler import handle_helper_call from .bpf_helper_handler import handle_helper_call
__all__ = [HelperHandlerRegistry, handle_helper_call]

View File

@ -1,10 +1,14 @@
import ast import ast
from llvmlite import ir from llvmlite import ir
from enum import Enum from enum import Enum
from .helper_utils import (HelperHandlerRegistry, from .helper_utils import (
get_or_create_ptr_from_arg, get_flags_val, HelperHandlerRegistry,
handle_fstring_print, simple_string_print, get_or_create_ptr_from_arg,
get_data_ptr_and_size) get_flags_val,
handle_fstring_print,
simple_string_print,
get_data_ptr_and_size,
)
class BPFHelperID(Enum): class BPFHelperID(Enum):
@ -18,9 +22,16 @@ class BPFHelperID(Enum):
@HelperHandlerRegistry.register("ktime") @HelperHandlerRegistry.register("ktime")
def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, def bpf_ktime_get_ns_emitter(
local_sym_tab=None, struct_sym_tab=None, call,
local_var_metadata=None): 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.
""" """
@ -34,27 +45,34 @@ def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func,
@HelperHandlerRegistry.register("lookup") @HelperHandlerRegistry.register("lookup")
def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func, def bpf_map_lookup_elem_emitter(
local_sym_tab=None, struct_sym_tab=None, call,
local_var_metadata=None): 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.
""" """
if not call.args or len(call.args) != 1: if not call.args or len(call.args) != 1:
raise ValueError("Map lookup expects exactly one argument (key), got " raise ValueError(
f"{len(call.args)}") "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) key_ptr = get_or_create_ptr_from_arg(call.args[0], builder, local_sym_tab)
map_void_ptr = builder.bitcast(map_ptr, ir.PointerType()) map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
fn_type = ir.FunctionType( fn_type = ir.FunctionType(
ir.PointerType(), # Return type: void* ir.PointerType(), # Return type: void*
[ir.PointerType(), ir.PointerType()], # Args: (void*, void*) [ir.PointerType(), ir.PointerType()], # Args: (void*, void*)
var_arg=False var_arg=False,
) )
fn_ptr_type = ir.PointerType(fn_type) fn_ptr_type = ir.PointerType(fn_type)
fn_addr = ir.Constant(ir.IntType( fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_MAP_LOOKUP_ELEM.value)
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)
@ -63,33 +81,46 @@ def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func,
@HelperHandlerRegistry.register("print") @HelperHandlerRegistry.register("print")
def bpf_printk_emitter(call, map_ptr, module, builder, func, def bpf_printk_emitter(
local_sym_tab=None, struct_sym_tab=None, call,
local_var_metadata=None): 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.""" """Emit LLVM IR for bpf_printk helper function call."""
if not hasattr(func, "_fmt_counter"): if not hasattr(func, "_fmt_counter"):
func._fmt_counter = 0 func._fmt_counter = 0
if not call.args: if not call.args:
raise ValueError( raise ValueError("bpf_printk expects at least one argument (format string)")
"bpf_printk expects at least one argument (format string)")
args = [] args = []
if isinstance(call.args[0], ast.JoinedStr): if isinstance(call.args[0], ast.JoinedStr):
args = handle_fstring_print(call.args[0], module, builder, func, args = handle_fstring_print(
local_sym_tab, struct_sym_tab, call.args[0],
local_var_metadata) module,
elif (isinstance(call.args[0], ast.Constant) and builder,
isinstance(call.args[0].value, str)): 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 only supporting single arguments for now. # TODO: We are only supporting single arguments for now.
# In case of multiple args, the first one will be taken. # In case of multiple args, the first one will be taken.
args = simple_string_print(call.args[0].value, module, builder, func) args = simple_string_print(call.args[0].value, module, builder, func)
else: else:
raise NotImplementedError( raise NotImplementedError(
"Only simple strings or f-strings are supported in bpf_printk.") "Only simple strings or f-strings are supported in bpf_printk."
)
fn_type = ir.FunctionType( fn_type = ir.FunctionType(
ir.IntType(64), [ir.PointerType(), ir.IntType(32)], var_arg=True) ir.IntType(64), [ir.PointerType(), ir.IntType(32)], var_arg=True
)
fn_ptr_type = ir.PointerType(fn_type) fn_ptr_type = ir.PointerType(fn_type)
fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_PRINTK.value) fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_PRINTK.value)
fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type) fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
@ -99,18 +130,25 @@ def bpf_printk_emitter(call, map_ptr, module, builder, func,
@HelperHandlerRegistry.register("update") @HelperHandlerRegistry.register("update")
def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func, def bpf_map_update_elem_emitter(
local_sym_tab=None, struct_sym_tab=None, call,
local_var_metadata=None): 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.
Expected call signature: map.update(key, value, flags=0) Expected call signature: map.update(key, value, flags=0)
""" """
if (not call.args or if not call.args or len(call.args) < 2 or len(call.args) > 3:
len(call.args) < 2 or raise ValueError(
len(call.args) > 3): "Map update expects 2 or 3 args (key, value, flags), "
raise ValueError("Map update expects 2 or 3 args (key, value, flags), " f"got {len(call.args)}"
f"got {len(call.args)}") )
key_arg = call.args[0] key_arg = call.args[0]
value_arg = call.args[1] value_arg = call.args[1]
@ -124,12 +162,11 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func,
fn_type = ir.FunctionType( fn_type = ir.FunctionType(
ir.IntType(64), ir.IntType(64),
[ir.PointerType(), ir.PointerType(), ir.PointerType(), ir.IntType(64)], [ir.PointerType(), ir.PointerType(), ir.PointerType(), ir.IntType(64)],
var_arg=False var_arg=False,
) )
fn_ptr_type = ir.PointerType(fn_type) fn_ptr_type = ir.PointerType(fn_type)
fn_addr = ir.Constant(ir.IntType( fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_MAP_UPDATE_ELEM.value)
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):
@ -138,22 +175,31 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func,
flags_const = flags_val flags_const = flags_val
result = builder.call( result = builder.call(
fn_ptr, [map_void_ptr, key_ptr, value_ptr, flags_const], tail=False) fn_ptr, [map_void_ptr, key_ptr, value_ptr, flags_const], tail=False
)
return result, None return result, None
@HelperHandlerRegistry.register("delete") @HelperHandlerRegistry.register("delete")
def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func, def bpf_map_delete_elem_emitter(
local_sym_tab=None, struct_sym_tab=None, call,
local_var_metadata=None): 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.
Expected call signature: map.delete(key) Expected call signature: map.delete(key)
""" """
if not call.args or len(call.args) != 1: if not call.args or len(call.args) != 1:
raise ValueError("Map delete expects exactly one argument (key), got " raise ValueError(
f"{len(call.args)}") "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) key_ptr = get_or_create_ptr_from_arg(call.args[0], builder, local_sym_tab)
map_void_ptr = builder.bitcast(map_ptr, ir.PointerType()) map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
@ -161,12 +207,11 @@ def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func,
fn_type = ir.FunctionType( fn_type = ir.FunctionType(
ir.IntType(64), # Return type: int64 (status code) ir.IntType(64), # Return type: int64 (status code)
[ir.PointerType(), ir.PointerType()], # Args: (void*, void*) [ir.PointerType(), ir.PointerType()], # Args: (void*, void*)
var_arg=False var_arg=False,
) )
fn_ptr_type = ir.PointerType(fn_type) fn_ptr_type = ir.PointerType(fn_type)
fn_addr = ir.Constant(ir.IntType( fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_MAP_DELETE_ELEM.value)
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)
result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False) result = builder.call(fn_ptr, [map_void_ptr, key_ptr], tail=False)
@ -175,15 +220,21 @@ def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func,
@HelperHandlerRegistry.register("pid") @HelperHandlerRegistry.register("pid")
def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func, def bpf_get_current_pid_tgid_emitter(
local_sym_tab=None, struct_sym_tab=None, call,
local_var_metadata=None): 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( helper_id = ir.Constant(ir.IntType(64), BPFHelperID.BPF_GET_CURRENT_PID_TGID.value)
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)
@ -196,18 +247,26 @@ def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func,
@HelperHandlerRegistry.register("output") @HelperHandlerRegistry.register("output")
def bpf_perf_event_output_handler(call, map_ptr, module, builder, func, def bpf_perf_event_output_handler(
local_sym_tab=None, struct_sym_tab=None, call,
local_var_metadata=None): map_ptr,
module,
builder,
func,
local_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
if len(call.args) != 1: if len(call.args) != 1:
raise ValueError("Perf event output expects exactly one argument, " raise ValueError(
f"got {len(call.args)}") "Perf event output expects exactly one argument, " f"got {len(call.args)}"
)
data_arg = call.args[0] data_arg = call.args[0]
ctx_ptr = func.args[0] # First argument to the function is ctx 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, data_ptr, size_val = get_data_ptr_and_size(
struct_sym_tab, data_arg, local_sym_tab, struct_sym_tab, local_var_metadata
local_var_metadata) )
# BPF_F_CURRENT_CPU is -1 in 32 bit # BPF_F_CURRENT_CPU is -1 in 32 bit
flags_val = ir.Constant(ir.IntType(64), 0xFFFFFFFF) flags_val = ir.Constant(ir.IntType(64), 0xFFFFFFFF)
@ -216,36 +275,56 @@ def bpf_perf_event_output_handler(call, map_ptr, module, builder, func,
data_void_ptr = builder.bitcast(data_ptr, ir.PointerType()) data_void_ptr = builder.bitcast(data_ptr, ir.PointerType())
fn_type = ir.FunctionType( fn_type = ir.FunctionType(
ir.IntType(64), ir.IntType(64),
[ir.PointerType(ir.IntType(8)), ir.PointerType(), ir.IntType(64), [
ir.PointerType(), ir.IntType(64)], ir.PointerType(ir.IntType(8)),
var_arg=False ir.PointerType(),
ir.IntType(64),
ir.PointerType(),
ir.IntType(64),
],
var_arg=False,
) )
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), fn_addr = ir.Constant(ir.IntType(64), BPFHelperID.BPF_PERF_EVENT_OUTPUT.value)
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(
fn_ptr, fn_ptr, [ctx_ptr, map_void_ptr, flags_val, data_void_ptr, size_val], tail=False
[ctx_ptr, map_void_ptr, flags_val, data_void_ptr, size_val], )
tail=False)
return result, None return result, None
def handle_helper_call(call, module, builder, func, def handle_helper_call(
local_sym_tab=None, map_sym_tab=None, call,
struct_sym_tab=None, local_var_metadata=None): module,
builder,
func,
local_sym_tab=None,
map_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
"""Process a BPF helper function call and emit the appropriate LLVM IR.""" """Process a BPF helper function call and emit the appropriate LLVM IR."""
# Helper function to get map pointer and invoke handler # Helper function to get map pointer and invoke handler
def invoke_helper(method_name, map_ptr=None): def invoke_helper(method_name, map_ptr=None):
handler = HelperHandlerRegistry.get_handler(method_name) handler = HelperHandlerRegistry.get_handler(method_name)
if not handler: if not handler:
raise NotImplementedError( raise NotImplementedError(
f"Helper function '{method_name}' is not implemented.") f"Helper function '{method_name}' is not implemented."
return handler(call, map_ptr, module, builder, func, )
local_sym_tab, struct_sym_tab, local_var_metadata) return handler(
call,
map_ptr,
module,
builder,
func,
local_sym_tab,
struct_sym_tab,
local_var_metadata,
)
# Handle direct function calls (e.g., print(), ktime()) # Handle direct function calls (e.g., print(), ktime())
if isinstance(call.func, ast.Name): if isinstance(call.func, ast.Name):
@ -265,7 +344,8 @@ def handle_helper_call(call, module, builder, func,
map_name = value.id map_name = value.id
else: else:
raise NotImplementedError( raise NotImplementedError(
f"Unsupported map access pattern: {ast.dump(value)}") f"Unsupported map access pattern: {ast.dump(value)}"
)
# Verify map exists and get pointer # Verify map exists and get pointer
if not map_sym_tab or map_name not in map_sym_tab: if not map_sym_tab or map_name not in map_sym_tab:

View File

@ -8,14 +8,17 @@ logger = logging.getLogger(__name__)
class HelperHandlerRegistry: class HelperHandlerRegistry:
"""Registry for BPF helpers""" """Registry for BPF helpers"""
_handlers = {} _handlers = {}
@classmethod @classmethod
def register(cls, helper_name): def register(cls, helper_name):
"""Decorator to register a handler function for a helper""" """Decorator to register a handler function for a helper"""
def decorator(func): def decorator(func):
cls._handlers[helper_name] = func cls._handlers[helper_name] = func
return func return func
return decorator return decorator
@classmethod @classmethod
@ -55,7 +58,8 @@ def get_or_create_ptr_from_arg(arg, builder, local_sym_tab):
ptr = create_int_constant_ptr(arg.value, builder) ptr = create_int_constant_ptr(arg.value, builder)
else: else:
raise NotImplementedError( raise NotImplementedError(
"Only simple variable names are supported as args in map helpers.") "Only simple variable names are supported as args in map helpers."
)
return ptr return ptr
@ -69,13 +73,13 @@ def get_flags_val(arg, builder, local_sym_tab):
flags_ptr = local_sym_tab[arg.id][0] flags_ptr = local_sym_tab[arg.id][0]
return builder.load(flags_ptr) return builder.load(flags_ptr)
else: else:
raise ValueError( raise ValueError(f"Variable '{arg.id}' not found in local symbol table")
f"Variable '{arg.id}' not found in local symbol table")
elif isinstance(arg, ast.Constant) and isinstance(arg.value, int): elif isinstance(arg, ast.Constant) and isinstance(arg.value, int):
return arg.value return arg.value
raise NotImplementedError( raise NotImplementedError(
"Only var names or int consts are supported as map helpers flags.") "Only var names or int consts are supported as map helpers flags."
)
def simple_string_print(string_value, module, builder, func): def simple_string_print(string_value, module, builder, func):
@ -87,9 +91,15 @@ def simple_string_print(string_value, module, builder, func):
return args return args
def handle_fstring_print(joined_str, module, builder, func, def handle_fstring_print(
local_sym_tab=None, struct_sym_tab=None, joined_str,
local_var_metadata=None): module,
builder,
func,
local_sym_tab=None,
struct_sym_tab=None,
local_var_metadata=None,
):
"""Handle f-string formatting for bpf_printk emitter.""" """Handle f-string formatting for bpf_printk emitter."""
fmt_parts = [] fmt_parts = []
exprs = [] exprs = []
@ -100,25 +110,34 @@ def handle_fstring_print(joined_str, module, builder, func,
if isinstance(value, ast.Constant): if isinstance(value, ast.Constant):
_process_constant_in_fstring(value, fmt_parts, exprs) _process_constant_in_fstring(value, fmt_parts, exprs)
elif isinstance(value, ast.FormattedValue): elif isinstance(value, ast.FormattedValue):
_process_fval(value, fmt_parts, exprs, _process_fval(
local_sym_tab, struct_sym_tab, value,
local_var_metadata) fmt_parts,
exprs,
local_sym_tab,
struct_sym_tab,
local_var_metadata,
)
else: else:
raise NotImplementedError( raise NotImplementedError(f"Unsupported f-string value type: {type(value)}")
f"Unsupported f-string value type: {type(value)}")
fmt_str = "".join(fmt_parts) fmt_str = "".join(fmt_parts)
args = simple_string_print(fmt_str, module, builder, func) args = simple_string_print(fmt_str, module, builder, func)
# NOTE: Process expressions (limited to 3 due to BPF constraints) # NOTE: Process expressions (limited to 3 due to BPF constraints)
if len(exprs) > 3: if len(exprs) > 3:
logger.warning( logger.warning("bpf_printk supports up to 3 args, extra args will be ignored.")
"bpf_printk supports up to 3 args, extra args will be ignored.")
for expr in exprs[:3]: for expr in exprs[:3]:
arg_value = _prepare_expr_args(expr, func, module, builder, arg_value = _prepare_expr_args(
local_sym_tab, struct_sym_tab, expr,
local_var_metadata) func,
module,
builder,
local_sym_tab,
struct_sym_tab,
local_var_metadata,
)
args.append(arg_value) args.append(arg_value)
return args return args
@ -133,24 +152,31 @@ def _process_constant_in_fstring(cst, fmt_parts, exprs):
exprs.append(ir.Constant(ir.IntType(64), cst.value)) exprs.append(ir.Constant(ir.IntType(64), cst.value))
else: else:
raise NotImplementedError( raise NotImplementedError(
f"Unsupported constant type in f-string: {type(cst.value)}") f"Unsupported constant type in f-string: {type(cst.value)}"
)
def _process_fval(fval, fmt_parts, exprs, def _process_fval(
local_sym_tab, struct_sym_tab, fval, fmt_parts, exprs, local_sym_tab, struct_sym_tab, local_var_metadata
local_var_metadata): ):
"""Process formatted values in f-string.""" """Process formatted values in f-string."""
logger.debug(f"Processing formatted value: {ast.dump(fval)}") logger.debug(f"Processing formatted value: {ast.dump(fval)}")
if isinstance(fval.value, ast.Name): if isinstance(fval.value, ast.Name):
_process_name_in_fval(fval.value, fmt_parts, exprs, local_sym_tab) _process_name_in_fval(fval.value, fmt_parts, exprs, local_sym_tab)
elif isinstance(fval.value, ast.Attribute): elif isinstance(fval.value, ast.Attribute):
_process_attr_in_fval(fval.value, fmt_parts, exprs, _process_attr_in_fval(
local_sym_tab, struct_sym_tab, fval.value,
local_var_metadata) fmt_parts,
exprs,
local_sym_tab,
struct_sym_tab,
local_var_metadata,
)
else: else:
raise NotImplementedError( raise NotImplementedError(
f"Unsupported formatted value in f-string: {type(fval.value)}") f"Unsupported formatted value in f-string: {type(fval.value)}"
)
def _process_name_in_fval(name_node, fmt_parts, exprs, local_sym_tab): def _process_name_in_fval(name_node, fmt_parts, exprs, local_sym_tab):
@ -160,34 +186,39 @@ def _process_name_in_fval(name_node, fmt_parts, exprs, local_sym_tab):
_populate_fval(var_type, name_node, fmt_parts, exprs) _populate_fval(var_type, name_node, fmt_parts, exprs)
def _process_attr_in_fval(attr_node, fmt_parts, exprs, def _process_attr_in_fval(
local_sym_tab, struct_sym_tab, attr_node, fmt_parts, exprs, local_sym_tab, struct_sym_tab, local_var_metadata
local_var_metadata): ):
"""Process attribute nodes in formatted values.""" """Process attribute nodes in formatted values."""
if (isinstance(attr_node.value, ast.Name) and if (
local_sym_tab and attr_node.value.id in local_sym_tab): isinstance(attr_node.value, ast.Name)
and local_sym_tab
and attr_node.value.id in local_sym_tab
):
var_name = attr_node.value.id var_name = attr_node.value.id
field_name = attr_node.attr field_name = attr_node.attr
if not local_var_metadata or var_name not in local_var_metadata: if not local_var_metadata or var_name not in local_var_metadata:
raise ValueError( raise ValueError(
f"Metadata for '{var_name}' not found in local var metadata") f"Metadata for '{var_name}' not found in local var metadata"
)
var_type = local_var_metadata[var_name] var_type = local_var_metadata[var_name]
if var_type not in struct_sym_tab: if var_type not in struct_sym_tab:
raise ValueError( raise ValueError(
f"Struct '{var_type}' for '{var_name}' not in symbol table") f"Struct '{var_type}' for '{var_name}' not in symbol table"
)
struct_info = struct_sym_tab[var_type] struct_info = struct_sym_tab[var_type]
if field_name not in struct_info.fields: if field_name not in struct_info.fields:
raise ValueError( raise ValueError(f"Field '{field_name}' not found in struct '{var_type}'")
f"Field '{field_name}' not found in struct '{var_type}'")
field_type = struct_info.field_type(field_name) field_type = struct_info.field_type(field_name)
_populate_fval(field_type, attr_node, fmt_parts, exprs) _populate_fval(field_type, attr_node, fmt_parts, exprs)
else: else:
raise NotImplementedError( raise NotImplementedError(
"Only simple attribute on local vars is supported in f-strings.") "Only simple attribute on local vars is supported in f-strings."
)
def _populate_fval(ftype, node, fmt_parts, exprs): def _populate_fval(ftype, node, fmt_parts, exprs):
@ -202,14 +233,14 @@ def _populate_fval(ftype, node, fmt_parts, exprs):
exprs.append(node) exprs.append(node)
else: else:
raise NotImplementedError( raise NotImplementedError(
f"Unsupported integer width in f-string: {ftype.width}") f"Unsupported integer width in f-string: {ftype.width}"
)
elif ftype == ir.PointerType(ir.IntType(8)): elif ftype == ir.PointerType(ir.IntType(8)):
# NOTE: We assume i8* is a string # NOTE: We assume i8* is a string
fmt_parts.append("%s") fmt_parts.append("%s")
exprs.append(node) exprs.append(node)
else: else:
raise NotImplementedError( raise NotImplementedError(f"Unsupported field type in f-string: {ftype}")
f"Unsupported field type in f-string: {ftype}")
def _create_format_string_global(fmt_str, func, module, builder): def _create_format_string_global(fmt_str, func, module, builder):
@ -218,11 +249,11 @@ def _create_format_string_global(fmt_str, func, module, builder):
func._fmt_counter += 1 func._fmt_counter += 1
fmt_gvar = ir.GlobalVariable( fmt_gvar = ir.GlobalVariable(
module, ir.ArrayType(ir.IntType(8), len(fmt_str)), name=fmt_name) module, ir.ArrayType(ir.IntType(8), len(fmt_str)), name=fmt_name
)
fmt_gvar.global_constant = True fmt_gvar.global_constant = True
fmt_gvar.initializer = ir.Constant( fmt_gvar.initializer = ir.Constant(
ir.ArrayType(ir.IntType(8), len(fmt_str)), ir.ArrayType(ir.IntType(8), len(fmt_str)), bytearray(fmt_str.encode("utf8"))
bytearray(fmt_str.encode("utf8"))
) )
fmt_gvar.linkage = "internal" fmt_gvar.linkage = "internal"
fmt_gvar.align = 1 fmt_gvar.align = 1
@ -230,13 +261,20 @@ def _create_format_string_global(fmt_str, func, module, builder):
return builder.bitcast(fmt_gvar, ir.PointerType()) return builder.bitcast(fmt_gvar, ir.PointerType())
def _prepare_expr_args(expr, func, module, builder, def _prepare_expr_args(
local_sym_tab, struct_sym_tab, expr, func, module, builder, local_sym_tab, struct_sym_tab, local_var_metadata
local_var_metadata): ):
"""Evaluate and prepare an expression to use as an arg for bpf_printk.""" """Evaluate and prepare an expression to use as an arg for bpf_printk."""
val, _ = eval_expr(func, module, builder, expr, val, _ = eval_expr(
local_sym_tab, None, struct_sym_tab, func,
local_var_metadata) module,
builder,
expr,
local_sym_tab,
None,
struct_sym_tab,
local_var_metadata,
)
if val: if val:
if isinstance(val.type, ir.PointerType): if isinstance(val.type, ir.PointerType):
@ -246,19 +284,19 @@ def _prepare_expr_args(expr, func, module, builder,
val = builder.sext(val, ir.IntType(64)) val = builder.sext(val, ir.IntType(64))
else: else:
logger.warning( logger.warning(
"Only int and ptr supported in bpf_printk args. " "Only int and ptr supported in bpf_printk args. " "Others default to 0."
"Others default to 0.") )
val = ir.Constant(ir.IntType(64), 0) val = ir.Constant(ir.IntType(64), 0)
return val return val
else: else:
logger.warning( logger.warning(
"Failed to evaluate expression for bpf_printk argument. " "Failed to evaluate expression for bpf_printk argument. "
"It will be converted to 0.") "It will be converted to 0."
)
return ir.Constant(ir.IntType(64), 0) return ir.Constant(ir.IntType(64), 0)
def get_data_ptr_and_size(data_arg, local_sym_tab, struct_sym_tab, def get_data_ptr_and_size(data_arg, local_sym_tab, struct_sym_tab, local_var_metadata):
local_var_metadata):
"""Extract data pointer and size information for perf event output.""" """Extract data pointer and size information for perf event output."""
if isinstance(data_arg, ast.Name): if isinstance(data_arg, ast.Name):
data_name = data_arg.id data_name = data_arg.id
@ -266,7 +304,8 @@ def get_data_ptr_and_size(data_arg, local_sym_tab, struct_sym_tab,
data_ptr = local_sym_tab[data_name][0] data_ptr = local_sym_tab[data_name][0]
else: else:
raise ValueError( raise ValueError(
f"Data variable {data_name} not found in local symbol table.") f"Data variable {data_name} not found in local symbol table."
)
# Check if data_name is a struct # Check if data_name is a struct
if local_var_metadata and data_name in local_var_metadata: if local_var_metadata and data_name in local_var_metadata:
@ -277,12 +316,14 @@ def get_data_ptr_and_size(data_arg, local_sym_tab, struct_sym_tab,
return data_ptr, size_val return data_ptr, size_val
else: else:
raise ValueError( raise ValueError(
f"Struct {data_type} for {data_name} not in symbol table.") f"Struct {data_type} for {data_name} not in symbol table."
)
else: else:
raise ValueError( raise ValueError(
f"Metadata for variable {data_name} " f"Metadata for variable {data_name} "
"not found in local variable metadata.") "not found in local variable metadata."
)
else: else:
raise NotImplementedError( raise NotImplementedError(
"Only simple object names are supported " "Only simple object names are supported " "as data in perf event output."
"as data in perf event output.") )