Merge branch 'master' into vmlinux-handler

Signed-off-by: varun-r-mallya <varunrmallya@gmail.com>

BCC-Examples/README.md

@@ -0,0 +1,20 @@
+## BCC examples ported to PythonBPF
+
+This folder contains examples of BCC tutorial examples that have been ported to use **PythonBPF**.
+
+## Requirements
+- install `pythonbpf` and `pylibbpf` using pip.
+- You will also need `matplotlib` for vfsreadlat.py example.
+- You will also need `rich` for vfsreadlat_rich.py example.
+- You will also need `plotly` and `dash` for vfsreadlat_plotly.py example.
+
+## Usage
+- You'll need root privileges to run these examples. If you are using a virtualenv, use the following command to run the scripts:
+  ```bash
+  sudo <path_to_virtualenv>/bin/python3 <script_name>.py
+  ```
+- For vfsreadlat_plotly.py, run the following command to start the Dash server:
+  ```bash
+  sudo <path_to_virtualenv>/bin/python3 vfsreadlat_plotly/bpf_program.py
+  ```
+  Then open your web browser and navigate to the given URL.

pythonbpf/codegen.py

@@ -37,6 +37,24 @@ def finalize_module(original_str):
     return re.sub(pattern, replacement, original_str)
 
 
+def bpf_passthrough_gen(module):
+    i32_ty = ir.IntType(32)
+    ptr_ty = ir.PointerType(ir.IntType(8))
+    fnty = ir.FunctionType(ptr_ty, [i32_ty, ptr_ty])
+
+    # Declare the intrinsic
+    passthrough = ir.Function(module, fnty, "llvm.bpf.passthrough.p0.p0")
+
+    # Set function attributes
+    # TODO: the ones commented are supposed to be there but cannot be added due to llvmlite limitations at the moment
+    # passthrough.attributes.add("nofree")
+    # passthrough.attributes.add("nosync")
+    passthrough.attributes.add("nounwind")
+    # passthrough.attributes.add("memory(none)")
+
+    return passthrough
+
+
 def find_bpf_chunks(tree):
     """Find all functions decorated with @bpf in the AST."""
     bpf_functions = []
@@ -57,6 +75,8 @@ def processor(source_code, filename, module):
     for func_node in bpf_chunks:
         logger.info(f"Found BPF function/struct: {func_node.name}")
 
+    bpf_passthrough_gen(module)
+
     vmlinux_symtab = vmlinux_proc(tree, module)
     if vmlinux_symtab:
         handler = VmlinuxHandler.initialize(vmlinux_symtab)

pythonbpf/expr/vmlinux_registry.py

@@ -1,5 +1,7 @@
 import ast
 
+from pythonbpf.vmlinux_parser.vmlinux_exports_handler import VmlinuxHandler
+
 
 class VmlinuxHandlerRegistry:
     """Registry for vmlinux handler operations"""
@@ -7,7 +9,7 @@ class VmlinuxHandlerRegistry:
     _handler = None
 
     @classmethod
-    def set_handler(cls, handler):
+    def set_handler(cls, handler: VmlinuxHandler):
         """Set the vmlinux handler"""
         cls._handler = handler
 
@@ -43,3 +45,10 @@ class VmlinuxHandlerRegistry:
         if cls._handler is None:
             return False
         return cls._handler.is_vmlinux_struct(name)
+
+    @classmethod
+    def get_struct_type(cls, name):
+        """Try to handle a struct name as vmlinux struct"""
+        if cls._handler is None:
+            return None
+        return cls._handler.get_vmlinux_struct_type(name)

pythonbpf/functions/functions_pass.py

@@ -7,7 +7,12 @@ from pythonbpf.helper import (
     reset_scratch_pool,
 )
 from pythonbpf.type_deducer import ctypes_to_ir
