add usdt probe type

src/ast/semantic_analyser.cpp

@@ -4,6 +4,7 @@
 #include "parser.tab.hh"
 #include "printf.h"
 #include "arch/arch.h"
+#include <sys/stat.h>
 
 #include "libbpf.h"
 
@@ -510,6 +511,13 @@ void SemanticAnalyser::visit(AttachPoint &ap)
     if (ap.func == "")
       err_ << "uprobes should be attached to a function" << std::endl;
   }
+  else if (ap.provider == "usdt") {
+    if (ap.target == "" || ap.func == "")
+      err_ << "usdt probe must have a target" << std::endl;
+    struct stat s;
+    if (stat(ap.target.c_str(), &s) != 0)
+      err_ << "usdt target file " << ap.target << " does not exist" << std::endl;
+  }
   else if (ap.provider == "tracepoint") {
     if (ap.target == "" || ap.func == "")
       err_ << "tracepoint probe must have a target" << std::endl;

src/attached_probe.cpp

@@ -9,6 +9,7 @@
 #include "attached_probe.h"
 #include "bpftrace.h"
 #include "bcc_syms.h"
+#include "bcc_usdt.h"
 #include "common.h"
 #include "libbpf.h"
 #include <linux/perf_event.h>
@@ -26,6 +27,7 @@ bpf_probe_attach_type attachtype(ProbeType t)
     case ProbeType::kretprobe: return BPF_PROBE_RETURN; break;
     case ProbeType::uprobe:    return BPF_PROBE_ENTRY;  break;
     case ProbeType::uretprobe: return BPF_PROBE_RETURN; break;
+    case ProbeType::usdt:      return BPF_PROBE_ENTRY; break;
     default: abort();
   }
 }
@@ -38,6 +40,7 @@ bpf_prog_type progtype(ProbeType t)
     case ProbeType::kretprobe:  return BPF_PROG_TYPE_KPROBE; break;
     case ProbeType::uprobe:     return BPF_PROG_TYPE_KPROBE; break;
     case ProbeType::uretprobe:  return BPF_PROG_TYPE_KPROBE; break;
+    case ProbeType::usdt:       return BPF_PROG_TYPE_KPROBE; break;
     case ProbeType::tracepoint: return BPF_PROG_TYPE_TRACEPOINT; break;
     case ProbeType::profile:      return BPF_PROG_TYPE_PERF_EVENT; break;
     case ProbeType::interval:      return BPF_PROG_TYPE_PERF_EVENT; break;
@@ -82,6 +85,20 @@ AttachedProbe::AttachedProbe(Probe &probe, std::tuple<uint8_t *, uintptr_t> func
   }
 }
 
+AttachedProbe::AttachedProbe(Probe &probe, std::tuple<uint8_t *, uintptr_t> func, int pid)
+  : probe_(probe), func_(func)
+{
+  load_prog();
+  switch (probe_.type)
+  {
+    case ProbeType::usdt:
+      attach_usdt(pid);
+      break;
+    default:
+      abort();
+  }
+}
+
 AttachedProbe::~AttachedProbe()
 {
   close(progfd_);
@@ -103,6 +120,7 @@ AttachedProbe::~AttachedProbe()
       break;
     case ProbeType::uprobe:
     case ProbeType::uretprobe:
+    case ProbeType::usdt:
       err = bpf_detach_uprobe(eventname().c_str());
       break;
     case ProbeType::tracepoint:
@@ -143,6 +161,7 @@ std::string AttachedProbe::eventname() const
       return eventprefix() + probe_.attach_point;
     case ProbeType::uprobe:
     case ProbeType::uretprobe:
+    case ProbeType::usdt:
       offset_str << std::hex << offset();
       return eventprefix() + sanitise(probe_.path) + "_" + offset_str.str();
     case ProbeType::tracepoint:
@@ -161,7 +180,7 @@ uint64_t AttachedProbe::offset() const
 {
   bcc_symbol sym;
   int err = bcc_resolve_symname(probe_.path.c_str(), probe_.attach_point.c_str(),
-      0, 0, nullptr, &sym);
+      probe_.loc, 0, nullptr, &sym);
 
   if (err)
     throw std::runtime_error("Could not resolve symbol: " + probe_.path + ":" + probe_.attach_point);
@@ -277,6 +296,57 @@ void AttachedProbe::attach_uprobe()
   perf_event_fds_.push_back(perf_event_fd);
 }
 
