format chore

BCC-Examples/container-monitor/file_io.bpf.py

@@ -0,0 +1,92 @@
+import logging
+
+from pythonbpf import bpf, map, section, bpfglobal, struct, compile
+from pythonbpf.maps import HashMap
+from pythonbpf.helper import get_current_cgroup_id
+from ctypes import c_int32, c_uint64
+from vmlinux import struct_pt_regs
+
+
+@bpf
+@struct
+class read_stats:
+    bytes: c_uint64
+    ops: c_uint64
+
+
+@bpf
+@struct
+class write_stats:
+    bytes: c_uint64
+    ops: c_uint64
+
+
+@bpf
+@map
+def read_map() -> HashMap:
+    return HashMap(key=c_uint64, value=read_stats, max_entries=1024)
+
+
+@bpf
+@map
+def write_map() -> HashMap:
+    return HashMap(key=c_uint64, value=write_stats, max_entries=1024)
+
+
+#
+# READ PROBE
+#
+@bpf
+@section("kprobe/vfs_read")
+def trace_read(ctx: struct_pt_regs) -> c_int32:
+    cg = get_current_cgroup_id()
+    count = c_uint64(ctx.dx)
+    ptr = read_map.lookup(cg)
+
+    if ptr:
+        s = read_stats()
+        s.bytes = ptr.bytes + count
+        s.ops = ptr.ops + 1
+        read_map.update(cg, ptr)
+    else:
+        print("read init")
+        s = read_stats()
+        s.bytes = count
+        s.ops = c_uint64(1)
+        read_map.update(cg, s)
+
+    return c_int32(0)
+
+
+#
+# WRITE PROBE
+#
+@bpf
+@section("kprobe/vfs_write")
+def trace_write(ctx1: struct_pt_regs) -> c_int32:
+    cg = get_current_cgroup_id()
+    count = c_uint64(ctx1.dx)
+    ptr = write_map.lookup(cg)
+
+    if ptr:
+        s = write_stats()
+        s.bytes = ptr.bytes + count
+        s.ops = ptr.ops + 1
+        write_map.update(cg, s)
+    else:
+        print("write init")
+        s = write_stats()
+        s.bytes = count
+        s.ops = c_uint64(1)
+        write_map.update(cg, s)
+
+    return c_int32(0)
+
+
+@bpf
+@bpfglobal
+def LICENSE() -> str:
+    return "GPL"
+
+
+compile(loglevel=logging.INFO)

pythonbpf/helper/__init__.py

@@ -16,6 +16,7 @@ from .helpers import (
     smp_processor_id,
     uid,
     skb_store_bytes,
+    get_current_cgroup_id,
     get_stack,
     XDP_DROP,
     XDP_PASS,
@@ -79,6 +80,7 @@ __all__ = [
     "handle_helper_call",
     "emit_probe_read_kernel_str_call",
     "emit_probe_read_kernel_call",
+    "get_current_cgroup_id",
     "ktime",
     "pid",
     "deref",

pythonbpf/helper/bpf_helper_handler.py

@@ -30,6 +30,7 @@ class BPFHelperID(Enum):
     BPF_SKB_STORE_BYTES = 9
     BPF_GET_CURRENT_PID_TGID = 14
     BPF_GET_CURRENT_UID_GID = 15
+    BPF_GET_CURRENT_CGROUP_ID = 80
     BPF_GET_CURRENT_COMM = 16
     BPF_PERF_EVENT_OUTPUT = 25
     BPF_GET_STACK = 67
@@ -68,6 +69,33 @@ def bpf_ktime_get_ns_emitter(
     return result, ir.IntType(64)
 
 
+@HelperHandlerRegistry.register(
+    "get_current_cgroup_id",
+    param_types=[],
+    return_type=ir.IntType(64),
+)
+def bpf_get_current_cgroup_id(
+    call,
+    map_ptr,
+    module,
+    builder,
+    func,
+    local_sym_tab=None,
+    struct_sym_tab=None,
+    map_sym_tab=None,
+):
+    """
+    Emit LLVM IR for bpf_get_current_cgroup_id helper function call.
+    """
+    # func is an arg to just have a uniform signature with other emitters
+    helper_id = ir.Constant(ir.IntType(64), BPFHelperID.BPF_GET_CURRENT_CGROUP_ID.value)
+    fn_type = ir.FunctionType(ir.IntType(64), [], var_arg=False)
+    fn_ptr_type = ir.PointerType(fn_type)
+    fn_ptr = builder.inttoptr(helper_id, fn_ptr_type)
+    result = builder.call(fn_ptr, [], tail=False)
+    return result, ir.IntType(64)
+
+
 @HelperHandlerRegistry.register(
     "lookup",
     param_types=[ir.PointerType(ir.IntType(64))],

pythonbpf/helper/helpers.py

@@ -57,6 +57,11 @@ def get_stack(buf, flags=0):
     return ctypes.c_int64(0)
 
 
+def get_current_cgroup_id():
+    """Get the current cgroup ID"""
+    return ctypes.c_int64(0)
+
+
 XDP_ABORTED = ctypes.c_int64(0)
 XDP_DROP = ctypes.c_int64(1)
 XDP_PASS = ctypes.c_int64(2)

tests/c-form/xdp_test.bpf.c

@@ -1,19 +1,18 @@
-#include "vmlinux.h"
-#include <bpf/bpf_helpers.h>
 #include <linux/bpf.h>
 #include <linux/if_ether.h>
 #include <linux/ip.h>
+#include <bpf/bpf_helpers.h>
 
 struct fake_iphdr {
-  unsigned short useless;
-  unsigned short tot_len;
-  unsigned short id;
-  unsigned short frag_off;
-  unsigned char ttl;
-  unsigned char protocol;
-  unsigned short check;
-  unsigned int saddr;
-  unsigned int daddr;
+    unsigned short useless;
+    unsigned short tot_len;
+    unsigned short id;
+    unsigned short frag_off;
+    unsigned char  ttl;
+    unsigned char  protocol;
+    unsigned short check;
+    unsigned int saddr;
+    unsigned int daddr;
 };
 
 SEC("xdp")
@@ -21,16 +20,14 @@ int xdp_prog(struct xdp_md *ctx) {
   unsigned long data = ctx->data;
   unsigned long data_end = ctx->data_end;
 
-  if (data + sizeof(struct ethhdr) + sizeof(struct fake_iphdr) <= data_end) {
-    struct fake_iphdr *iph = (void *)data + sizeof(struct ethhdr);
-
-    bpf_printk("%d", iph->saddr);
-
-    return XDP_PASS;
-  } else {
+  if (data + sizeof(struct ethhdr) + sizeof(struct fake_iphdr) > data_end) {
     return XDP_ABORTED;
   }
-  struct task_struct * a = btf_bpf_get_current_task_btf();
+  struct fake_iphdr *iph = (void *)data + sizeof(struct ethhdr);
+
+   bpf_printk("%d", iph->saddr);
+
+   return XDP_PASS;
 }
 
 char _license[] SEC("license") = "GPL";