Add python support for attaching bpf programs to uprobes

This adds a similar set of functions to kprobes for userspace probes.
The calling convention is different, however, since the user must
provide a library/binary name and function symbol or address. Add two
simple test cases for both.

I don't see an easy way in this api to add auto-loading support, as in
kprobe functions that start with "kprobe__". Such niceties can come
later.

Signed-off-by: Brenden Blanco <[email protected]>

src/python/bcc/__init__.py

@@ -20,7 +20,9 @@
 import json
 import multiprocessing
 import os
+import re
 from subprocess import Popen, PIPE
+import struct
 import sys
 basestring = (unicode if sys.version_info[0] < 3 else str)
 
@@ -97,6 +99,14 @@
  ct.c_int, ct.c_int, _CB_TYPE, ct.py_object]
 lib.bpf_detach_kprobe.restype = ct.c_int
 lib.bpf_detach_kprobe.argtypes = [ct.c_char_p]
+lib.bpf_attach_uprobe.restype = ct.c_void_p
+_CB_TYPE = ct.CFUNCTYPE(None, ct.py_object, ct.c_int,
+ ct.c_ulonglong, ct.POINTER(ct.c_ulonglong))
+_RAW_CB_TYPE = ct.CFUNCTYPE(None, ct.py_object, ct.c_void_p, ct.c_int)
+lib.bpf_attach_uprobe.argtypes = [ct.c_int, ct.c_char_p, ct.c_char_p, ct.c_int,
+ ct.c_int, ct.c_int, _CB_TYPE, ct.py_object]
+lib.bpf_detach_uprobe.restype = ct.c_int
+lib.bpf_detach_uprobe.argtypes = [ct.c_char_p]
 lib.bpf_open_perf_buffer.restype = ct.c_void_p
 lib.bpf_open_perf_buffer.argtypes = [_RAW_CB_TYPE, ct.py_object, ct.c_int, ct.c_int]
 lib.perf_reader_poll.restype = ct.c_int
@@ -107,6 +117,7 @@
 lib.perf_reader_fd.argtypes = [ct.c_void_p]
 
 open_kprobes = {}
+open_uprobes = {}
 tracefile = None
 TRACEFS = "/sys/kernel/debug/tracing"
 KALLSYMS = "/proc/kallsyms"
@@ -122,7 +133,13 @@ def cleanup_kprobes():
  if isinstance(k, str):
  desc = "-:kprobes/%s" % k
  lib.bpf_detach_kprobe(desc.encode("ascii"))
+ for k, v in open_uprobes.items():
+ lib.perf_reader_free(v)
+ if isinstance(k, str):
+ desc = "-:uprobes/%s" % k
+ lib.bpf_detach_uprobe(desc.encode("ascii"))
  open_kprobes.clear()
+ open_uprobes.clear()
  if tracefile:
  tracefile.close()
 
@@ -137,6 +154,11 @@ class BPF(object):
  PROG_ARRAY = 3
  PERF_EVENT_ARRAY = 4
 
+ _probe_repl = re.compile("[^a-zA-Z0-9_]")
+ _libsearch_cache = {}
+ _lib_load_address_cache = {}
+ _lib_symbol_cache = {}
+
  class Function(object):
  def __init__(self, bpf, name, fd):
  self.bpf = bpf
@@ -670,6 +692,141 @@ def detach_kretprobe(event):
  raise Exception("Failed to detach BPF from kprobe")
  del open_kprobes[ev_name]
 