-from pythonbpf.expr import eval_expr, handle_expr, convert_to_bool
+from pythonbpf.expr import (
+    eval_expr,
+    handle_expr,
+    convert_to_bool,
+    VmlinuxHandlerRegistry,
+)
 from pythonbpf.assign_pass import (
     handle_variable_assignment,
     handle_struct_field_assignment,
@@ -324,6 +329,35 @@ def process_func_body(
 
     local_sym_tab = {}
 
+    # Add the context parameter (first function argument) to the local symbol table
+    if func_node.args.args and len(func_node.args.args) > 0:
+        context_arg = func_node.args.args[0]
+        context_name = context_arg.arg
+
+        if hasattr(context_arg, "annotation") and context_arg.annotation:
+            if isinstance(context_arg.annotation, ast.Name):
+                context_type_name = context_arg.annotation.id
+            elif isinstance(context_arg.annotation, ast.Attribute):
+                context_type_name = context_arg.annotation.attr
+            else:
+                raise TypeError(
+                    f"Unsupported annotation type: {ast.dump(context_arg.annotation)}"
+                )
+            if VmlinuxHandlerRegistry.is_vmlinux_struct(context_type_name):
+                resolved_type = VmlinuxHandlerRegistry.get_struct_type(
+                    context_type_name
+                )
+                context_type = {"type": ir.PointerType(resolved_type), "ptr": True}
+            else:
+                try:
+                    resolved_type = ctypes_to_ir(context_type_name)
+                    context_type = {"type": ir.PointerType(resolved_type), "ptr": True}
+                except Exception:
+                    raise TypeError(f"Type '{context_type_name}' not declared")
+
+            local_sym_tab[context_name] = context_type
+            logger.info(f"Added argument '{context_name}' to local symbol table")
+
     # pre-allocate dynamic variables
     local_sym_tab = allocate_mem(
         module,

pythonbpf/helper/bpf_helper_handler.py

@@ -78,9 +78,9 @@ def bpf_map_lookup_elem_emitter(
     map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
 
     # TODO: I have changed the return type to i64*, as we are
-    # allocating space for that type in allocate_mem. This is
+    #   allocating space for that type in allocate_mem. This is
-    # temporary, and we will honour other widths later. But this
+    #   temporary, and we will honour other widths later. But this
-    # allows us to have cool binary ops on the returned value.
+    #   allows us to have cool binary ops on the returned value.
     fn_type = ir.FunctionType(
         ir.PointerType(ir.IntType(64)),  # Return type: void*
         [ir.PointerType(), ir.PointerType()],  # Args: (void*, void*)

pythonbpf/vmlinux_parser/import_detector.py

@@ -86,19 +86,19 @@ def vmlinux_proc(tree: ast.AST, module):
 
     if not import_statements:
         logger.info("No vmlinux imports found")
-        return
+        return None
 
     # Import vmlinux module directly
     try:
         vmlinux_mod = importlib.import_module("vmlinux")
     except ImportError:
         logger.warning("Could not import vmlinux module")
-        return
+        return None
 
     source_file = inspect.getsourcefile(vmlinux_mod)
     if source_file is None:
         logger.warning("Cannot find source for vmlinux module")
-        return
+        return None
 
     with open(source_file, "r") as f:
         mod_ast = ast.parse(f.read(), filename=source_file)

pythonbpf/vmlinux_parser/vmlinux_exports_handler.py

@@ -39,6 +39,16 @@ class VmlinuxHandler:
             and self.vmlinux_symtab[name]["value_type"] == AssignmentType.CONSTANT
         )
 
+    def get_vmlinux_struct_type(self, name):
+        """Check if name is a vmlinux struct type"""
+        if (
+            name in self.vmlinux_symtab
+            and self.vmlinux_symtab[name]["value_type"] == AssignmentType.STRUCT
+        ):
+            return self.vmlinux_symtab[name]["python_type"]
+        else:
+            raise ValueError(f"{name} is not a vmlinux struct type")
+
     def is_vmlinux_struct(self, name):
         """Check if name is a vmlinux struct"""
         return (

tests/c-form/ex5.bpf.c

@@ -1,25 +0,0 @@
-#define __TARGET_ARCH_arm64
-
-#include "vmlinux.h"
-#include <bpf/bpf_helpers.h>
-#include <bpf/bpf_tracing.h>
-#include <bpf/bpf_core_read.h>
-
-// Map: key = struct request*, value = u64 timestamp
-struct {
-    __uint(type, BPF_MAP_TYPE_HASH);
-    __type(key, struct request *);
-    __type(value, u64);
-    __uint(max_entries, 1024);
-} start SEC(".maps");
-
-// Attach to kprobe for blk_start_request
-SEC("kprobe/blk_start_request")
-int BPF_KPROBE(trace_start, struct request *req)
-{
-    u64 ts = bpf_ktime_get_ns();
-    bpf_map_update_elem(&start, &req, &ts, BPF_ANY);
-    return 0;
-}
-
-char LICENSE[] SEC("license") = "GPL";

tests/c-form/struct_field_tests.bpf.c

@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include "vmlinux.h"
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+
+/*
+Information gained from reversing this (multiple kernel versions):
+There is no point of
+```llvm
+ tail call void @llvm.dbg.value(metadata ptr %0, metadata !60, metadata !DIExpression()), !dbg !70
+ ```
+ and the first argument of passthrough is fucking useless. It just needs to be a distinct integer:
+ ```llvm
+   %9 = tail call ptr @llvm.bpf.passthrough.p0.p0(i32 3, ptr %8)
+ ```
+*/
+
+SEC("tp/syscalls/sys_enter_execve")
+int handle_setuid_entry(struct trace_event_raw_sys_enter *ctx) {
+  // Access each argument separately with clear variable assignments
+  unsigned long arg0 = ctx->args[0];
+  bpf_printk("args[0]: %u", arg0);
+
+  unsigned long arg1 = ctx->args[1];
+  bpf_printk("args[1]: %u", arg1);
+
+  // Remove the duplicate access to args[1]
+
+  unsigned long arg2 = ctx->args[2];
+  bpf_printk("args[3]: %u", arg2);
+  bpf_printk("args[4]: %u", ctx->args[2]);
+
+  return 0;
+}
+
+char LICENSE[] SEC("license") = "GPL";

tests/c-form/xdp_modify.bpf.c

@@ -0,0 +1,21 @@
+// xdp_rewrite.c
+#include <linux/bpf.h>
+#include <bpf/bpf_helpers.h>
+#include <linux/if_ether.h>
+
+SEC("xdp")
+int xdp_rewrite_mac(struct xdp_md *ctx)
+{
+    void *data_end = (void *)(long)ctx->data_end;
+    void *data     = (void *)(long)ctx->data;
+
+    struct ethhdr *eth = data;
+    if ((void*)(eth + 1) > data_end)
+        return XDP_PASS;
+    __u8 new_src[ETH_ALEN] = {0x02,0x00,0x00,0x00,0x00,0x02};
+    for (int i = 0; i < ETH_ALEN; i++) eth->h_source[i] = new_src[i];
+
+    return XDP_PASS;
+}
+
+char _license[] SEC("license") = "GPL";

tests/failing_tests/vmlinux/struct_field_access.py

@@ -0,0 +1,31 @@
+import logging
+
+from pythonbpf import bpf, section, bpfglobal, compile_to_ir
+from pythonbpf import compile  # noqa: F401
+from vmlinux import TASK_COMM_LEN  # noqa: F401
+from vmlinux import struct_trace_event_raw_sys_enter  # noqa: F401
+from ctypes import c_int64, c_void_p  # noqa: F401
+
+
+# from vmlinux import struct_uinput_device
+# from vmlinux import struct_blk_integrity_iter
+
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_execve")
+def hello_world(ctx: struct_trace_event_raw_sys_enter) -> c_int64:
+    a = 2 + TASK_COMM_LEN + TASK_COMM_LEN
+    b = ctx.id
+    print(f"Hello, World{TASK_COMM_LEN} and {a}")
+    print(f"This is context field {b}")
+    return c_int64(TASK_COMM_LEN + 2)
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+compile_to_ir("struct_field_access.py", "struct_field_access.ll", loglevel=logging.INFO)
+# compile()