From 5a888c78762c59d6b62807158b49ff11d999f485 Mon Sep 17 00:00:00 2001 From: Brenden Blanco Date: Mon, 24 Aug 2015 21:33:25 -0700 Subject: [PATCH] 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 /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 --- src/python/bpf/__init__.py | 33 +++++++++++++++++++++++++++++++-- tests/cc/CMakeLists.txt | 2 ++ tests/cc/test_trace4.py | 40 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 73 insertions(+), 2 deletions(-) create mode 100755 tests/cc/test_trace4.py diff --git a/src/python/bpf/__init__.py b/src/python/bpf/__init__.py index 10ee07d..b53c721 100644 --- 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 --- 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 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() -- 2.7.4