1#!/usr/bin/env python
2#
3# USAGE: test_usdt.py
4#
5# Copyright 2017 Facebook, Inc
6# Licensed under the Apache License, Version 2.0 (the "License")
7
8from __future__ import print_function
9from bcc import BPF, USDT
10from unittest import main, TestCase
11from subprocess import Popen, PIPE
12from tempfile import NamedTemporaryFile
13import ctypes as ct
14import inspect
15import os
16import signal
17
18class TestUDST(TestCase):
19    def setUp(self):
20        # Application, minimum, to define three trace points
21        app_text = b"""
22#include <unistd.h>
23#include <stdint.h>
24#include <stdio.h>
25#include <string.h>
26#include "folly/tracing/StaticTracepoint.h"
27
28int main() {
29  char s[100];
30  int i, a = 200, b = 40;
31  for (i = 0; i < 100; i++) s[i] = (i & 7) + (i & 6);
32  uint64_t j = 0;
33  char s1[64];
34  const char* str = "str";
35  size_t len = strlen(str);
36  while (1) {
37    FOLLY_SDT(test, probe_point_1, s[7], b);
38    FOLLY_SDT(test, probe_point_3, a, b);
39    FOLLY_SDT(test, probe_point_1, s[4], a);
40    FOLLY_SDT(test, probe_point_2, 5, s[10]);
41    FOLLY_SDT(test, probe_point_3, s[4], s[7]);
42
43    memset(&s1, '\0', sizeof(s1));
44    strncpy(s1, str, len);
45    snprintf(s1 + len, sizeof(s1) - len, "%d", j);
46    FOLLY_SDT(test, probe_point_4, j++, &s1);
47
48    memset(&s1, '\0', sizeof(s1));
49    strncpy(s1, str, len);
50    snprintf(s1 + len, sizeof(s1) - len, "%d", j);
51    FOLLY_SDT(test, probe_point_5, &s1, j++);
52
53    sleep(1);
54  }
55  return 1;
56}
57"""
58        # BPF program
59        self.bpf_text = """
60#include <linux/blkdev.h>
61#include <uapi/linux/ptrace.h>
62
63struct probe_result_t1 {
64  char v1;
65  int  v2;
66};
67
68struct probe_result_t2 {
69  int  v1;
70  char v2;
71};
72
73struct probe_result_t3 {
74  int v1;
75  int v2;
76};
77
78struct probe_result_t4 {
79  u64  v1;
80  char v2[8];
81};
82
83struct probe_result_t5 {
84  char v1[8];
85  u64  v2;
86};
87
88BPF_PERF_OUTPUT(event1);
89BPF_PERF_OUTPUT(event2);
90BPF_PERF_OUTPUT(event3);
91BPF_PERF_OUTPUT(event4);
92BPF_PERF_OUTPUT(event5);
93
94int do_trace1(struct pt_regs *ctx) {
95    struct probe_result_t1 result = {};
96    bpf_usdt_readarg(1, ctx, &result.v1);
97    bpf_usdt_readarg(2, ctx, &result.v2);
98    event1.perf_submit(ctx, &result, sizeof(result));
99    return 0;
100};
101int do_trace2(struct pt_regs *ctx) {
102    struct probe_result_t2 result = {};
103    bpf_usdt_readarg(1, ctx, &result.v1);
104    bpf_usdt_readarg(2, ctx, &result.v2);
105    event2.perf_submit(ctx, &result, sizeof(result));
106    return 0;
107}
108int do_trace3(struct pt_regs *ctx) {
109    struct probe_result_t3 result = {};
110    bpf_usdt_readarg(1, ctx, &result.v1);
111    bpf_usdt_readarg(2, ctx, &result.v2);
112    event3.perf_submit(ctx, &result, sizeof(result));
113    return 0;
114}
115int do_trace4(struct pt_regs *ctx) {
116    struct probe_result_t4 result = {};
117    bpf_usdt_readarg(1, ctx, &result.v1);
118    bpf_usdt_readarg_p(2, ctx, &result.v2, sizeof(result.v2));
119    event4.perf_submit(ctx, &result, sizeof(result));
120    return 0;
121}
122int do_trace5(struct pt_regs *ctx) {
123    struct probe_result_t5 result = {};
124    bpf_usdt_readarg_p(1, ctx, &result.v1, sizeof(result.v1));
125    bpf_usdt_readarg(2, ctx, &result.v2);
126    event5.perf_submit(ctx, &result, sizeof(result));
127    return 0;
128}
129"""
130
131        # Compile and run the application
132        self.ftemp = NamedTemporaryFile(delete=False)
133        self.ftemp.close()
134        comp = Popen(["gcc", "-I", "%s/include" % os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))),
135                      "-x", "c", "-o", self.ftemp.name, "-"],
136                     stdin=PIPE)
137        comp.stdin.write(app_text)
138        comp.stdin.close()
139        self.assertEqual(comp.wait(), 0)
140        self.app = Popen([self.ftemp.name])
141
142    def test_attach1(self):
143        # enable USDT probe from given PID and verifier generated BPF programs
144        u = USDT(pid=int(self.app.pid))
145        u.enable_probe(probe="probe_point_1", fn_name="do_trace1")
146        u.enable_probe(probe="probe_point_2", fn_name="do_trace2")
147        u.enable_probe(probe="probe_point_3", fn_name="do_trace3")
148        u.enable_probe(probe="probe_point_4", fn_name="do_trace4")
149        u.enable_probe(probe="probe_point_5", fn_name="do_trace5")
150        b = BPF(text=self.bpf_text, usdt_contexts=[u], debug=4)
151
152        # Event states for each event:
153        # 0 - probe not caught, 1 - probe caught with correct value,
154        # 2 - probe caught with incorrect value
155        self.evt_st_1 = 0
156        self.evt_st_2 = 0
157        self.evt_st_3 = 0
158
159        # define output data structure in Python
160        class Data1(ct.Structure):
161            _fields_ = [("v1", ct.c_char),
162                        ("v2", ct.c_int)]
163
164        class Data2(ct.Structure):
165            _fields_ = [("v1", ct.c_int),
166                        ("v2", ct.c_char)]
167
168        class Data3(ct.Structure):
169            _fields_ = [("v1", ct.c_int),
170                        ("v2", ct.c_int)]
171
172        class Data4(ct.Structure):
173            _fields_ = [("v1", ct.c_ulonglong),
174                        ("v2", ct.c_char * 64)]
175
176        class Data5(ct.Structure):
177            _fields_ = [("v1", ct.c_char * 64),
178                        ("v2", ct.c_ulonglong)]
179
180        def check_event_val(event, event_state, v1, v2, v3, v4):
181            if ((event.v1 == v1 and event.v2 == v2) or (event.v1 == v3 and event.v2 == v4)):
182                if (event_state == 0 or event_state == 1):
183                    return 1
184            return 2
185
186        def print_event1(cpu, data, size):
187            event = ct.cast(data, ct.POINTER(Data1)).contents
188            self.evt_st_1 = check_event_val(event, self.evt_st_1, b'\x0d', 40, b'\x08', 200)
189
190        def print_event2(cpu, data, size):
191            event = ct.cast(data, ct.POINTER(Data2)).contents
192            # pretend we have two identical probe points to simplify the code
193            self.evt_st_2 = check_event_val(event, self.evt_st_2, 5, b'\x04', 5, b'\x04')
194
195        def print_event3(cpu, data, size):
196            event = ct.cast(data, ct.POINTER(Data3)).contents
197            self.evt_st_3 = check_event_val(event, self.evt_st_3, 200, 40, 8, 13)
198
199        def print_event4(cpu, data, size):
200            event = ct.cast(data, ct.POINTER(Data4)).contents
201            print("%s" % event.v2)
202
203        def print_event5(cpu, data, size):
204            event = ct.cast(data, ct.POINTER(Data5)).contents
205            print("%s" % event.v1)
206
207        # loop with callback to print_event
208        b["event1"].open_perf_buffer(print_event1)
209        b["event2"].open_perf_buffer(print_event2)
210        b["event3"].open_perf_buffer(print_event3)
211        b["event4"].open_perf_buffer(print_event4)
212        b["event5"].open_perf_buffer(print_event5)
213
214        # three iterations to make sure we get some probes and have time to process them
215        for i in range(3):
216            b.perf_buffer_poll()
217        self.assertTrue(self.evt_st_1 == 1 and self.evt_st_2 == 1 and self.evt_st_3 == 1)
218
219    def tearDown(self):
220        # kill the subprocess, clean the environment
221        self.app.kill()
222        self.app.wait()
223        os.unlink(self.ftemp.name)
224
225if __name__ == "__main__":
226    main()
227