diff --git a/pythonbpf/functions/functions_pass.py b/pythonbpf/functions/functions_pass.py index 6eff2e7..7d86257 100644 --- a/pythonbpf/functions/functions_pass.py +++ b/pythonbpf/functions/functions_pass.py @@ -9,7 +9,7 @@ from pythonbpf.type_deducer import ctypes_to_ir from pythonbpf.binary_ops import handle_binary_op from pythonbpf.expr_pass import eval_expr, handle_expr -from .return_utils import _handle_none_return +from .return_utils import _handle_none_return, _handle_xdp_return logger = logging.getLogger(__name__) @@ -397,14 +397,7 @@ def handle_return(builder, stmt, local_sym_tab, ret_type): else: raise ValueError("Failed to evaluate return expression") elif isinstance(stmt.value, ast.Name): - if stmt.value.id == "XDP_PASS": - builder.ret(ir.Constant(ret_type, 2)) - return True - elif stmt.value.id == "XDP_DROP": - builder.ret(ir.Constant(ret_type, 1)) - return True - else: - raise ValueError("Failed to evaluate return expression") + return _handle_xdp_return(stmt, builder, ret_type) else: raise ValueError("Unsupported return value") diff --git a/pythonbpf/functions/return_utils.py b/pythonbpf/functions/return_utils.py index 12a2d01..74df7d8 100644 --- a/pythonbpf/functions/return_utils.py +++ b/pythonbpf/functions/return_utils.py @@ -1,11 +1,39 @@ import logging -import ir +import ast + +from llvmlite import ir logger: logging.Logger = logging.getLogger(__name__) +XDP_ACTIONS = { + "XDP_ABORTED": 0, + "XDP_DROP": 1, + "XDP_PASS": 2, + "XDP_TX": 3, + "XDP_REDIRECT": 4, +} + def _handle_none_return(builder) -> bool: """Handle return or return None -> returns 0.""" builder.ret(ir.Constant(ir.IntType(64), 0)) logger.debug("Generated default return: 0") return True + + +def _handle_xdp_return(stmt: ast.Return, builder, ret_type) -> bool: + """Handle XDP returns""" + if not isinstance(stmt.value, ast.Name): + return False + + action_name = stmt.value.id + + if action_name not in XDP_ACTIONS: + raise ValueError( + f"Unknown XDP action: {action_name}. Available: {XDP_ACTIONS.keys()}" + ) + + value = XDP_ACTIONS[action_name] + builder.ret(ir.Constant(ret_type, value)) + logger.debug(f"Generated XDP action return: {action_name} = {value}") + return True