+ @classmethod
+ def find_library(cls, name):
+ if name in cls._libsearch_cache:
+ return cls._libsearch_cache[name]
+
+ if struct.calcsize("l") == 4:
+ machine = os.uname()[4] + "-32"
+ else:
+ machine = os.uname()[4] + "-64"
+ mach_map = {
+ "x86_64-64": "libc6,x86-64",
+ "ppc64-64": "libc6,64bit",
+ "sparc64-64": "libc6,64bit",
+ "s390x-64": "libc6,64bit",
+ "ia64-64": "libc6,IA-64",
+ }
+ abi_type = mach_map.get(machine, "libc6")
+ expr = r"\s+lib%s\.[^\s]+\s+\(%s, [^)]+[^/]+([^\s]+)" % (name, abi_type)
+ with os.popen("/sbin/ldconfig -p 2>/dev/null") as f:
+ data = f.read()
+ res = re.search(expr, data)
+ if not res:
+ return None
+ path = res.group(1)
+ cls._libsearch_cache[name] = path
+ return path
+
+ @classmethod
+ def find_load_address(cls, path):
+ if path in cls._lib_load_address_cache:
+ return cls._lib_load_address_cache[path]
+
+ # "LOAD off 0x0000000000000000 vaddr 0x0000000000400000 paddr 0x..."
+ with os.popen("""/usr/bin/objdump -x %s | awk '$1 == "LOAD" && $3 ~ /^[0x]*$/ { print $5 }'""" % path) as f:
+ data = f.read().rstrip()
+ if not data:
+ return None
+ addr = int(data, 16)
+ cls._lib_load_address_cache[path] = addr
+ cls._lib_symbol_cache[path] = {}
+ return addr
+
+ @classmethod
+ def find_symbol(cls, path, sym):
+ # initialized in find_load_address
+ symbols = cls._lib_symbol_cache[path]
+ if sym in symbols:
+ return symbols[sym]
+
+ with os.popen("""objdump -tT %s | awk -v sym=%s '$NF == sym && $4 == ".text" { print $1; exit }'""" % (path, sym)) as f:
+ data = f.read().rstrip()
+ if not data:
+ return None
+ addr = int(data, 16)
+ symbols[sym] = addr
+ return addr
+
+ @classmethod
+ def _check_path_symbol(cls, name, sym, addr):
+ if name.startswith("/"):
+ path = name
+ else:
+ path = BPF.find_library(name)
+ if not path:
+ raise Exception("could not find library %s" % name)
+ path = os.path.realpath(path)
+ load_addr = BPF.find_load_address(path)
+
+ if not addr and sym:
+ addr = BPF.find_symbol(path, sym)
+ if not addr:
+ raise Exception("could not determine address of symbol %s" % sym)
+
+ return (path, load_addr+addr)
+
+ def attach_uprobe(self, name="", sym="", addr=None,
+ fn_name="", pid=-1, cpu=0, group_fd=-1):
+
+ (path, addr) = BPF._check_path_symbol(name, sym, addr)
+
+ fn = self.load_func(fn_name, BPF.KPROBE)
+ ev_name = "p_%s_0x%x" % (self._probe_repl.sub("_", path), addr)
+ desc = "p:uprobes/%s %s:0x%x" % (ev_name, path, addr)
+ res = lib.bpf_attach_uprobe(fn.fd, ev_name.encode("ascii"),
+ desc.encode("ascii"), pid, cpu, group_fd,
+ self._reader_cb_impl, ct.cast(id(self), ct.py_object))
+ res = ct.cast(res, ct.c_void_p)
+ if res == None:
+ raise Exception("Failed to attach BPF to uprobe")
+ open_uprobes[ev_name] = res
+ return self
+
+ @classmethod
+ def detach_uprobe(cls, name="", sym="", addr=None):
+ (path, addr) = BPF._check_path_symbol(name, sym, addr)
+ ev_name = "p_%s_0x%x" % (cls._probe_repl.sub("_", path), addr)
+ if ev_name not in open_uprobes:
+ raise Exception("Uprobe %s is not attached" % event)
+ lib.perf_reader_free(open_uprobes[ev_name])
+ desc = "-:uprobes/%s" % ev_name
+ res = lib.bpf_detach_uprobe(desc.encode("ascii"))
+ if res < 0:
+ raise Exception("Failed to detach BPF from uprobe")
+ del open_uprobes[ev_name]
+
+ def attach_uretprobe(self, name="", sym="", addr=None,
+ fn_name="", pid=-1, cpu=0, group_fd=-1):
+
+ (path, addr) = BPF._check_path_symbol(name, sym, addr)
+
+ fn = self.load_func(fn_name, BPF.KPROBE)
+ ev_name = "r_%s_0x%x" % (self._probe_repl.sub("_", path), addr)
+ desc = "r:uprobes/%s %s:0x%x" % (ev_name, path, addr)
+ res = lib.bpf_attach_uprobe(fn.fd, ev_name.encode("ascii"),
+ desc.encode("ascii"), pid, cpu, group_fd,
+ self._reader_cb_impl, ct.cast(id(self), ct.py_object))
+ res = ct.cast(res, ct.c_void_p)
+ if res == None:
+ raise Exception("Failed to attach BPF to uprobe")
+ open_uprobes[ev_name] = res
+ return self
+
+ @classmethod
+ def detach_uretprobe(cls, name="", sym="", addr=None):
+ (path, addr) = BPF._check_path_symbol(name, sym, addr)
+ ev_name = "r_%s_0x%x" % (cls._probe_repl.sub("_", path), addr)
+ if ev_name not in open_uprobes:
+ raise Exception("Kretprobe %s is not attached" % event)
+ lib.perf_reader_free(open_uprobes[ev_name])
+ desc = "-:uprobes/%s" % ev_name
+ res = lib.bpf_detach_uprobe(desc.encode("ascii"))
+ if res < 0:
+ raise Exception("Failed to detach BPF from uprobe")
+ del open_uprobes[ev_name]
+
  def _trace_autoload(self):
  # Cater to one-liner case where attach_kprobe is omitted and C function
  # name matches that of the kprobe.

