Add setuid C example

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

demo/pybpf4.py

@@ -0,0 +1,56 @@
+import time
+
+from pythonbpf import bpf, map, section, bpfglobal, BPF
+from pythonbpf.helpers import pid
+from pythonbpf.maps import HashMap
+from pylibbpf import *
+from ctypes import c_void_p, c_int64, c_uint64, c_int32
+import matplotlib.pyplot as plt
+
+# This program attaches an eBPF tracepoint to sys_enter_clone,
+# counts per-PID clone syscalls, stores them in a hash map,
+# and then plots the distribution as a histogram using matplotlib.
+# It provides a quick view of process creation activity over 10 seconds.
+# Everything is done with Python only code and with the new pylibbpf library.
+# Run `sudo /path/to/python/binary/ pybpf4.py`
+
+@bpf
+@map
+def hist() -> HashMap:
+    return HashMap(key=c_int32, value=c_uint64, max_entries=4096)
+
+@bpf
+@section("tracepoint/syscalls/sys_enter_clone")
+def hello(ctx: c_void_p) -> c_int64:
+    process_id = pid()
+    one = 1
+    prev = hist().lookup(process_id)
+    if prev:
+        previous_value = prev + 1
+        print(f"count: {previous_value} with {process_id}")
+        hist().update(process_id, previous_value)
+        return c_int64(0)
+    else:
+        hist().update(process_id, one)
+    return c_int64(0)
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+b = BPF()
+b.load_and_attach()
+hist = BpfMap(b, hist)
+print("Recording")
+time.sleep(10)
+
+counts = list(hist.values())
+
+plt.hist(counts, bins=20)
+plt.xlabel("Clone calls per PID")
+plt.ylabel("Frequency")
+plt.title("Syscall clone counts")
+plt.show()

examples/c-form/ex7.bpf.c

@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: GPL-2.0
+#include <linux/bpf.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+
+
+struct event {
+    __u32 pid;
+    __u32 uid;
+    __u64 ts;
+};
+
+struct {
+    __uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);
+    __uint(key_size, sizeof(int));
+    __uint(value_size, sizeof(int));
+} events SEC(".maps");
+
+SEC("tp/syscalls/sys_enter_setuid")
+int handle_setuid_entry(struct trace_event_raw_sys_enter *ctx)
+{
+    struct event data = {};
+    
+    // Extract UID from the syscall arguments
+    data.uid = (unsigned int)ctx->args[0];
+    data.ts = bpf_ktime_get_ns();
+    data.pid = bpf_get_current_pid_tgid() >> 32;
+    
+    bpf_perf_event_output(ctx, &events, BPF_F_CURRENT_CPU, &data, sizeof(data));
+    
+    return 0;
+}
+
+char LICENSE[] SEC("license") = "GPL";

examples/execve5.py

@@ -27,6 +27,7 @@ def hello(ctx: c_void_p) -> c_int32:
     dataobj.pid = process_id
     dataobj.ts = ts
     print(f"clone called at {ts} by pid {process_id}")
+    events.output(dataobj)
     return c_int32(0)
 
 

pyproject.toml

@@ -16,7 +16,8 @@ requires-python = ">=3.8"
 
 dependencies = [
   "llvmlite",
-  "astpretty"
+  "astpretty",
+  "pylibbpf"
 ]
 
 [tool.setuptools.packages.find]

pythonbpf/__init__.py

@@ -1,2 +1,2 @@
 from .decorators import bpf, map, section, bpfglobal, struct
-from .codegen import compile_to_ir, compile
+from .codegen import compile_to_ir, compile, BPF

pythonbpf/bpf_helper_handler.py

@@ -3,7 +3,7 @@ from llvmlite import ir
 from .expr_pass import eval_expr
 
 
-def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, local_sym_tab=None):
+def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, local_var_metadata=None):
     """
     Emit LLVM IR for bpf_ktime_get_ns helper function call.
     """
@@ -16,7 +16,7 @@ def bpf_ktime_get_ns_emitter(call, map_ptr, module, builder, func, local_sym_tab
     return result
 
 
-def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, local_sym_tab=None):
+def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
     """
     Emit LLVM IR for bpf_map_lookup_elem helper function call.
     """
@@ -63,7 +63,7 @@ def bpf_map_lookup_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No
     return result
 
 
-def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None):
+def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, local_var_metadata=None):
     if not hasattr(func, "_fmt_counter"):
         func._fmt_counter = 0
 
@@ -172,7 +172,7 @@ def bpf_printk_emitter(call, map_ptr, module, builder, func, local_sym_tab=None)
                 ir.IntType(32), len(fmt_str))], tail=True)
 
 
