Merge branch 'master' into bcc_examples

Author: varun-r-mallya

README.md

@@ -40,6 +40,12 @@ Python-BPF is an LLVM IR generator for eBPF programs written in Python. It uses
 
 ---
 
+## Try It Out!
+Run
+```bash
+curl -s https://raw.githubusercontent.com/pythonbpf/Python-BPF/refs/heads/master/tools/setup.sh | sudo bash
+```
+
 ## Installation
 
 Dependencies:

pyproject.toml

@@ -4,12 +4,26 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pythonbpf"
-version = "0.1.4"
+version = "0.1.5"
 description = "Reduced Python frontend for eBPF"
 authors = [
   { name = "r41k0u", email="[email protected]" },
   { name = "varun-r-mallya", email="[email protected]" }
 ]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Operating System :: POSIX :: Linux",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+    "Topic :: System :: Operating System Kernels :: Linux",
+]
 readme = "README.md"
 license = {text = "Apache-2.0"}
 requires-python = ">=3.8"

pythonbpf/allocation_pass.py

@@ -5,6 +5,7 @@
 from dataclasses import dataclass
 from typing import Any
 from pythonbpf.helper import HelperHandlerRegistry
+from .expr import VmlinuxHandlerRegistry
 from pythonbpf.type_deducer import ctypes_to_ir
 
 logger = logging.getLogger(__name__)
@@ -64,6 +65,15 @@ def handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab):
         if var_name in local_sym_tab:
             logger.debug(f"Variable {var_name} already allocated, skipping")
             continue
+            
+        # When allocating a variable, check if it's a vmlinux struct type
+        if isinstance(stmt.value, ast.Name) and VmlinuxHandlerRegistry.is_vmlinux_struct(
+            stmt.value.id
+        ):
+            # Handle vmlinux struct allocation
+            # This requires more implementation
+            print(stmt.value)
+            pass
 
         # Determine type and allocate based on rval
         if isinstance(rval, ast.Call):
@@ -85,7 +95,6 @@ def handle_assign_allocation(builder, stmt, local_sym_tab, structs_sym_tab):
                 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):
     """Allocate memory for variable assigned from a call."""
 

pythonbpf/codegen.py

@@ -5,6 +5,8 @@
 from .maps import maps_proc
 from .structs import structs_proc
 from .vmlinux_parser import vmlinux_proc
+from pythonbpf.vmlinux_parser.vmlinux_exports_handler import VmlinuxHandler
+from .expr import VmlinuxHandlerRegistry
 from .globals_pass import (
     globals_list_creation,
     globals_processing,
@@ -19,10 +21,20 @@
 import tempfile
 from logging import Logger
 import logging
+import re
 
 logger: Logger = logging.getLogger(__name__)
 
-VERSION = "v0.1.4"
+VERSION = "v0.1.5"
+
+
+def finalize_module(original_str):
+    """After all IR generation is complete, we monkey patch btf_ama attribute"""
+
+    # Create a string with applied transformation of btf_ama attribute addition to BTF struct field accesses.
+    pattern = r'(@"llvm\.[^"]+:[^"]*" = external global i64, !llvm\.preserve\.access\.index ![0-9]+)'
+    replacement = r'\1 "btf_ama"'
+    return re.sub(pattern, replacement, original_str)
 
 
 def find_bpf_chunks(tree):
@@ -45,11 +57,14 @@ def processor(source_code, filename, module):
     for func_node in bpf_chunks:
         logger.info(f"Found BPF function/struct: {func_node.name}")
 
-    vmlinux_proc(tree, module)
+    vmlinux_symtab = vmlinux_proc(tree, module)
+    if vmlinux_symtab:
+        handler = VmlinuxHandler.initialize(vmlinux_symtab)
+        VmlinuxHandlerRegistry.set_handler(handler)
+
     populate_global_symbol_table(tree, module)
     license_processing(tree, module)
     globals_processing(tree, module)
-
     structs_sym_tab = structs_proc(tree, module, bpf_chunks)
     map_sym_tab = maps_proc(tree, module, bpf_chunks)
     func_proc(tree, module, bpf_chunks, map_sym_tab, structs_sym_tab)
@@ -122,10 +137,12 @@ def compile_to_ir(filename: str, output: str, loglevel=logging.INFO):
 
     module.add_named_metadata("llvm.ident", [f"PythonBPF {VERSION}"])
 
+    module_string = finalize_module(str(module))
+
     logger.info(f"IR written to {output}")
     with open(output, "w") as f:
         f.write(f'source_filename = "{filename}"\n')
-        f.write(str(module))
+        f.write(module_string)
         f.write("\n")
 
     return output, structs_sym_tab, maps_sym_tab

pythonbpf/debuginfo/debug_info_generator.py

@@ -81,6 +81,20 @@ def create_array_type(self, base_type: Any, count: int) -> Any:
             },
         )
 
