Add _handle_xdp_return

Author: r41k0u

pythonbpf/functions/functions_pass.py

@@ -9,7 +9,7 @@
 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")
 

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