1 /*
2  * RecordMySQLQuery Record MySQL queries by probing the alloc_query() function
3  *                  in mysqld. For Linux, uses BCC, eBPF. Embedded C.
4  *
5  * Basic example of BCC and uprobes.
6  *
7  * Copyright (c) Facebook, Inc.
8  * Licensed under the Apache License, Version 2.0 (the "License")
9  */
10 
11 #include <unistd.h>
12 #include <algorithm>
13 #include <cstdlib>
14 #include <iostream>
15 #include <string>
16 
17 #include "BPF.h"
18 
19 const std::string BPF_PROGRAM = R"(
20 #include <linux/ptrace.h>
21 
22 struct query_probe_t {
23   uint64_t ts;
24   pid_t pid;
25   char query[100];
26 };
27 
28 BPF_HASH(queries, struct query_probe_t, int);
29 
30 int probe_mysql_query(struct pt_regs *ctx, void* thd, char* query, size_t len) {
31   if (query) {
32     struct query_probe_t key = {};
33 
34     key.ts = bpf_ktime_get_ns();
35     key.pid = bpf_get_current_pid_tgid();
36 
37     bpf_probe_read_str(&key.query, sizeof(key.query), query);
38 
39     int one = 1;
40     queries.update(&key, &one);
41   }
42   return 0;
43 }
44 )";
45 const std::string ALLOC_QUERY_FUNC = "_Z11alloc_queryP3THDPKcj";
46 
47 // Define the same struct to use in user space.
48 struct query_probe_t {
49   uint64_t ts;
50   pid_t pid;
51   char query[100];
52 };
53 
main(int argc,char ** argv)54 int main(int argc, char** argv) {
55   if (argc < 2) {
56     std::cout << "USAGE: RecordMySQLQuery PATH_TO_MYSQLD [duration]"
57               << std::endl;
58     exit(1);
59   }
60 
61   std::string mysql_path(argv[1]);
62   std::cout << "Using mysqld path: " << mysql_path << std::endl;
63 
64   ebpf::BPF bpf;
65   auto init_res = bpf.init(BPF_PROGRAM);
66   if (init_res.code() != 0) {
67     std::cerr << init_res.msg() << std::endl;
68     return 1;
69   }
70 
71   auto attach_res =
72       bpf.attach_uprobe(mysql_path, ALLOC_QUERY_FUNC, "probe_mysql_query");
73   if (attach_res.code() != 0) {
74     std::cerr << attach_res.msg() << std::endl;
75     return 1;
76   }
77 
78   int probe_time = 10;
79   if (argc >= 3)
80     probe_time = atoi(argv[2]);
81   std::cout << "Probing for " << probe_time << " seconds" << std::endl;
82   sleep(probe_time);
83 
84   auto table_handle = bpf.get_hash_table<query_probe_t, int>("queries");
85   auto table = table_handle.get_table_offline();
86   std::sort(
87       table.begin(), table.end(),
88       [](std::pair<query_probe_t, int> a, std::pair<query_probe_t, int> b) {
89         return a.first.ts < b.first.ts;
90       });
91   std::cout << table.size() << " queries recorded:" << std::endl;
92   for (auto it : table) {
93     std::cout << "Time: " << it.first.ts << " PID: " << it.first.pid
94               << " Query: " << it.first.query << std::endl;
95   }
96 
97   auto detach_res = bpf.detach_uprobe(mysql_path, ALLOC_QUERY_FUNC);
98   if (detach_res.code() != 0) {
99     std::cerr << detach_res.msg() << std::endl;
100     return 1;
101   }
102 
103   return 0;
104 }
105