1#!/usr/bin/env python
2#
3# USAGE: test_usdt2.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 <stdlib.h>
23#include <unistd.h>
24#include "folly/tracing/StaticTracepoint.h"
25
26int main(int argc, char **argv) {
27  int t = atoi(argv[1]);
28  while (1) {
29    FOLLY_SDT(test, probe_point_1, t);
30    FOLLY_SDT(test, probe_point_2, t + 1);
31    FOLLY_SDT(test, probe_point_3, t + 2);
32    sleep(1);
33  }
34  return 1;
35}
36"""
37        # BPF program
38        self.bpf_text = """
39#include <uapi/linux/ptrace.h>
40
41BPF_PERF_OUTPUT(event1);
42BPF_PERF_OUTPUT(event2);
43BPF_PERF_OUTPUT(event3);
44BPF_PERF_OUTPUT(event4);
45BPF_PERF_OUTPUT(event5);
46BPF_PERF_OUTPUT(event6);
47
48int do_trace1(struct pt_regs *ctx) {
49    u32 pid = bpf_get_current_pid_tgid();
50    int result = 0;
51    bpf_usdt_readarg(1, ctx, &result);
52    if (FILTER)
53      event1.perf_submit(ctx, &result, sizeof(result));
54    else
55      event4.perf_submit(ctx, &result, sizeof(result));
56    return 0;
57};
58int do_trace2(struct pt_regs *ctx) {
59    u32 pid = bpf_get_current_pid_tgid();
60    int result = 0;
61    bpf_usdt_readarg(1, ctx, &result);
62    if (FILTER)
63      event2.perf_submit(ctx, &result, sizeof(result));
64    else
65      event5.perf_submit(ctx, &result, sizeof(result));
66    return 0;
67}
68int do_trace3(struct pt_regs *ctx) {
69    u32 pid = bpf_get_current_pid_tgid();
70    int result = 0;
71    bpf_usdt_readarg(1, ctx, &result);
72    if (FILTER)
73      event3.perf_submit(ctx, &result, sizeof(result));
74    else
75      event6.perf_submit(ctx, &result, sizeof(result));
76    return 0;
77}
78"""
79
80        # Compile and run the application
81        self.ftemp = NamedTemporaryFile(delete=False)
82        self.ftemp.close()
83        comp = Popen(["gcc", "-I", "%s/include" % os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))),
84                      "-x", "c", "-o", self.ftemp.name, "-"],
85                     stdin=PIPE)
86        comp.stdin.write(app_text)
87        comp.stdin.close()
88        self.assertEqual(comp.wait(), 0)
89
90        # create 3 applications, 2 applications will have usdt attached and
91        # the third one does not, and the third one should not call into
92        # bpf program.
93        self.app = Popen([self.ftemp.name, "1"])
94        self.app2 = Popen([self.ftemp.name, "11"])
95        self.app3 = Popen([self.ftemp.name, "21"])
96
97    def test_attach1(self):
98        # Enable USDT probe from given PID and verifier generated BPF programs.
99        u = USDT(pid=int(self.app.pid))
100        u.enable_probe(probe="probe_point_1", fn_name="do_trace1")
101        u.enable_probe(probe="probe_point_2", fn_name="do_trace2")
102        u2 = USDT(pid=int(self.app2.pid))
103        u2.enable_probe(probe="probe_point_2", fn_name="do_trace2")
104        u2.enable_probe(probe="probe_point_3", fn_name="do_trace3")
105        self.bpf_text = self.bpf_text.replace("FILTER", "pid == %d" % self.app.pid)
106        b = BPF(text=self.bpf_text, usdt_contexts=[u, u2])
107
108        # Event states for each event:
109        # 0 - probe not caught, 1 - probe caught with correct value,
110        # 2 - probe caught with incorrect value
111        self.evt_st_1 = 0
112        self.evt_st_2 = 0
113        self.evt_st_3 = 0
114        self.evt_st_4 = 0
115        self.evt_st_5 = 0
116        self.evt_st_6 = 0
117
118        def check_event_val(data, event_state, expected_val):
119            result = ct.cast(data, ct.POINTER(ct.c_int)).contents
120            if result.value == expected_val:
121                if (event_state == 0 or event_state == 1):
122                    return 1
123            return 2
124
125        def print_event1(cpu, data, size):
126            self.evt_st_1 = check_event_val(data, self.evt_st_1, 1)
127
128        def print_event2(cpu, data, size):
129            self.evt_st_2 = check_event_val(data, self.evt_st_2, 2)
130
131        def print_event3(cpu, data, size):
132            self.evt_st_3 = check_event_val(data, self.evt_st_3, 3)
133
134        def print_event4(cpu, data, size):
135            self.evt_st_4 = check_event_val(data, self.evt_st_4, 11)
136
137        def print_event5(cpu, data, size):
138            self.evt_st_5 = check_event_val(data, self.evt_st_5, 12)
139
140        def print_event6(cpu, data, size):
141            self.evt_st_6 = check_event_val(data, self.evt_st_6, 13)
142
143        # loop with callback to print_event
144        b["event1"].open_perf_buffer(print_event1)
145        b["event2"].open_perf_buffer(print_event2)
146        b["event3"].open_perf_buffer(print_event3)
147        b["event4"].open_perf_buffer(print_event4)
148        b["event5"].open_perf_buffer(print_event5)
149        b["event6"].open_perf_buffer(print_event6)
150
151        # three iterations to make sure we get some probes and have time to process them
152        for i in range(3):
153            b.perf_buffer_poll()
154
155        # note that event1 and event4 do not really fire, so their state should be 0
156        # use separate asserts so that if test fails we know which one is the culprit
157        self.assertTrue(self.evt_st_1 == 1)
158        self.assertTrue(self.evt_st_2 == 1)
159        self.assertTrue(self.evt_st_3 == 0)
160        self.assertTrue(self.evt_st_4 == 0)
161        self.assertTrue(self.evt_st_5 == 1)
162        self.assertTrue(self.evt_st_6 == 1)
163
164    def tearDown(self):
165        # kill the subprocess, clean the environment
166        self.app.kill()
167        self.app.wait()
168        self.app2.kill()
169        self.app2.wait()
170        self.app3.kill()
171        self.app3.wait()
172        os.unlink(self.ftemp.name)
173
174if __name__ == "__main__":
175    main()
176