-def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=None):
+def bpf_map_update_elem_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
     """
     Emit LLVM IR for bpf_map_update_elem helper function call.
     Expected call signature: map.update(key, value, flags=0)
@@ -268,7 +268,7 @@ def bpf_map_update_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No
     return result
 
 
-def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, local_sym_tab=None):
+def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
     """
     Emit LLVM IR for bpf_map_delete_elem helper function call.
     Expected call signature: map.delete(key)
@@ -323,7 +323,7 @@ def bpf_map_delete_elem_emitter(call, map_ptr, module, builder, local_sym_tab=No
     return result
 
 
-def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func, local_sym_tab=None):
+def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func, local_sym_tab=None, local_var_metadata=None):
     """
     Emit LLVM IR for bpf_get_current_pid_tgid helper function call.
     """
@@ -340,6 +340,58 @@ def bpf_get_current_pid_tgid_emitter(call, map_ptr, module, builder, func, local
     return pid
 
 
+def bpf_perf_event_output_handler(call, map_ptr, module, builder, func, local_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
+    if len(call.args) != 1:
+        raise ValueError("Perf event output expects exactly one argument (data), got "
+                         f"{len(call.args)}")
+    data_arg = call.args[0]
+    ctx_ptr = func.args[0]  # First argument to the function is ctx
+
+    if isinstance(data_arg, ast.Name):
+        data_name = data_arg.id
+        if local_sym_tab and data_name in local_sym_tab:
+            data_ptr = local_sym_tab[data_name]
+        else:
+            raise ValueError(
+                f"Data variable {data_name} not found in local symbol table.")
+        # Check is data_name is a struct
+        if local_var_metadata and data_name in local_var_metadata:
+            data_type = local_var_metadata[data_name]
+            if data_type in struct_sym_tab:
+                struct_info = struct_sym_tab[data_type]
+                size_val = ir.Constant(ir.IntType(64), struct_info["size"])
+            else:
+                raise ValueError(
+                    f"Struct type {data_type} for variable {data_name} not found in struct symbol table.")
+        else:
+            raise ValueError(
+                f"Metadata for variable {data_name} not found in local variable metadata.")
+
+        # BPF_F_CURRENT_CPU is -1 in 32 bit
+        flags_val = ir.Constant(ir.IntType(64), 0xFFFFFFFF)
+
+        map_void_ptr = builder.bitcast(map_ptr, ir.PointerType())
+        data_void_ptr = builder.bitcast(data_ptr, ir.PointerType())
+        fn_type = ir.FunctionType(
+            ir.IntType(64),
+            [ir.PointerType(ir.IntType(8)), ir.PointerType(), ir.IntType(64),
+             ir.PointerType(), ir.IntType(64)],
+            var_arg=False
+        )
+        fn_ptr_type = ir.PointerType(fn_type)
+
+        # helper id
+        fn_addr = ir.Constant(ir.IntType(64), 25)
+        fn_ptr = builder.inttoptr(fn_addr, fn_ptr_type)
+
+        result = builder.call(
+            fn_ptr, [ctx_ptr, map_void_ptr, flags_val, data_void_ptr, size_val], tail=False)
+        return result
+    else:
+        raise NotImplementedError(
+            "Only simple object names are supported as data in perf event output.")
+
+
 helper_func_list = {
     "lookup": bpf_map_lookup_elem_emitter,
     "print": bpf_printk_emitter,
@@ -347,10 +399,11 @@ helper_func_list = {
     "update": bpf_map_update_elem_emitter,
     "delete": bpf_map_delete_elem_emitter,
     "pid": bpf_get_current_pid_tgid_emitter,
+    "output": bpf_perf_event_output_handler,
 }
 
 
-def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_tab=None):
+def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_tab=None, struct_sym_tab=None, local_var_metadata=None):
     if isinstance(call.func, ast.Name):
         func_name = call.func.id
         if func_name in helper_func_list:
@@ -367,14 +420,29 @@ def handle_helper_call(call, module, builder, func, local_sym_tab=None, map_sym_
             if map_sym_tab and map_name in map_sym_tab:
                 map_ptr = map_sym_tab[map_name]
                 if method_name in helper_func_list:
+                    print(local_var_metadata)
                     return helper_func_list[method_name](
-                        call, map_ptr, module, builder, local_sym_tab)
+                        call, map_ptr, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata)
                 else:
                     raise NotImplementedError(
                         f"Map method {method_name} is not implemented as a helper function.")
             else:
                 raise ValueError(
                     f"Map variable {map_name} not found in symbol tables.")
+        elif isinstance(call.func.value, ast.Name):
+            obj_name = call.func.value.id
+            method_name = call.func.attr
+            if map_sym_tab and obj_name in map_sym_tab:
+                map_ptr = map_sym_tab[obj_name]
+                if method_name in helper_func_list:
+                    return helper_func_list[method_name](
+                        call, map_ptr, module, builder, func, local_sym_tab, struct_sym_tab, local_var_metadata)
+                else:
+                    raise NotImplementedError(
+                        f"Map method {method_name} is not implemented as a helper function.")
+            else:
+                raise ValueError(
+                    f"Map variable {obj_name} not found in symbol tables.")
         else:
             raise NotImplementedError(
                 "Attribute not supported for map method calls.")

pythonbpf/codegen.py

@@ -9,6 +9,7 @@ import os
 import subprocess
 import inspect
 from pathlib import Path
+from pylibbpf import BpfProgram
 
 
 def find_bpf_chunks(tree):
@@ -116,3 +117,17 @@ def compile():
     ], check=True)
 
     print(f"Object written to {o_file}, {ll_file} can be removed")
