Add wildcard matches for probe functions

src/ast.cpp

@@ -106,5 +106,14 @@ std::string Probe::name() const
   return n.substr(0, n.size()-1);
 }
 
+std::string Probe::name(const std::string &attach_point) const
+{
+  std::string n = type;
+  if (path != "")
+    n += ":" + path;
+  n += ":" + attach_point;
+  return n;
+}
+
 } // namespace ast
 } // namespace bpftrace

src/ast.h

@@ -155,6 +155,7 @@ public:
 
   void accept(Visitor &v) override;
   std::string name() const;
+  std::string name(const std::string &attach_point) const;
 };
 using ProbeList = std::vector<Probe *>;
 

src/attached_probe.cpp

@@ -163,7 +163,7 @@ void AttachedProbe::attach_kprobe()
       pid, cpu, group_fd, cb, cb_cookie);
 
   if (perf_reader_ == nullptr)
-    throw std::runtime_error("Error attaching probe: " + probe_.name);
+    throw std::runtime_error("Error attaching probe: '" + probe_.name + "'");
 }
 
 void AttachedProbe::attach_uprobe()

src/bpftrace.cpp

+#include <fstream>
 #include <iomanip>
 #include <iostream>
+#include <regex>
 #include <sstream>
 #include <sys/epoll.h>
 
@@ -21,6 +23,7 @@ int BPFtrace::add_probe(ast::Probe &p)
     probe.path = "/proc/self/exe";
     probe.attach_point = "BEGIN_trigger";
     probe.type = probetype(p.type);
+    probe.prog_name = p.name();
     probe.name = p.name();
     special_probes_.push_back(probe);
     return 0;
@@ -31,23 +34,76 @@ int BPFtrace::add_probe(ast::Probe &p)
     probe.path = "/proc/self/exe";
     probe.attach_point = "END_trigger";
     probe.type = probetype(p.type);
+    probe.prog_name = p.name();
     probe.name = p.name();
     special_probes_.push_back(probe);
     return 0;
   }
 
+  ast::AttachPointList expanded_attach_points;
   for (std::string attach_point : *p.attach_points)
+  {
+    if (attach_point.find("*") != std::string::npos)
+    {
+      std::string file_name;
+      switch (probetype(p.type))
+      {
+        case ProbeType::kprobe:
+        case ProbeType::kretprobe:
+          file_name = "/sys/kernel/debug/tracing/available_filter_functions";
+          break;
+        default:
+          std::cerr << "Wildcard matches aren't available on probe type '"
+                    << p.type << "'" << std::endl;
+          return 1;
+      }
+      auto matches = find_wildcard_matches(attach_point, file_name);
+      expanded_attach_points.insert(expanded_attach_points.end(),
+          matches.begin(), matches.end());
+      continue;
+    }
+
+    expanded_attach_points.push_back(attach_point);
+  }
+
+  for (std::string attach_point : expanded_attach_points)
   {
     Probe probe;
     probe.path = p.path;
     probe.attach_point = attach_point;
     probe.type = probetype(p.type);
-    probe.name = p.name();
+    probe.prog_name = p.name();
+    probe.name = p.name(attach_point);
     probes_.push_back(probe);
   }
   return 0;
 }
 