+void AttachedProbe::attach_usdt(int pid)
+{
+  struct bcc_usdt_location loc = {};
+  int err, i;
+  std::ostringstream offset_str;
+  void *ctx;
+
+  if (pid)
+  {
+    ctx = bcc_usdt_new_frompid(pid, probe_.path.c_str());
+    if (!ctx)
+      throw std::runtime_error("Error initializing context for probe: " + probe_.name + ", for PID: " + std::to_string(pid));
+  }
+  else
+  {
+    ctx = bcc_usdt_new_frompath(probe_.path.c_str());
+    if (!ctx)
+      throw std::runtime_error("Error initializing context for probe: " + probe_.name);
+  }
+
+  // TODO: fn_name may need a unique suffix for each attachment on the same probe:
+  std::string fn_name = "probe_" + probe_.attach_point + "_1";
+  err = bcc_usdt_enable_probe(ctx, probe_.attach_point.c_str(), fn_name.c_str());
+  if (err)
+    throw std::runtime_error("Error finding or enabling probe: " + probe_.name);
+
+  std::string provider_name;
+  if ((i = probe_.path.rfind("/")) != std::string::npos)
+     provider_name = probe_.path.substr(i + 1);
+  else
+     provider_name = probe_.path;
+
+  err = bcc_usdt_get_location(ctx, provider_name.c_str(), probe_.attach_point.c_str(), 0, &loc);
+  if (err)
+    throw std::runtime_error("Error finding location for probe: " + probe_.name);
+  probe_.loc = loc.address;
+
+  int perf_event_fd = bpf_attach_uprobe(progfd_, attachtype(probe_.type),
+      eventname().c_str(), probe_.path.c_str(), loc.address - 0x400000, pid == 0 ? -1 : pid);
+
+  if (perf_event_fd < 0)
+  {
+    if (pid)
+      throw std::runtime_error("Error attaching probe: " + probe_.name + ", to PID: " + std::to_string(pid));
+    else
+      throw std::runtime_error("Error attaching probe: " + probe_.name);
+  }
+
+  perf_event_fds_.push_back(perf_event_fd);
+}
+
 void AttachedProbe::attach_tracepoint()
 {
   int perf_event_fd = bpf_attach_tracepoint(progfd_, probe_.path.c_str(),

src/attached_probe.h

@@ -13,6 +13,7 @@ class AttachedProbe
 {
 public:
   AttachedProbe(Probe &probe, std::tuple<uint8_t *, uintptr_t> func);
+  AttachedProbe(Probe &probe, std::tuple<uint8_t *, uintptr_t> func, int pid);
   ~AttachedProbe();
   AttachedProbe(const AttachedProbe &) = delete;
   AttachedProbe& operator=(const AttachedProbe &) = delete;
@@ -25,6 +26,7 @@ private:
   void load_prog();
   void attach_kprobe();
   void attach_uprobe();
+  void attach_usdt(int pid);
   void attach_tracepoint();
   void attach_profile();
   void attach_interval();

src/bpftrace.cpp

@@ -259,7 +259,10 @@ std::unique_ptr<AttachedProbe> BPFtrace::attach_probe(Probe &probe, const BpfOrc
   }
   try
   {
-    return std::make_unique<AttachedProbe>(probe, func->second);
+    if (probe.type == ProbeType::usdt)
+      return std::make_unique<AttachedProbe>(probe, func->second, pid_);
+    else
+      return std::make_unique<AttachedProbe>(probe, func->second);
   }
   catch (std::runtime_error e)
   {

src/bpftrace.h

@@ -37,6 +37,7 @@ public:
   std::string get_stack(uint32_t stackid, bool ustack, int indent=0);
   std::string resolve_sym(uintptr_t addr, bool show_offset=false);
   std::string resolve_usym(uintptr_t addr) const;
+  int pid_;
 
   std::map<std::string, std::unique_ptr<IMap>> maps_;
   std::map<std::string, Struct> structs_;

src/main.cpp

@@ -19,6 +19,7 @@ void usage()
   std::cerr << "OPTIONS:" << std::endl;
   std::cerr << "    -l [search]    list probes" << std::endl;
   std::cerr << "    -e 'program'   execute this program" << std::endl;
+  std::cerr << "    -p PID    PID for enabling USDT probes" << std::endl;
   std::cerr << "    -v    verbose messages" << std::endl;
   std::cerr << "    -d    debug info dry run" << std::endl << std::endl;
   std::cerr << "EXAMPLES:" << std::endl;
@@ -34,11 +35,12 @@ int main(int argc, char *argv[])
 {
   int err;
   Driver driver;
+  char *pid_str = NULL;
   bool listing = false;
 
   std::string script, search;
   int c;
-  while ((c = getopt(argc, argv, "de:lv")) != -1)
+  while ((c = getopt(argc, argv, "de:lp:v")) != -1)
   {
     switch (c)
     {
@@ -51,6 +53,9 @@ int main(int argc, char *argv[])
       case 'e':
         script = optarg;
         break;
+      case 'p':
+        pid_str = optarg;
+        break;
       case 'l':
         listing = true;
         break;
@@ -108,6 +113,13 @@ int main(int argc, char *argv[])
 
   BPFtrace bpftrace;
 
+  // PID is currently only used for USDT probes that need enabling. Future work:
+  // - make PID a filter for all probe types: pass to perf_event_open(), etc.
+  // - provide PID in USDT probe specification as a way to override -p.
+  bpftrace.pid_ = 0;
+  if (pid_str)
+    bpftrace.pid_ = atoi(pid_str);
+
   if (bt_debug)
   {
     ast::Printer p(std::cout);

src/types.cpp

@@ -54,6 +54,8 @@ ProbeType probetype(const std::string &type)
     return ProbeType::uprobe;
   else if (type == "uretprobe")
     return ProbeType::uretprobe;
+  else if (type == "usdt")
+    return ProbeType::usdt;
   else if (type == "BEGIN")
     return ProbeType::uprobe;
   else if (type == "END")

src/types.h

@@ -54,6 +54,7 @@ enum class ProbeType
   kretprobe,
   uprobe,
   uretprobe,
+  usdt,
   tracepoint,
   profile,
   interval,
@@ -72,6 +73,7 @@ public:
   std::string attach_point;
   std::string prog_name;
   std::string name;
+  uint64_t loc;
   int freq;
 };
 

tests/ast.cpp

@@ -48,6 +48,19 @@ TEST(ast, probe_name_uprobe)
   EXPECT_EQ(uprobe2.name(), "uprobe:/bin/sh:readline,uprobe:/bin/sh:somefunc");
 }
 
+TEST(ast, probe_name_usdt)
+{
+  AttachPoint ap1("usdt", "/bin/sh", "probe1");
+  AttachPointList attach_points1 = { &ap1 };
+  Probe usdt1(&attach_points1, nullptr, nullptr);
+  EXPECT_EQ(usdt1.name(), "usdt:/bin/sh:probe1");
+
+  AttachPoint ap2("usdt", "/bin/sh", "probe2");
+  AttachPointList attach_points2 = { &ap1, &ap2 };
+  Probe usdt2(&attach_points2, nullptr, nullptr);
+  EXPECT_EQ(usdt2.name(), "usdt:/bin/sh:probe1,usdt:/bin/sh:probe2");
+}
+
 TEST(ast, attach_point_name)
 {
   AttachPoint ap1("kprobe", "sys_read");

tests/bpftrace.cpp

@@ -43,6 +43,14 @@ void check_uprobe(Probe &p, const std::string &path, const std::string &attach_p
   EXPECT_EQ("uprobe:" + path + ":" + attach_point, p.name);
 }
 
+void check_usdt(Probe &p, const std::string &path, const std::string &attach_point, const std::string &prog_name)
+{
+  EXPECT_EQ(ProbeType::usdt, p.type);
+  EXPECT_EQ(attach_point, p.attach_point);
+  EXPECT_EQ(prog_name, p.prog_name);
+  EXPECT_EQ("usdt:" + path + ":" + attach_point, p.name);
+}
+
 void check_tracepoint(Probe &p, const std::string &target, const std::string &func, const std::string &prog_name)
 {
   EXPECT_EQ(ProbeType::tracepoint, p.type);
@@ -244,6 +252,20 @@ TEST(bpftrace, add_probes_uprobe)
   check_uprobe(bpftrace.get_probes().at(0), "/bin/sh", "foo", "uprobe:/bin/sh:foo");
 }
 
+TEST(bpftrace, add_probes_usdt)
+{
+  ast::AttachPoint a("usdt", "/bin/sh", "foo");
+  ast::AttachPointList attach_points = { &a };
+  ast::Probe probe(&attach_points, nullptr, nullptr);
+
+  StrictMock<MockBPFtrace> bpftrace;
+
+  EXPECT_EQ(0, bpftrace.add_probe(probe));
+  EXPECT_EQ(1, bpftrace.get_probes().size());
+  EXPECT_EQ(0, bpftrace.get_special_probes().size());
+  check_usdt(bpftrace.get_probes().at(0), "/bin/sh", "foo", "usdt:/bin/sh:foo");
+}
+
 TEST(bpftrace, add_probes_uprobe_wildcard)
 {
   ast::AttachPoint a("uprobe", "/bin/sh", "foo*");

tests/parser.cpp

@@ -263,6 +263,14 @@ TEST(Parser, uprobe)
       "  int: 1\n");
 }
 
+TEST(Parser, usdt)
+{
+  test("usdt:/my/program:probe { 1; }",
+      "Program\n"
+      " usdt:/my/program:probe\n"
+      "  int: 1\n");
+}
+
 TEST(Parser, escape_chars)
 {
   test("kprobe:sys_open { \"newline\\nand tab\\tbackslash\\\\quote\\\"here\" }",

tests/semantic_analyser.cpp

@@ -421,6 +421,13 @@ TEST(semantic_analyser, uprobe)
   test("uretprobe { 1 }", 1);
 }
 
+TEST(semantic_analyser, usdt)
+{
+  test("usdt:/bin/sh:probe { 1 }", 0);
+  test("usdt:/notexistfile:probe { 1 }", 1);
+  test("usdt { 1 }", 1);
+}
+
 TEST(semantic_analyser, begin_end_probes)
 {
   test("BEGIN { 1 }", 0);