Add regex support to attach_k[ret]probe

Add new event_re parameter to kprobe functions. This searches through
the list of functions/symbols in <tracefs>/available_filter_functions
for all matching functions. Every one is attached.

Add a test case for attach_k[ret]probe.

Fixes: #141
Signed-off-by: Brenden Blanco <bblanco@plumgrid.com>

diff --git a/src/python/bpf/__init__.py b/src/python/bpf/__init__.py
index 10ee07d..b53c721 100644 (file)
--- a/src/python/bpf/__init__.py
+++ b/src/python/bpf/__init__.py
@@ -18,6 +18,7 @@ import ctypes as ct
 import fcntl
 import json
 import os
+from subprocess import Popen, PIPE
 import sys
 basestring = (unicode if sys.version_info[0] < 3 else str)
 
@@ -380,7 +381,23 @@ class BPF(object):
                     % (dev, errstr))
         fn.sock = sock
 
-    def attach_kprobe(self, event="", fn_name="", pid=0, cpu=-1, group_fd=-1):
+    @staticmethod
+    def _get_kprobe_functions(event_re):
+        p = Popen(["awk", "$1 ~ /%s/ { print $1 }" % event_re,
+            "%s/available_filter_functions" % TRACEFS], stdout=PIPE)
+        lines = p.communicate()[0].decode().split()
+        return [line.rstrip() for line in lines if line != "\n"]
+
+    def attach_kprobe(self, event="", fn_name="", event_re="",
+            pid=0, cpu=-1, group_fd=-1):
+
+        # allow the caller to glob multiple functions together
+        if event_re:
+            for line in BPF._get_kprobe_functions(event_re):
+                self.attach_kprobe(event=line, fn_name=fn_name, pid=pid,
+                        cpu=cpu, group_fd=group_fd)
+            return
+
         fn = self.load_func(fn_name, BPF.KPROBE)
         ev_name = "p_" + event.replace("+", "_")
         desc = "p:kprobes/%s %s" % (ev_name, event)
@@ -403,7 +420,16 @@ class BPF(object):
             raise Exception("Failed to detach BPF from kprobe")
         del open_kprobes[ev_name]
 
-    def attach_kretprobe(self, event="", fn_name="", pid=-1, cpu=0, group_fd=-1):
+    def attach_kretprobe(self, event="", fn_name="", event_re="",
+            pid=-1, cpu=0, group_fd=-1):
+
+        # allow the caller to glob multiple functions together
+        if event_re:
+            for line in BPF._get_kprobe_functions(event_re):
+                self.attach_kretprobe(event=line, fn_name=fn_name, pid=pid,
+                        cpu=cpu, group_fd=group_fd)
+            return
+
         fn = self.load_func(fn_name, BPF.KPROBE)
         ev_name = "r_" + event.replace("+", "_")
         desc = "r:kprobes/%s %s" % (ev_name, event)
@@ -452,6 +478,8 @@ class BPF(object):
         """
         line = BPF.trace_readline(nonblocking)
         if line:
+            # don't print messages related to lost events
+            if line.startswith("CPU:"): return
             task = line[:16].lstrip()
             line = line[17:]
             ts_end = line.find(":")
@@ -493,6 +521,7 @@ class BPF(object):
         while True:
             if fmt:
                 fields = BPF.trace_readline_fields(nonblocking=False)
+                if not fields: continue
                 line = fmt.format(*fields)
             else:
                 line = BPF.trace_readline(nonblocking=False)

diff --git a/tests/cc/CMakeLists.txt b/tests/cc/CMakeLists.txt
index 2c42e71..5eb82bd 100644 (file)
--- a/tests/cc/CMakeLists.txt
+++ b/tests/cc/CMakeLists.txt
@@ -30,6 +30,8 @@ add_test(NAME py_test_trace2 WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_trace2 sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_trace2.py)
 add_test(NAME py_test_trace3_c WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_trace3_c sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_trace3.py test_trace3.c)
+add_test(NAME py_test_trace4 WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+  COMMAND ${TEST_WRAPPER} py_trace4 sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_trace4.py)
 add_test(NAME py_test_brb WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_brb_c sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_brb.py test_brb.c)
 add_test(NAME py_test_brb2 WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}

diff --git a/tests/cc/test_trace4.py b/tests/cc/test_trace4.py
new file mode 100755 (executable)
index 0000000..38f5441
--- /dev/null
+++ b/tests/cc/test_trace4.py
@@ -0,0 +1,40 @@
+#!/usr/bin/env python
+# Copyright (c) PLUMgrid, Inc.
+# Licensed under the Apache License, Version 2.0 (the "License")
+
+from bpf import BPF
+import os
+from socket import socket, AF_INET, SOCK_DGRAM
+import sys
+from unittest import main, TestCase
+
+class TestKprobeRgx(TestCase):
+    def setUp(self):
+        self.b = BPF(text="""
+        typedef struct { int idx; } Key;
+        typedef struct { u64 val; } Val;
+        BPF_TABLE("array", Key, Val, stats, 3);
+        int hello(void *ctx) {
+          stats.lookup_or_init(&(Key){1}, &(Val){0})->val++;
+          return 0;
+        }
+        int goodbye(void *ctx) {
+          stats.lookup_or_init(&(Key){2}, &(Val){0})->val++;
+          return 0;
+        }
+        """)
+        self.b.attach_kprobe(event_re="^SyS_send.*", fn_name="hello",
+                pid=0, cpu=-1)
+        self.b.attach_kretprobe(event_re="^SyS_send.*", fn_name="goodbye",
+                pid=1, cpu=-1)
+
+    def test_send1(self):
+        udp = socket(AF_INET, SOCK_DGRAM)
+        udp.sendto(b"a" * 10, ("127.0.0.1", 5000))
+        udp.close()
+        k1 = self.b["stats"].Key(1)
+        k2 = self.b["stats"].Key(2)
+        self.assertEqual(self.b["stats"][k1].val, self.b["stats"][k2].val)
+
+if __name__ == "__main__":
+    main()