+std::set<std::string> BPFtrace::find_wildcard_matches(std::string attach_point, std::string file_name)
+{
+  // Turn glob into a regex
+  attach_point = "^" + std::regex_replace(attach_point, std::regex("\\*"), "[^\\s]*");
+  std::regex attach_point_regex(attach_point);
+  std::smatch match;
+
+  std::ifstream file(file_name);
+  std::string line;
+  std::set<std::string> matches;
+  while (std::getline(file, line))
+  {
+    if (std::regex_search(line, match, attach_point_regex))
+    {
+      matches.insert(match[0]);
+    }
+  }
+  return matches;
+}
+
+int BPFtrace::num_probes() const
+{
+  return special_probes_.size() + probes_.size();
+}
+
 void perf_event_printer(void *cb_cookie, void *data, int size)
 {
   auto bpftrace = static_cast<BPFtrace*>(cb_cookie);
@@ -120,7 +176,7 @@ void perf_event_lost(uint64_t lost)
 
 std::unique_ptr<AttachedProbe> BPFtrace::attach_probe(Probe &probe)
 {
-  auto func = sections_.find(probe.name);
+  auto func = sections_.find(probe.prog_name);
   if (func == sections_.end())
   {
     std::cerr << "Code not generated for probe: " << probe.name << std::endl;

src/bpftrace.h

@@ -2,6 +2,7 @@
 
 #include <map>
 #include <memory>
+#include <set>
 #include <vector>
 
 #include "syms.h"
@@ -18,6 +19,7 @@ class BPFtrace
 public:
   virtual ~BPFtrace() { }
   virtual int add_probe(ast::Probe &p);
+  int num_probes() const;
   int run();
   int print_maps();
   std::string get_stack(uint32_t stackid, bool ustack, int indent=0);
@@ -30,9 +32,12 @@ public:
   std::unique_ptr<Map> stackid_map_;
   std::unique_ptr<Map> perf_event_map_;
 
-private:
+protected:
+  virtual std::set<std::string> find_wildcard_matches(std::string attach_point, std::string file_name);
   std::vector<Probe> probes_;
   std::vector<Probe> special_probes_;
+
+private:
   std::vector<std::unique_ptr<AttachedProbe>> attached_probes_;
   std::vector<std::unique_ptr<AttachedProbe>> special_attached_probes_;
   KSyms ksyms;

src/lexer.l

@@ -18,7 +18,7 @@ using namespace bpftrace;
 %}
 
 ident  [_a-zA-Z][_a-zA-Z0-9]*
-map    @[_a-zA-Z0-9]*
+map    @{ident}|@
 var    ${ident}
 int    [0-9]+|0[xX][0-9a-fA-F]+
 hspace [ \t]

src/main.cpp

@@ -95,7 +95,16 @@ int main(int argc, char *argv[])
   act.sa_handler = [](int) { };
   sigaction(SIGINT, &act, NULL);
 
-  std::cout << "Running... press Ctrl-C to stop" << std::endl;
+  int num_probes = bpftrace.num_probes();
+  if (num_probes == 0)
+  {
+    std::cout << "No probes to attach" << std::endl;
+    return 1;
+  }
+  else if (num_probes == 1)
+    std::cout << "Attaching " << bpftrace.num_probes() << " probe..." << std::endl;
+  else
+    std::cout << "Attaching " << bpftrace.num_probes() << " probes..." << std::endl;
 
   err = bpftrace.run();
   if (err)

src/parser.yy

@@ -85,6 +85,7 @@ void yyerror(bpftrace::Driver &driver, const char *s);
 %type <ast::Variable *> var
 %type <ast::ExpressionList *> vargs
 %type <ast::AttachPointList *> attach_points
+%type <std::string> attach_point
 
 %right ASSIGN
 %left LOR
@@ -114,15 +115,23 @@ probe : IDENT pred block                    { $$ = new ast::Probe($1, $2, $3); }
       | IDENT PATH attach_points pred block { $$ = new ast::Probe($1, $2.substr(1, $2.size()-2), $3, $4, $5); }
       ;
 
-attach_points : attach_points "," IDENT { $$ = $1; $1->push_back($3); }
-              | IDENT                   { $$ = new ast::AttachPointList; $$->push_back($1); }
+attach_points : attach_points "," attach_point { $$ = $1; $1->push_back($3); }
+              | attach_point                   { $$ = new ast::AttachPointList; $$->push_back($1); }
               ;
 
+attach_point : attach_point IDENT { $$ = $1 + $2; }
+             | attach_point MUL   { $$ = $1 + "*"; }
+             | IDENT              { $$ = $1; }
+             | MUL                { $$ = "*"; }
+             ;
+
 pred : DIV expr ENDPRED { $$ = new ast::Predicate($2); }
      |                  { $$ = nullptr; }
+     ;
 
 block : "{" stmts "}"     { $$ = $2; }
       | "{" stmts ";" "}" { $$ = $2; }
+      ;
 
 stmts : stmts ";" stmt { $$ = $1; $1->push_back($3); }
       | stmt           { $$ = new ast::StatementList; $$->push_back($1); }
@@ -163,6 +172,7 @@ expr : INT             { $$ = new ast::Integer($1); }
 
 call : IDENT "(" ")"       { $$ = new ast::Call($1); }
      | IDENT "(" vargs ")" { $$ = new ast::Call($1, $3); }
+     ;
 
 map : MAP               { $$ = new ast::Map($1); }
     | MAP "[" vargs "]" { $$ = new ast::Map($1, $3); }

src/types.h

@@ -58,6 +58,7 @@ public:
   ProbeType type;
   std::string path;
   std::string attach_point;
+  std::string prog_name;
   std::string name;
 };
 

tests/CMakeLists.txt

@@ -6,6 +6,8 @@ include_directories(${CMAKE_SOURCE_DIR}/src)
 include_directories(${CMAKE_BINARY_DIR}/src)
 
 add_executable(bpftrace_test
+  ast.cpp
+  bpftrace.cpp
   main.cpp
   parser.cpp
   semantic_analyser.cpp

tests/ast.cpp

+#include "gtest/gtest.h"
+#include "ast.h"
+
+namespace bpftrace {
+namespace test {
+namespace ast {
+
+using bpftrace::ast::AttachPointList;
+using bpftrace::ast::Probe;
+
+TEST(ast, probe_name_special)
+{
+  Probe begin("BEGIN", nullptr, nullptr);
+  EXPECT_EQ(begin.name(), "BEGIN");
+
+  Probe end("END", nullptr, nullptr);
+  EXPECT_EQ(end.name(), "END");
+}
+
+TEST(ast, probe_name_kprobe)
+{
+  AttachPointList attach_points1 = { "sys_read" };
+  Probe kprobe1("kprobe", &attach_points1, nullptr, nullptr);
+  EXPECT_EQ(kprobe1.name(), "kprobe:sys_read");
+
+  AttachPointList attach_points2 = { "sys_read", "sys_write" };
+  Probe kprobe2("kprobe", &attach_points2, nullptr, nullptr);
+  EXPECT_EQ(kprobe2.name(), "kprobe:sys_read,sys_write");
+}
+
+TEST(ast, probe_name_uprobe)
+{
+  AttachPointList attach_points1 = { "readline" };
+  Probe uprobe1("uprobe", "/bin/sh", &attach_points1, nullptr, nullptr);
+  EXPECT_EQ(uprobe1.name(), "uprobe:/bin/sh:readline");
+
+  AttachPointList attach_points2 = { "readline", "somefunc" };
+  Probe uprobe2("uprobe", "/bin/sh", &attach_points2, nullptr, nullptr);
+  EXPECT_EQ(uprobe2.name(), "uprobe:/bin/sh:readline,somefunc");
+}
+
+TEST(ast, probe_name_singular)
+{
+  AttachPointList attach_points = { "sys_read", "sys_thisone", "readline" };
+  Probe kprobe("kprobe", &attach_points, nullptr, nullptr);
+  EXPECT_EQ(kprobe.name("sys_thisone"), "kprobe:sys_thisone");
+
+  Probe uprobe("uprobe", "/bin/sh", &attach_points, nullptr, nullptr);
+  EXPECT_EQ(uprobe.name("readline"), "uprobe:/bin/sh:readline");
+}
+
+} // namespace ast
+} // namespace test
+} // namespace bpftrace

tests/bpftrace.cpp

+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include "bpftrace.h"
+
+namespace bpftrace {
+namespace test {
+namespace bpftrace {
+
+class MockBPFtrace : public BPFtrace {
+public:
+  MOCK_METHOD2(find_wildcard_matches, std::set<std::string>(std::string attach_point, std::string file));
+  std::vector<Probe> get_probes()
+  {
+    return probes_;
+  }
+  std::vector<Probe> get_special_probes()
+  {
+    return special_probes_;
+  }
+};
+
+using ::testing::_;
+using ::testing::Return;
+using ::testing::StrictMock;
+
+void check_kprobe(Probe &p, const std::string &attach_point, const std::string &prog_name)
+{
+  EXPECT_EQ(p.type, ProbeType::kprobe);
+  EXPECT_EQ(p.attach_point, attach_point);
+  EXPECT_EQ(p.prog_name, prog_name);
+  EXPECT_EQ(p.name, "kprobe:" + attach_point);
+}
+
+void check_uprobe(Probe &p, const std::string &path, const std::string &attach_point, const std::string &prog_name)
+{
+  EXPECT_EQ(p.type, ProbeType::uprobe);
+  EXPECT_EQ(p.attach_point, attach_point);
+  EXPECT_EQ(p.prog_name, prog_name);
+  EXPECT_EQ(p.name, "uprobe:" + path + ":" + attach_point);
+}
+
+void check_special_probe(Probe &p, const std::string &attach_point, const std::string &prog_name)
+{
+  EXPECT_EQ(p.type, ProbeType::uprobe);
+  EXPECT_EQ(p.attach_point, attach_point);
+  EXPECT_EQ(p.prog_name, prog_name);
+  EXPECT_EQ(p.name, prog_name);
+}
+
+TEST(bpftrace, add_begin_probe)
+{
+  ast::Probe probe("BEGIN", nullptr, nullptr);
+
+  StrictMock<MockBPFtrace> bpftrace;
+  EXPECT_EQ(bpftrace.add_probe(probe), 0);
+  EXPECT_EQ(bpftrace.get_probes().size(), 0);
+  EXPECT_EQ(bpftrace.get_special_probes().size(), 1);
+
+  check_special_probe(bpftrace.get_special_probes().at(0), "BEGIN_trigger", "BEGIN");
+}
+
+TEST(bpftrace, add_end_probe)
+{
+  ast::Probe probe("END", nullptr, nullptr);
+
+  StrictMock<MockBPFtrace> bpftrace;
+  EXPECT_EQ(bpftrace.add_probe(probe), 0);
+  EXPECT_EQ(bpftrace.get_probes().size(), 0);
+  EXPECT_EQ(bpftrace.get_special_probes().size(), 1);
+
+  check_special_probe(bpftrace.get_special_probes().at(0), "END_trigger", "END");
+}
+
+TEST(bpftrace, add_probes_single)
+{
+  ast::AttachPointList attach_points = {"sys_read"};
+  ast::Probe probe("kprobe", &attach_points, nullptr, nullptr);
+
+  StrictMock<MockBPFtrace> bpftrace;
+  EXPECT_EQ(bpftrace.add_probe(probe), 0);
+  EXPECT_EQ(bpftrace.get_probes().size(), 1);
+  EXPECT_EQ(bpftrace.get_special_probes().size(), 0);
+
+  check_kprobe(bpftrace.get_probes().at(0), "sys_read", "kprobe:sys_read");
+}
+
+TEST(bpftrace, add_probes_multiple)
+{
+  ast::AttachPointList attach_points = {"sys_read", "sys_write"};
+  ast::Probe probe("kprobe", &attach_points, nullptr, nullptr);
+
+  StrictMock<MockBPFtrace> bpftrace;
+  EXPECT_EQ(bpftrace.add_probe(probe), 0);
+  EXPECT_EQ(bpftrace.get_probes().size(), 2);
+  EXPECT_EQ(bpftrace.get_special_probes().size(), 0);
+
+  std::string probe_prog_name = "kprobe:sys_read,sys_write";
+  check_kprobe(bpftrace.get_probes().at(0), "sys_read", probe_prog_name);
+  check_kprobe(bpftrace.get_probes().at(1), "sys_write", probe_prog_name);
+}
+
+TEST(bpftrace, add_probes_wildcard)
+{
+  ast::AttachPointList attach_points = {"sys_read", "my_*", "sys_write"};
+  ast::Probe probe("kprobe", &attach_points, nullptr, nullptr);
+
+  StrictMock<MockBPFtrace> bpftrace;
+  std::set<std::string> matches = { "my_one", "my_two" };
+  ON_CALL(bpftrace, find_wildcard_matches(_, _))
+    .WillByDefault(Return(matches));
+  EXPECT_CALL(bpftrace,
+      find_wildcard_matches("my_*",
+        "/sys/kernel/debug/tracing/available_filter_functions"))
+    .Times(1);
+
+  EXPECT_EQ(bpftrace.add_probe(probe), 0);
+  EXPECT_EQ(bpftrace.get_probes().size(), 4);
+  EXPECT_EQ(bpftrace.get_special_probes().size(), 0);
+
+  std::string probe_prog_name = "kprobe:sys_read,my_*,sys_write";
+  check_kprobe(bpftrace.get_probes().at(0), "sys_read", probe_prog_name);
+  check_kprobe(bpftrace.get_probes().at(1), "my_one", probe_prog_name);
+  check_kprobe(bpftrace.get_probes().at(2), "my_two", probe_prog_name);
+  check_kprobe(bpftrace.get_probes().at(3), "sys_write", probe_prog_name);
+}
+
+TEST(bpftrace, add_probes_wildcard_no_matches)
+{
+  ast::AttachPointList attach_points = {"sys_read", "my_*", "sys_write"};
+  ast::Probe probe("kprobe", &attach_points, nullptr, nullptr);
+
+  StrictMock<MockBPFtrace> bpftrace;
+  std::set<std::string> matches;
+  ON_CALL(bpftrace, find_wildcard_matches(_, _))
+    .WillByDefault(Return(matches));
+  EXPECT_CALL(bpftrace,
+      find_wildcard_matches("my_*",
+        "/sys/kernel/debug/tracing/available_filter_functions"))
+    .Times(1);
+
+  EXPECT_EQ(bpftrace.add_probe(probe), 0);
+  EXPECT_EQ(bpftrace.get_probes().size(), 2);
+  EXPECT_EQ(bpftrace.get_special_probes().size(), 0);
+
+  std::string probe_prog_name = "kprobe:sys_read,my_*,sys_write";
+  check_kprobe(bpftrace.get_probes().at(0), "sys_read", probe_prog_name);
+  check_kprobe(bpftrace.get_probes().at(1), "sys_write", probe_prog_name);
+}
+
+TEST(bpftrace, add_probes_uprobe)
+{
+  ast::AttachPointList attach_points = {"foo"};
+  ast::Probe probe("uprobe", "/bin/sh", &attach_points, nullptr, nullptr);
+
+  StrictMock<MockBPFtrace> bpftrace;
+
+  EXPECT_EQ(bpftrace.add_probe(probe), 0);
+  EXPECT_EQ(bpftrace.get_probes().size(), 1);
+  EXPECT_EQ(bpftrace.get_special_probes().size(), 0);
+  check_uprobe(bpftrace.get_probes().at(0), "/bin/sh", "foo", "uprobe:/bin/sh:foo");
+}
+
+TEST(bpftrace, add_probes_uprobe_wildcard)
+{
+  ast::AttachPointList attach_points = {"foo*"};
+  ast::Probe probe("uprobe", "/bin/sh", &attach_points, nullptr, nullptr);
+
+  StrictMock<MockBPFtrace> bpftrace;
+
+  EXPECT_NE(bpftrace.add_probe(probe), 0);
+  EXPECT_EQ(bpftrace.get_probes().size(), 0);
+  EXPECT_EQ(bpftrace.get_special_probes().size(), 0);
+}
+
+} // namespace bpftrace
+} // namespace test
+} // namespace bpftrace

tests/parser.cpp

@@ -228,6 +228,51 @@ TEST(Parser, multiple_attach_points)
       "  int: 1\n");
 }
 
