3 Commits

3 changed files with 86 additions and 73 deletions

View File

@ -14,10 +14,8 @@ def last() -> HashMap:
@bpf @bpf
@section("tracepoint/syscalls/sys_enter_sync") @section("tracepoint/syscalls/sys_enter_sync")
def do_trace(ctx: c_void_p) -> c_int64: def do_trace(ctx: c_void_p) -> c_int64:
ts_key = 0 ts_key, cnt_key = 0, 1
cnt_key = 1 tsp, cntp = last.lookup(ts_key), last.lookup(cnt_key)
tsp = last.lookup(ts_key)
cntp = last.lookup(cnt_key)
if not cntp: if not cntp:
last.update(cnt_key, 0) last.update(cnt_key, 0)
cntp = last.lookup(cnt_key) cntp = last.lookup(cnt_key)

View File

@ -22,44 +22,60 @@ class LocalSymbol:
yield self.metadata yield self.metadata
def create_targets_and_rvals(stmt):
"""Create lists of targets and right-hand values from an assignment statement."""
if isinstance(stmt.targets[0], ast.Tuple):
if not isinstance(stmt.value, ast.Tuple):
logger.warning("Mismatched multi-target assignment, skipping allocation")
return
targets, rvals = stmt.targets[0].elts, stmt.value.elts
if len(targets) != len(rvals):
logger.warning("length of LHS != length of RHS, skipping allocation")
return
return targets, rvals
return stmt.targets, [stmt.value]
def handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab): def handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab):
"""Handle memory allocation for assignment statements.""" """Handle memory allocation for assignment statements."""
# Validate assignment logger.info(f"Handling assignment for allocation: {ast.dump(stmt)}")
if len(stmt.targets) != 1:
logger.warning("Multi-target assignment not supported, skipping allocation")
return
target = stmt.targets[0] # NOTE: Support multi-target assignments (e.g.: a, b = 1, 2)
targets, rvals = create_targets_and_rvals(stmt)
# Skip non-name targets (e.g., struct field assignments) for target, rval in zip(targets, rvals):
if isinstance(target, ast.Attribute): # Skip non-name targets (e.g., struct field assignments)
logger.debug(f"Struct field assignment to {target.attr}, no allocation needed") if isinstance(target, ast.Attribute):
return logger.debug(
f"Struct field assignment to {target.attr}, no allocation needed"
)
continue
if not isinstance(target, ast.Name): if not isinstance(target, ast.Name):
logger.warning(f"Unsupported assignment target type: {type(target).__name__}") logger.warning(
return f"Unsupported assignment target type: {type(target).__name__}"
)
continue
var_name = target.id var_name = target.id
rval = stmt.value
# Skip if already allocated # Skip if already allocated
if var_name in local_sym_tab: if var_name in local_sym_tab:
logger.debug(f"Variable {var_name} already allocated, skipping") logger.debug(f"Variable {var_name} already allocated, skipping")
return continue
# Determine type and allocate based on rval # Determine type and allocate based on rval
if isinstance(rval, ast.Call): if isinstance(rval, ast.Call):
_allocate_for_call(builder, var_name, rval, local_sym_tab, structs_sym_tab) _allocate_for_call(builder, var_name, rval, local_sym_tab, structs_sym_tab)
elif isinstance(rval, ast.Constant): elif isinstance(rval, ast.Constant):
_allocate_for_constant(builder, var_name, rval, local_sym_tab) _allocate_for_constant(builder, var_name, rval, local_sym_tab)
elif isinstance(rval, ast.BinOp): elif isinstance(rval, ast.BinOp):
_allocate_for_binop(builder, var_name, local_sym_tab) _allocate_for_binop(builder, var_name, local_sym_tab)
else: else:
logger.warning( logger.warning(
f"Unsupported assignment value type for {var_name}: {type(rval).__name__}" f"Unsupported assignment value type for {var_name}: {type(rval).__name__}"
) )
def _allocate_for_call(builder, var_name, rval, local_sym_tab, structs_sym_tab): def _allocate_for_call(builder, var_name, rval, local_sym_tab, structs_sym_tab):

View File

@ -12,7 +12,11 @@ from pythonbpf.assign_pass import (
handle_variable_assignment, handle_variable_assignment,
handle_struct_field_assignment, handle_struct_field_assignment,
) )
from pythonbpf.allocation_pass import handle_assign_allocation, allocate_temp_pool from pythonbpf.allocation_pass import (
handle_assign_allocation,
allocate_temp_pool,
create_targets_and_rvals,
)
from .return_utils import handle_none_return, handle_xdp_return, is_xdp_name from .return_utils import handle_none_return, handle_xdp_return, is_xdp_name
from .function_metadata import get_probe_string, is_global_function, infer_return_type from .function_metadata import get_probe_string, is_global_function, infer_return_type
@ -140,48 +144,43 @@ def handle_assign(
): ):
"""Handle assignment statements in the function body.""" """Handle assignment statements in the function body."""
# TODO: Support this later # NOTE: Support multi-target assignments (e.g.: a, b = 1, 2)
# GH #37 targets, rvals = create_targets_and_rvals(stmt)
if len(stmt.targets) != 1:
logger.error("Multi-target assignment is not supported for now")
return
target = stmt.targets[0] for target, rval in zip(targets, rvals):
rval = stmt.value if isinstance(target, ast.Name):
# NOTE: Simple variable assignment case: x = 5
var_name = target.id
result = handle_variable_assignment(
func,
module,
builder,
var_name,
rval,
local_sym_tab,
map_sym_tab,
structs_sym_tab,
)
if not result:
logger.error(f"Failed to handle assignment to {var_name}")
continue
if isinstance(target, ast.Name): if isinstance(target, ast.Attribute):
# NOTE: Simple variable assignment case: x = 5 # NOTE: Struct field assignment case: pkt.field = value
var_name = target.id handle_struct_field_assignment(
result = handle_variable_assignment( func,
func, module,
module, builder,
builder, target,
var_name, rval,
rval, local_sym_tab,
local_sym_tab, map_sym_tab,
map_sym_tab, structs_sym_tab,
structs_sym_tab, )
) continue
if not result:
logger.error(f"Failed to handle assignment to {var_name}")
return
if isinstance(target, ast.Attribute): # Unsupported target type
# NOTE: Struct field assignment case: pkt.field = value logger.error(f"Unsupported assignment target: {ast.dump(target)}")
handle_struct_field_assignment(
func,
module,
builder,
target,
rval,
local_sym_tab,
map_sym_tab,
structs_sym_tab,
)
return
# Unsupported target type
logger.error(f"Unsupported assignment target: {ast.dump(target)}")
def handle_cond( def handle_cond(