+
+def BPF() -> BpfProgram:
+    caller_frame = inspect.stack()[1]
+    caller_file = Path(caller_frame.filename).resolve()
+    ll_file = Path("/tmp") / caller_file.with_suffix(".ll").name
+    o_file = Path("/tmp") / caller_file.with_suffix(".o").name
+    compile_to_ir(str(caller_file), str(ll_file))
+
+    subprocess.run([
+        "llc", "-march=bpf", "-filetype=obj", "-O2",
+        str(ll_file), "-o", str(o_file)
+    ], check=True)
+    
+    return BpfProgram(str(o_file))

pythonbpf/expr_pass.py

@@ -2,7 +2,7 @@ import ast
 from llvmlite import ir
 
 
-def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab):
+def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab, structs_sym_tab=None, local_var_metadata=None):
     print(f"Evaluating expression: {expr}")
     if isinstance(expr, ast.Name):
         if expr.id in local_sym_tab:
@@ -50,22 +50,31 @@ def eval_expr(func, module, builder, expr, local_sym_tab, map_sym_tab):
             # check for helpers
             if expr.func.id in helper_func_list:
                 return handle_helper_call(
-                    expr, module, builder, func, local_sym_tab, map_sym_tab)
+                    expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
         elif isinstance(expr.func, ast.Attribute):
+            print(f"Handling method call: {ast.dump(expr.func)}")
             if isinstance(expr.func.value, ast.Call) and isinstance(expr.func.value.func, ast.Name):
                 method_name = expr.func.attr
                 if method_name in helper_func_list:
                     return handle_helper_call(
-                        expr, module, builder, func, local_sym_tab, map_sym_tab)
+                        expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
+            elif isinstance(expr.func.value, ast.Name):
+                obj_name = expr.func.value.id
+                method_name = expr.func.attr
+                if obj_name in map_sym_tab:
+                    if method_name in helper_func_list:
+                        return handle_helper_call(
+                            expr, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
     print("Unsupported expression evaluation")
     return None
 
 
-def handle_expr(func, module, builder, expr, local_sym_tab, map_sym_tab):
+def handle_expr(func, module, builder, expr, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata):
     """Handle expression statements in the function body."""
     print(f"Handling expression: {ast.dump(expr)}")
     call = expr.value
     if isinstance(call, ast.Call):
-        eval_expr(func, module, builder, call, local_sym_tab, map_sym_tab)
+        eval_expr(func, module, builder, call, local_sym_tab,
+                  map_sym_tab, structs_sym_tab, local_var_metadata)
     else:
         print("Unsupported expression type")

pythonbpf/functions_pass.py

@@ -57,7 +57,7 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc
                                  ir.Constant(ir.IntType(32), field_idx)],
                     inbounds=True)
                 val = eval_expr(func, module, builder, rval,
-                                local_sym_tab, map_sym_tab)
+                                local_sym_tab, map_sym_tab, structs_sym_tab)
                 if val is None:
                     print("Failed to evaluate struct field assignment")
                     return
@@ -100,14 +100,14 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc
                 # var = builder.alloca(ir.IntType(64), name=var_name)
                 # var.align = 8
                 val = handle_helper_call(
-                    rval, module, builder, None, local_sym_tab, map_sym_tab)
+                    rval, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
                 builder.store(val, local_sym_tab[var_name])
                 # local_sym_tab[var_name] = var
                 print(f"Assigned constant {rval.func.id} to {var_name}")
             elif call_type == "deref" and len(rval.args) == 1:
                 print(f"Handling deref assignment {ast.dump(rval)}")
                 val = eval_expr(func, module, builder, rval,
-                                local_sym_tab, map_sym_tab)
+                                local_sym_tab, map_sym_tab, structs_sym_tab)
                 if val is None:
                     print("Failed to evaluate deref argument")
                     return