+TEST(Parser, wildcard_attach_points)
+{
+  test("kprobe:sys_* { 1 }",
+      "Program\n"
+      " kprobe:sys_*\n"
+      "  int: 1\n");
+  test("kprobe:*blah { 1 }",
+      "Program\n"
+      " kprobe:*blah\n"
+      "  int: 1\n");
+  test("kprobe:sys*blah { 1 }",
+      "Program\n"
+      " kprobe:sys*blah\n"
+      "  int: 1\n");
+  test("kprobe:* { 1 }",
+      "Program\n"
+      " kprobe:*\n"
+      "  int: 1\n");
+  test("kprobe:sys_* { @x = y*z }",
+      "Program\n"
+      " kprobe:sys_*\n"
+      "  =\n"
+      "   map: @x\n"
+      "   *\n"
+      "    builtin: y\n"
+      "    builtin: z\n");
+  test("kprobe:sys_* { @x = *arg0 }",
+      "Program\n"
+      " kprobe:sys_*\n"
+      "  =\n"
+      "   map: @x\n"
+      "   dereference\n"
+      "    builtin: arg0\n");
+}
+
+TEST(Parser, short_map_name)
+{
+  test("kprobe:sys_read { @ = 1 }",
+      "Program\n"
+      " kprobe:sys_read\n"
+      "  =\n"
+      "   map: @\n"
+      "   int: 1\n");
+}
+
 } // namespace parser
 } // namespace test
 } // namespace bpftrace