Add _handle_xdp_return

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