tests/cc/CMakeLists.txt

@@ -46,3 +46,5 @@ add_test(NAME py_test_callchain WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
  COMMAND ${TEST_WRAPPER} py_callchain sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_callchain.py)
 add_test(NAME py_array WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
  COMMAND ${TEST_WRAPPER} py_array sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_array.py)
+add_test(NAME py_uprobes WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+ COMMAND ${TEST_WRAPPER} py_uprobes sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_uprobes.py)

tests/cc/test_uprobes.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python
+# Copyright (c) PLUMgrid, Inc.
+# Licensed under the Apache License, Version 2.0 (the "License")
+
+import bcc
+import ctypes
+import os
+import re
+import struct
+import time
+import unittest
+
+## code from ctypes impl
+#if struct.calcsize("l") == 4:
+# machine = os.uname()[4] + "-32"
+#else:
+# machine = os.uname()[4] + "-64"
+#mach_map = {
+# "x86_64-64": "libc6,x86-64",
+# "ppc64-64": "libc6,64bit",
+# "sparc64-64": "libc6,64bit",
+# "s390x-64": "libc6,64bit",
+# "ia64-64": "libc6,IA-64",
+#}
+#abi_type = mach_map.get(machine, "libc6")
+#
+#def find_library_fullpath(name):
+# expr = r"\s+lib%s\.[^\s]+\s+\(%s, [^)]+[^/]+([^\s]+)" % (name, abi_type)
+# with os.popen("/sbin/ldconfig -p 2>/dev/null") as f:
+# data = f.read()
+# res = re.search(expr, data)
+# if not res:
+# return None
+# return res.group(1)
+
+class TestUprobes(unittest.TestCase):
+ def test_simple_library(self):
+ text = """
+#include <uapi/linux/ptrace.h>
+BPF_TABLE("array", int, u64, stats, 1);
+static void incr(int idx) {
+ u64 *ptr = stats.lookup(&idx);
+ if (ptr)
+ ++(*ptr);
+}
+int count(struct pt_regs *ctx) {
+ u32 pid = bpf_get_current_pid_tgid();
+ if (pid == PID)
+ incr(0);
+ return 0;
+}"""
+ text = text.replace("PID", "%d" % os.getpid())
+ b = bcc.BPF(text=text)
+ b.attach_uprobe(name="c", sym="malloc_stats", fn_name="count")
+ b.attach_uretprobe(name="c", sym="malloc_stats", fn_name="count")
+ libc = ctypes.CDLL("libc.so.6")
+ libc.malloc_stats.restype = None
+ libc.malloc_stats.argtypes = []
+ libc.malloc_stats()
+ self.assertEqual(b["stats"][ctypes.c_int(0)].value, 2)
+ b.detach_uprobe(name="c", sym="malloc_stats")
+
+ def test_simple_binary(self):
+ text = """
+#include <uapi/linux/ptrace.h>
+BPF_TABLE("array", int, u64, stats, 1);
+static void incr(int idx) {
+ u64 *ptr = stats.lookup(&idx);
+ if (ptr)
+ ++(*ptr);
+}
+int count(struct pt_regs *ctx) {
+ u32 pid = bpf_get_current_pid_tgid();
+ incr(0);
+ return 0;
+}"""
+ text = text.replace("PID", "%d" % os.getpid())
+ b = bcc.BPF(text=text)
+ b.attach_uprobe(name="/usr/bin/python2", sym="main", fn_name="count")
+ b.attach_uretprobe(name="/usr/bin/python2", sym="main", fn_name="count")
+ with os.popen("/usr/bin/python2 -V") as f:
+ pass
+ self.assertGreater(b["stats"][ctypes.c_int(0)].value, 0)
+ b.detach_uprobe(name="/usr/bin/python2", sym="main")
+
+if __name__ == "__main__":
+ unittest.main()