Commit 88fda697 authored by 4ast's avatar 4ast

Merge pull request #162 from iovisor/bblanco_dev

Add regex support to attach_k[ret]probe
parents f39d745a 5a888c78
...@@ -18,6 +18,7 @@ import ctypes as ct ...@@ -18,6 +18,7 @@ import ctypes as ct
import fcntl import fcntl
import json import json
import os import os
from subprocess import Popen, PIPE
import sys import sys
basestring = (unicode if sys.version_info[0] < 3 else str) basestring = (unicode if sys.version_info[0] < 3 else str)
...@@ -380,7 +381,23 @@ class BPF(object): ...@@ -380,7 +381,23 @@ class BPF(object):
% (dev, errstr)) % (dev, errstr))
fn.sock = sock 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) fn = self.load_func(fn_name, BPF.KPROBE)
ev_name = "p_" + event.replace("+", "_") ev_name = "p_" + event.replace("+", "_")
desc = "p:kprobes/%s %s" % (ev_name, event) desc = "p:kprobes/%s %s" % (ev_name, event)
...@@ -403,7 +420,16 @@ class BPF(object): ...@@ -403,7 +420,16 @@ class BPF(object):
raise Exception("Failed to detach BPF from kprobe") raise Exception("Failed to detach BPF from kprobe")
del open_kprobes[ev_name] 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) fn = self.load_func(fn_name, BPF.KPROBE)
ev_name = "r_" + event.replace("+", "_") ev_name = "r_" + event.replace("+", "_")
desc = "r:kprobes/%s %s" % (ev_name, event) desc = "r:kprobes/%s %s" % (ev_name, event)
...@@ -452,6 +478,8 @@ class BPF(object): ...@@ -452,6 +478,8 @@ class BPF(object):
""" """
line = BPF.trace_readline(nonblocking) line = BPF.trace_readline(nonblocking)
if line: if line:
# don't print messages related to lost events
if line.startswith("CPU:"): return
task = line[:16].lstrip() task = line[:16].lstrip()
line = line[17:] line = line[17:]
ts_end = line.find(":") ts_end = line.find(":")
...@@ -493,6 +521,7 @@ class BPF(object): ...@@ -493,6 +521,7 @@ class BPF(object):
while True: while True:
if fmt: if fmt:
fields = BPF.trace_readline_fields(nonblocking=False) fields = BPF.trace_readline_fields(nonblocking=False)
if not fields: continue
line = fmt.format(*fields) line = fmt.format(*fields)
else: else:
line = BPF.trace_readline(nonblocking=False) line = BPF.trace_readline(nonblocking=False)
......
...@@ -30,6 +30,8 @@ add_test(NAME py_test_trace2 WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} ...@@ -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) 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} 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) 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} 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) 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} add_test(NAME py_test_brb2 WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
......
#!/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()
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment