Merge pull request #49 from iovisor/bblanco_dev

Add kprobe auto-cleanup and header fix

examples/hello_world.py

@@ -19,4 +19,7 @@ int hello(void *ctx) {
 b = BPF(text=prog)
 fn = b.load_func("hello", BPF.KPROBE)
 BPF.attach_kprobe(fn, "sys_clone")
-call(["cat", "/sys/kernel/debug/tracing/trace_pipe"])
+try:
+    call(["cat", "/sys/kernel/debug/tracing/trace_pipe"])
+except KeyboardInterrupt:
+    pass

src/cc/export/helpers.h

@@ -17,7 +17,7 @@
 #define __BPF_HELPERS_H
 
 #include <uapi/linux/bpf.h>
-#include <linux/if_packet.h>
+#include <uapi/linux/if_packet.h>
 #include <linux/version.h>
 
 /* helper macro to place programs, maps, license in

src/cc/libbpf.c

@@ -229,3 +229,25 @@ cleanup:
   return rc;
 }
 
+int bpf_detach_kprobe(const char *event_desc) {
+  int rc = -1, kfd = -1;
+
+  kfd = open("/sys/kernel/debug/tracing/kprobe_events", O_WRONLY | O_APPEND, 0);
+  if (kfd < 0) {
+    perror("open(kprobe_events)");
+    goto cleanup;
+  }
+
+  if (write(kfd, event_desc, strlen(event_desc)) < 0) {
+    perror("write(kprobe_events)");
+    goto cleanup;
+  }
+  rc = 0;
+
+cleanup:
+  if (kfd >= 0)
+    close(kfd);
+
+  return rc;
+}
+

src/libbpf.h

@@ -41,6 +41,7 @@ int bpf_attach_socket(int sockfd, int progfd);
 int bpf_open_raw_sock(const char *name);
 
 int bpf_attach_kprobe(int progfd, const char *event, const char *event_desc, pid_t pid, int cpu, int group_fd);
+int bpf_detach_kprobe(const char *event_desc);
 
 #define LOG_BUF_SIZE 65536
 extern char bpf_log_buf[LOG_BUF_SIZE];

src/python/bpf/__init__.py

@@ -12,10 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import atexit
 import ctypes as ct
 import os
 
-lib = ct.cdll.LoadLibrary("libbpfprog.so")
+lib = ct.CDLL("libbpfprog.so")
 
 # keep in sync with bpf_common.h
 lib.bpf_module_create.restype = ct.c_void_p
@@ -54,6 +55,17 @@ lib.bpf_prog_load.argtypes = [ct.c_int, ct.c_void_p, ct.c_size_t,
         ct.c_char_p, ct.c_uint]
 lib.bpf_attach_kprobe.restype = ct.c_int
 lib.bpf_attach_kprobe.argtypes = [ct.c_int, ct.c_char_p, ct.c_char_p, ct.c_int, ct.c_int, ct.c_int]
+lib.bpf_detach_kprobe.restype = ct.c_int
+lib.bpf_detach_kprobe.argtypes = [ct.c_char_p]
+
+open_kprobes = {}
+
+@atexit.register
+def cleanup_kprobes():
+    for k, v in open_kprobes.items():
+        os.close(v)
+        desc = "-:kprobes/%s" % k
+        lib.bpf_detach_kprobe(desc.encode("ascii"))
 
 class BPF(object):
     SOCKET_FILTER = 1
@@ -185,8 +197,21 @@ class BPF(object):
                 desc.encode("ascii"), pid, cpu, group_fd)
         if res < 0:
             raise Exception("Failed to attach BPF to kprobe")
+        open_kprobes[ev_name] = res
         return res
 
+    @staticmethod
+    def detach_kprobe(event):
+        ev_name = "p_" + event.replace("+", "_")
+        if ev_name not in open_kprobes:
+            raise Exception("Kprobe %s is not attached" % event)
+        os.close(open_kprobes[ev_name])
+        desc = "-:kprobes/%s" % ev_name
+        res = lib.bpf_detach_kprobe(desc.encode("ascii"))
+        if res < 0:
+            raise Exception("Failed to detach BPF from kprobe")
+        del open_kprobes[ev_name]
+
     @staticmethod
     def attach_kretprobe(fn, event, pid=-1, cpu=0, group_fd=-1):
         if not isinstance(fn, BPF.Function):
@@ -197,5 +222,18 @@ class BPF(object):
                 desc.encode("ascii"), pid, cpu, group_fd)
         if res < 0:
             raise Exception("Failed to attach BPF to kprobe")
+        open_kprobes[ev_name] = res
         return res
 
+    @staticmethod
+    def detach_kretprobe(event):
+        ev_name = "r_" + event.replace("+", "_")
+        if ev_name not in open_kprobes:
+            raise Exception("Kretprobe %s is not attached" % event)
+        os.close(open_kprobes[ev_name])
+        desc = "-:kprobes/%s" % ev_name
+        res = lib.bpf_detach_kprobe(desc.encode("ascii"))
+        if res < 0:
+            raise Exception("Failed to detach BPF from kprobe")
+        del open_kprobes[ev_name]
+