+    def create_array_type_vmlinux(self, type_info: Any, count: int) -> Any:
+        """Create an array type of the given base type with specified count"""
+        base_type, type_sizing = type_info
+        subrange = self.module.add_debug_info("DISubrange", {"count": count})
+        return self.module.add_debug_info(
+            "DICompositeType",
+            {
+                "tag": dc.DW_TAG_array_type,
+                "baseType": base_type,
+                "size": type_sizing,
+                "elements": [subrange],
+            },
+        )
+
     @staticmethod
     def _compute_array_size(base_type: Any, count: int) -> int:
         # Extract size from base_type if possible
@@ -101,6 +115,23 @@ def create_struct_member(self, name: str, base_type: Any, offset: int) -> Any:
             },
         )
 
+    def create_struct_member_vmlinux(
+        self, name: str, base_type_with_size: Any, offset: int
+    ) -> Any:
+        """Create a struct member with the given name, type, and offset"""
+        base_type, type_size = base_type_with_size
+        return self.module.add_debug_info(
+            "DIDerivedType",
+            {
+                "tag": dc.DW_TAG_member,
+                "name": name,
+                "file": self.module._file_metadata,
+                "baseType": base_type,
+                "size": type_size,
+                "offset": offset,
+            },
+        )
+
     def create_struct_type(
         self, members: List[Any], size: int, is_distinct: bool
     ) -> Any:
@@ -116,6 +147,22 @@ def create_struct_type(
             is_distinct=is_distinct,
         )
 
+    def create_struct_type_with_name(
+        self, name: str, members: List[Any], size: int, is_distinct: bool
+    ) -> Any:
+        """Create a struct type with the given members and size"""
+        return self.module.add_debug_info(
+            "DICompositeType",
+            {
+                "name": name,
+                "tag": dc.DW_TAG_structure_type,
+                "file": self.module._file_metadata,
+                "size": size,
+                "elements": members,
+            },
+            is_distinct=is_distinct,
+        )
+
     def create_global_var_debug_info(
         self, name: str, var_type: Any, is_local: bool = False
     ) -> Any:

pythonbpf/expr/__init__.py

@@ -2,6 +2,7 @@
 from .type_normalization import convert_to_bool, get_base_type_and_depth
 from .ir_ops import deref_to_depth
 from .call_registry import CallHandlerRegistry
+from .vmlinux_registry import VmlinuxHandlerRegistry
 
 __all__ = [
     "eval_expr",
@@ -11,4 +12,5 @@
     "deref_to_depth",
     "get_operand_value",
     "CallHandlerRegistry",
+    "VmlinuxHandlerRegistry",
 ]

pythonbpf/expr/expr_pass.py

@@ -12,6 +12,7 @@
     get_base_type_and_depth,
     deref_to_depth,
 )
+from .vmlinux_registry import VmlinuxHandlerRegistry
 
 logger: Logger = logging.getLogger(__name__)
 
