tests/python/test_usdt.py

   1 #!/usr/bin/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
   8 from __future__ import print_function
   9 from bcc import BPF, USDT
  10 from unittest import main, TestCase
  11 from subprocess import Popen, PIPE
  12 from tempfile import NamedTemporaryFile
  13 import ctypes as ct
  14 import inspect
  15 import os
  16 import signal
  17
  18 class 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
  28 int 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
  63 struct probe_result_t1 {
  64   char v1;
  65   int  v2;
  66 };
  67
  68 struct probe_result_t2 {
  69   int  v1;
  70   char v2;
  71 };
  72
  73 struct probe_result_t3 {
  74   int v1;
  75   int v2;
  76 };
  77
  78 struct probe_result_t4 {
  79   u64  v1;
  80   char v2[8];
  81 };
  82
  83 struct probe_result_t5 {
  84   char v1[8];
  85   u64  v2;
  86 };
  87
  88 BPF_PERF_OUTPUT(event1);
  89 BPF_PERF_OUTPUT(event2);
  90 BPF_PERF_OUTPUT(event3);
  91 BPF_PERF_OUTPUT(event4);
  92 BPF_PERF_OUTPUT(event5);
  93
  94 int 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 };
 101 int 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 }
 108 int 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 }
 115 int 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 }
 122 int 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.kprobe_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
 225 if __name__ == "__main__":
 226     main()