@@ -139,7 +139,7 @@ def handle_assign(func, module, builder, stmt, map_sym_tab, local_sym_tab, struc
                     map_ptr = map_sym_tab[map_name]
                     if method_name in helper_func_list:
                         val = handle_helper_call(
-                            rval, module, builder, func, local_sym_tab, map_sym_tab)
+                            rval, module, builder, func, local_sym_tab, map_sym_tab, structs_sym_tab, local_var_metadata)
                         # var = builder.alloca(ir.IntType(64), name=var_name)
                         # var.align = 8
                         builder.store(val, local_sym_tab[var_name])
@@ -222,7 +222,7 @@ def handle_cond(func, module, builder, cond, local_sym_tab, map_sym_tab):
         return None
 
 
-def handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab):
+def handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab, structs_sym_tab=None):
     """Handle if statements in the function body."""
     print("Handling if statement")
     start = builder.block.parent
@@ -243,7 +243,7 @@ def handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab):
     builder.position_at_end(then_block)
     for s in stmt.body:
         process_stmt(func, module, builder, s,
-                     local_sym_tab, map_sym_tab, False)
+                     local_sym_tab, map_sym_tab, structs_sym_tab, False)
     if not builder.block.is_terminated:
         builder.branch(merge_block)
 
@@ -251,7 +251,7 @@ def handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab):
         builder.position_at_end(else_block)
         for s in stmt.orelse:
             process_stmt(func, module, builder, s,
-                         local_sym_tab, map_sym_tab, False)
+                         local_sym_tab, map_sym_tab, structs_sym_tab, False)
         if not builder.block.is_terminated:
             builder.branch(merge_block)
 
@@ -261,14 +261,16 @@ def handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab):
 def process_stmt(func, module, builder, stmt, local_sym_tab, map_sym_tab, structs_sym_tab, did_return, ret_type=ir.IntType(64)):
     print(f"Processing statement: {ast.dump(stmt)}")
     if isinstance(stmt, ast.Expr):
-        handle_expr(func, module, builder, stmt, local_sym_tab, map_sym_tab)
+        handle_expr(func, module, builder, stmt, local_sym_tab,
+                    map_sym_tab, structs_sym_tab, local_var_metadata)
     elif isinstance(stmt, ast.Assign):
         handle_assign(func, module, builder, stmt, map_sym_tab,
                       local_sym_tab, structs_sym_tab)
     elif isinstance(stmt, ast.AugAssign):
         raise SyntaxError("Augmented assignment not supported")
     elif isinstance(stmt, ast.If):
-        handle_if(func, module, builder, stmt, map_sym_tab, local_sym_tab)
+        handle_if(func, module, builder, stmt, map_sym_tab,
+                  local_sym_tab, structs_sym_tab)
     elif isinstance(stmt, ast.Return):
         if stmt.value is None:
             builder.ret(ir.Constant(ir.IntType(32), 0))

pythonbpf/maps.py

@@ -30,3 +30,6 @@ class PerfEventArray:
         self.key_type = key_size
         self.value_type = value_size
         self.entries = {}
+
+    def output(self, data):
+        pass  # Placeholder for output method

pythonbpf/structs_pass.py

@@ -31,9 +31,27 @@ def process_bpf_struct(cls_node, module):
             field_names.append(item.target.id)
             field_types.append(ctypes_to_ir(item.annotation.id))
 
+    curr_offset = 0
+    for ftype in field_types:
+        if isinstance(ftype, ir.IntType):
+            fsize = ftype.width // 8
+            alignment = fsize
+        elif isinstance(ftype, ir.PointerType):
+            fsize = 8
+            alignment = 8
+        else:
+            print(f"Unsupported field type in struct {struct_name}")
+            return
+        padding = (alignment - (curr_offset % alignment)) % alignment
+        curr_offset += padding
+        curr_offset += fsize
+    final_padding = (8 - (curr_offset % 8)) % 8
+    total_size = curr_offset + final_padding
+
     struct_type = ir.LiteralStructType(field_types)
     structs_sym_tab[struct_name] = {
         "type": struct_type,
-        "fields": {name: idx for idx, name in enumerate(field_names)}
+        "fields": {name: idx for idx, name in enumerate(field_names)},
+        "size": total_size
     }
     print(f"Created struct {struct_name} with fields {field_names}")