@@ -27,8 +28,12 @@ def _handle_name_expr(expr: ast.Name, local_sym_tab: Dict, builder: ir.IRBuilder
         val = builder.load(var)
         return val, local_sym_tab[expr.id].ir_type
     else:
-        logger.info(f"Undefined variable {expr.id}")
-        return None
+        # Check if it's a vmlinux enum/constant
+        vmlinux_result = VmlinuxHandlerRegistry.handle_name(expr.id)
+        if vmlinux_result is not None:
+            return vmlinux_result
+
+        raise SyntaxError(f"Undefined variable {expr.id}")
 
 
 def _handle_constant_expr(module, builder, expr: ast.Constant):
@@ -74,6 +79,13 @@ def _handle_attribute_expr(
                 val = builder.load(gep)
                 field_type = metadata.field_type(attr_name)
                 return val, field_type
+
+        # Try vmlinux handler as fallback
+        vmlinux_result = VmlinuxHandlerRegistry.handle_attribute(
+            expr, local_sym_tab, None, builder
+        )
+        if vmlinux_result is not None:
+            return vmlinux_result
     return None
 
 
@@ -130,7 +142,12 @@ def get_operand_value(
             logger.info(f"var is {var}, base_type is {base_type}, depth is {depth}")
             val = deref_to_depth(func, builder, var, depth)
             return val
-        raise ValueError(f"Undefined variable: {operand.id}")
+        else:
+            # Check if it's a vmlinux enum/constant
+            vmlinux_result = VmlinuxHandlerRegistry.handle_name(operand.id)
+            if vmlinux_result is not None:
+                val, _ = vmlinux_result
+                return val
     elif isinstance(operand, ast.Constant):
         if isinstance(operand.value, int):
             cst = ir.Constant(ir.IntType(64), int(operand.value))
@@ -332,6 +349,7 @@ def _handle_unary_op(
         neg_one = ir.Constant(ir.IntType(64), -1)
         result = builder.mul(operand, neg_one)
         return result, ir.IntType(64)
+    return None
 
 
 # ============================================================================

pythonbpf/expr/vmlinux_registry.py

@@ -0,0 +1,45 @@
+import ast
+
+
+class VmlinuxHandlerRegistry:
+    """Registry for vmlinux handler operations"""
+
+    _handler = None
+
+    @classmethod
+    def set_handler(cls, handler):
+        """Set the vmlinux handler"""
+        cls._handler = handler
+
+    @classmethod
+    def get_handler(cls):
+        """Get the vmlinux handler"""
+        return cls._handler
+
+    @classmethod
+    def handle_name(cls, name):
+        """Try to handle a name as vmlinux enum/constant"""
+        if cls._handler is None:
+            return None
+        return cls._handler.handle_vmlinux_enum(name)
+
+    @classmethod
+    def handle_attribute(cls, expr, local_sym_tab, module, builder):
+        """Try to handle an attribute access as vmlinux struct field"""
+        if cls._handler is None:
+            return None
+
+        if isinstance(expr.value, ast.Name):
+            var_name = expr.value.id
+            field_name = expr.attr
+            return cls._handler.handle_vmlinux_struct_field(
+                var_name, field_name, module, builder, local_sym_tab
+            )
+        return None
+
+    @classmethod
+    def is_vmlinux_struct(cls, name):
+        """Check if a name refers to a vmlinux struct"""
+        if cls._handler is None:
+            return False
+        return cls._handler.is_vmlinux_struct(name)

pythonbpf/functions/functions_pass.py

@@ -310,7 +310,13 @@ def process_stmt(
 
 
 def process_func_body(
-    module, builder, func_node, func, ret_type, map_sym_tab, structs_sym_tab
+    module,
+    builder,
+    func_node,
+    func,
+    ret_type,
+    map_sym_tab,
+    structs_sym_tab,
 ):
     """Process the body of a bpf function"""
     # TODO: A lot.  We just have print -> bpf_trace_printk for now
@@ -383,7 +389,13 @@ def process_bpf_chunk(func_node, module, return_type, map_sym_tab, structs_sym_t
     builder = ir.IRBuilder(block)
 
     process_func_body(
-        module, builder, func_node, func, ret_type, map_sym_tab, structs_sym_tab
+        module,
+        builder,
+        func_node,
+        func,
+        ret_type,
+        map_sym_tab,
+        structs_sym_tab,
     )
     return func
 

pythonbpf/helper/printk_formatter.py

@@ -3,6 +3,7 @@
 
 from llvmlite import ir
 from pythonbpf.expr import eval_expr, get_base_type_and_depth, deref_to_depth
+from pythonbpf.expr.vmlinux_registry import VmlinuxHandlerRegistry
 
 logger = logging.getLogger(__name__)
 
@@ -108,6 +109,16 @@ def _process_name_in_fval(name_node, fmt_parts, exprs, local_sym_tab):
     if local_sym_tab and name_node.id in local_sym_tab:
         _, var_type, tmp = local_sym_tab[name_node.id]
         _populate_fval(var_type, name_node, fmt_parts, exprs)
+    else:
+        # Try to resolve through vmlinux registry if not in local symbol table
+        result = VmlinuxHandlerRegistry.handle_name(name_node.id)
+        if result:
+            val, var_type = result
+            _populate_fval(var_type, name_node, fmt_parts, exprs)
+        else:
+            raise ValueError(
+                f"Variable '{name_node.id}' not found in symbol table or vmlinux"
+            )
 
 
 def _process_attr_in_fval(attr_node, fmt_parts, exprs, local_sym_tab, struct_sym_tab):