Merge pull request #162 from iovisor/bblanco_dev

Add regex support to attach_k[ret]probe

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)

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}

tests/cc/test_trace4.py

+#!/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()