Refactor find_wildcard_matches() to allow for proper testing

Some tests weren't running correctly (add_probes_character_class)
and some were system-dependent (e.g. add_probes_uprobe_wildcard).

This also makes MockBPFtrace available for use from all test cases.

src/ast/codegen_llvm.cpp

@@ -1298,40 +1298,11 @@ void CodegenLLVM::visit(Probe &probe)
     int starting_time_id_ = time_id_;
     int starting_join_id_ = join_id_;
 
-    for (auto &attach_point : *probe.attach_points) {
+    for (auto attach_point : *probe.attach_points) {
       current_attach_point_ = attach_point;
-      std::set<std::string> matches;
-      switch (probetype(attach_point->provider))
-      {
-        case ProbeType::kprobe:
-        case ProbeType::kretprobe:
-          matches = bpftrace_.find_wildcard_matches(attach_point->target,
-                                                    attach_point->func,
-                                                    "/sys/kernel/debug/tracing/available_filter_functions");
-          break;
-        case ProbeType::uprobe:
-        case ProbeType::uretprobe:
-        {
-          auto symbol_stream = std::istringstream(bpftrace_.extract_func_symbols_from_path(attach_point->target));
-          matches = bpftrace_.find_wildcard_matches("", attach_point->func, symbol_stream);
-          break;
-        }
-        case ProbeType::tracepoint:
-          matches = bpftrace_.find_wildcard_matches(attach_point->target,
-                                                    attach_point->func,
-                                                    "/sys/kernel/debug/tracing/available_events");
-          break;
-        case ProbeType::usdt:
-        {
-          auto usdt_symbol_stream = USDTHelper::probe_stream(bpftrace_.pid_, attach_point->target);
-          matches = bpftrace_.find_usdt_wildcard_matches(attach_point->ns, attach_point->func, usdt_symbol_stream);
-          break;
-        }
-        default:
-          std::cerr << "Wildcard matches aren't available on probe type '"
-                    << attach_point->provider << "'" << std::endl;
-          return;
-      }
+
+      auto matches = bpftrace_.find_wildcard_matches(*attach_point);
+
       tracepoint_struct_ = "";
       for (auto &match_ : matches) {
         printf_id_ = starting_printf_id_;

src/ast/semantic_analyser.cpp

@@ -168,12 +168,14 @@ void SemanticAnalyser::visit(Builtin &builtin)
        * 2. sets is_tparg so that codegen does the real type setting after
        *    expansion.
        */
-      std::set<std::string> matches;
-      matches = bpftrace_.find_wildcard_matches(attach_point->target,
-                                                attach_point->func,
-                                                "/sys/kernel/debug/tracing/available_events");
+      auto symbol_stream = bpftrace_.get_symbols_from_file(
+          "/sys/kernel/debug/tracing/available_events");
+      auto matches = bpftrace_.find_wildcard_matches(attach_point->target,
+                                                     attach_point->func,
+                                                     *symbol_stream);
       for (auto &match : matches) {
-        std::string tracepoint_struct = TracepointFormatParser::get_struct_name(attach_point->target, match);
+        std::string tracepoint_struct = TracepointFormatParser::get_struct_name(
+            attach_point->target, match);
         Struct &cstruct = bpftrace_.structs_[tracepoint_struct];
         builtin.type = SizedType(Type::cast, cstruct.size, tracepoint_struct);
         builtin.type.is_pointer = true;
@@ -786,9 +788,11 @@ void SemanticAnalyser::visit(FieldAccess &acc)
     for (AttachPoint *attach_point : *probe_->attach_points) {
       assert(probetype(attach_point->provider) == ProbeType::tracepoint);
 
-      std::set<std::string> matches = bpftrace_.find_wildcard_matches(
-          attach_point->target, attach_point->func,
+      auto symbol_stream = bpftrace_.get_symbols_from_file(
           "/sys/kernel/debug/tracing/available_events");
+      auto matches = bpftrace_.find_wildcard_matches(attach_point->target,
+                                                     attach_point->func,
+                                                     *symbol_stream);
       for (auto &match : matches) {
         std::string tracepoint_struct =
             TracepointFormatParser::get_struct_name(attach_point->target,

src/bpftrace.cpp

@@ -98,38 +98,15 @@ int BPFtrace::add_probe(ast::Probe &p)
     if (attach_point->need_expansion && (has_wildcard(attach_point->func) || underspecified_usdt_probe))
     {
       std::set<std::string> matches;
-      switch (probetype(attach_point->provider))
+      try
       {
-        case ProbeType::kprobe:
-        case ProbeType::kretprobe:
-          matches = find_wildcard_matches(attach_point->target,
-                                          attach_point->func,
-                                          "/sys/kernel/debug/tracing/available_filter_functions");
-          break;
-        case ProbeType::uprobe:
-        case ProbeType::uretprobe:
-        {
-            auto symbol_stream = std::istringstream(extract_func_symbols_from_path(attach_point->target));
-            matches = find_wildcard_matches("", attach_point->func, symbol_stream);
-            break;
-        }
-        case ProbeType::tracepoint:
-          matches = find_wildcard_matches(attach_point->target,
-                                          attach_point->func,
-                                          "/sys/kernel/debug/tracing/available_events");
-          break;
-        case ProbeType::usdt:
-        {
-          auto usdt_symbol_stream = USDTHelper::probe_stream(pid_, attach_point->target);
-          matches = find_usdt_wildcard_matches(attach_point->ns, attach_point->func, usdt_symbol_stream);
-          break;
-        }
-        default:
-          std::cerr << "Wildcard matches aren't available on probe type '"
-                    << attach_point->provider << "'" << std::endl;
-          return 1;
+        matches = find_wildcard_matches(*attach_point);
+      }
+      catch (const WildcardException &e)
+      {
+        std::cerr << e.what() << std::endl;
+        return 1;
       }
-
       attach_funcs.insert(attach_funcs.end(), matches.begin(), matches.end());
     }
     else
@@ -175,39 +152,71 @@ int BPFtrace::add_probe(ast::Probe &p)
   return 0;
 }
 
-// FIXME should this really be a separate function?
-std::set<std::string> BPFtrace::find_usdt_wildcard_matches(const std::string &prefix, const std::string &func, std::istream &symbol_name_stream)
+std::set<std::string> BPFtrace::find_wildcard_matches(
+    const ast::AttachPoint &attach_point) const
 {
-  // Turn glob into a regex
-  std::string search_str = func;
-  if (prefix == "")
-    search_str = "*:" + func;
-  else
-    search_str = prefix + ":" + func;
-  auto regex_str = "(" + std::regex_replace(search_str, std::regex("\\*"), "[^\\s]*") + ")";
-  regex_str = "^" + regex_str + "$";
-
-  std::regex func_regex(regex_str);
-  std::smatch match;
+  std::unique_ptr<std::istream> symbol_stream;
+  std::string prefix, func;
 
-  std::string line;
-  std::set<std::string> matches;
-  while (std::getline(symbol_name_stream, line))
+  switch (probetype(attach_point.provider))
   {
-    if (std::regex_search(line, match, func_regex))
+    case ProbeType::kprobe:
+    case ProbeType::kretprobe:
     {
-      assert(match.size() == 2);
-      // skip the ".part.N" kprobe variants, as they can't be traced:
-      if (std::strstr(match.str(1).c_str(), ".part.") == NULL)
-      {
-        matches.insert(match[1]);
-      }
+      symbol_stream = get_symbols_from_file(
+          "/sys/kernel/debug/tracing/available_filter_functions");
+      prefix = "";
+      func = attach_point.func;
+      break;
+    }
+    case ProbeType::uprobe:
+    case ProbeType::uretprobe:
+    {
+      symbol_stream = std::make_unique<std::istringstream>(
+          extract_func_symbols_from_path(attach_point.target));
+      prefix = "";
+      func = attach_point.func;
+      break;
+    }
+    case ProbeType::tracepoint:
+    {
+      symbol_stream = get_symbols_from_file(
+          "/sys/kernel/debug/tracing/available_events");
+      prefix = attach_point.target;
+      func = attach_point.func;
+      break;
+    }
+    case ProbeType::usdt:
+    {
+      symbol_stream = get_symbols_from_usdt(pid_, attach_point.target);
+      prefix = "";
+      if (attach_point.ns == "")
+        func = "*:" + attach_point.func;
+      else
+        func = attach_point.ns + ":" + attach_point.func;
+      break;
+    }
+    default:
+    {
+      throw WildcardException("Wildcard matches aren't available on probe type '"
+          + attach_point.provider + "'");
     }
   }
-  return matches;
+
+  return find_wildcard_matches(prefix, func, *symbol_stream);
 }
 
-std::set<std::string> BPFtrace::find_wildcard_matches(const std::string &prefix, const std::string &func, std::istream &symbol_name_stream)
+/*
+ * Finds all matches of func in the provided input stream.
+ *
+ * If an optional prefix is provided, lines must start with it to count as a
+ * match, but the prefix is stripped from entries in the result set.
+ * Wildcard tokens ("*") are accepted in func.
+ */
+std::set<std::string> BPFtrace::find_wildcard_matches(
+    const std::string &prefix,
+    const std::string &func,
+    std::istream &symbol_stream) const
 {
   if (!has_wildcard(func))
     return std::set<std::string>({func});
@@ -220,7 +229,7 @@ std::set<std::string> BPFtrace::find_wildcard_matches(const std::string &prefix,
   std::string line;
   std::set<std::string> matches;
   std::string full_prefix = prefix.empty() ? "" : (prefix + ":");
-  while (std::getline(symbol_name_stream, line))
+  while (std::getline(symbol_stream, line))
   {
     if (!full_prefix.empty()) {
       if (line.find(full_prefix, 0) != 0)
@@ -240,26 +249,39 @@ std::set<std::string> BPFtrace::find_wildcard_matches(const std::string &prefix,
   return matches;
 }
 
-std::set<std::string> BPFtrace::find_wildcard_matches(const std::string &prefix, const std::string &func, const std::string &file_name)
+std::unique_ptr<std::istream> BPFtrace::get_symbols_from_file(const std::string &path) const
 {
-  if (!has_wildcard(func))
-    return std::set<std::string>({func});
-  std::ifstream file(file_name);
-  if (file.fail())
+  auto file = std::make_unique<std::ifstream>(path);
+  if (file->fail())
   {
-    throw std::runtime_error("Could not read symbols from \"" + file_name + "\", err=" + std::to_string(errno));
+    throw std::runtime_error("Could not read symbols from " + path +
+                             ": " + strerror(errno));
   }
 
-  std::stringstream symbol_name_stream;
-  std::string line;
-  while (file >> line)
+  return file;
+}
+
+std::unique_ptr<std::istream> BPFtrace::get_symbols_from_usdt(
+    int pid,
+    const std::string &target) const
+{
+  std::string probes;
+  usdt_probe_list usdt_probes;
+
+  if (pid > 0)
+    usdt_probes = USDTHelper::probes_for_pid(pid);
+  else
+    usdt_probes = USDTHelper::probes_for_path(target);
+
+  for (auto const& usdt_probe : usdt_probes)
   {
-    symbol_name_stream << line << std::endl;
+    std::string path     = std::get<USDT_PATH_INDEX>(usdt_probe);
+    std::string provider = std::get<USDT_PROVIDER_INDEX>(usdt_probe);
+    std::string fname    = std::get<USDT_FNAME_INDEX>(usdt_probe);
+    probes += provider + ":" + fname + "\n";
   }
 
-  file.close();
-
-  return find_wildcard_matches(prefix, func, symbol_name_stream);
+  return std::make_unique<std::istringstream>(probes);
 }
 
 int BPFtrace::num_probes() const

src/bpftrace.h

@@ -48,6 +48,20 @@ inline DebugLevel operator++(DebugLevel& level, int)
   return level;
 }
 
+class WildcardException : public std::exception
+{
+public:
+  WildcardException(const std::string &msg) : msg_(msg) {}
+
+  const char *what() const noexcept override
+  {
+    return msg_.c_str();
+  }
+
+private:
+  std::string msg_;
+};
+
 class BPFtrace
 {
 public:
@@ -70,7 +84,7 @@ class BPFtrace
   std::string resolve_uid(uintptr_t addr) const;
   uint64_t resolve_kname(const std::string &name) const;
   uint64_t resolve_uname(const std::string &name, const std::string &path) const;
-  std::string extract_func_symbols_from_path(const std::string &path) const;
+  virtual std::string extract_func_symbols_from_path(const std::string &path) const;
   std::string resolve_probe(uint64_t probe_id) const;
   uint64_t resolve_cgroupid(const std::string &path) const;
   std::vector<std::unique_ptr<IPrintable>> get_arg_values(const std::vector<Field> &args, uint8_t* arg_data);
@@ -101,11 +115,20 @@ class BPFtrace
   bool demangle_cpp_symbols = true;
   bool safe_mode = true;
 
-  static void sort_by_key(std::vector<SizedType> key_args,
-      std::vector<std::pair<std::vector<uint8_t>, std::vector<uint8_t>>> &values_by_key);
-  virtual std::set<std::string> find_usdt_wildcard_matches(const std::string &prefix, const std::string &func, std::istream &symbol_name_stream);
-  virtual std::set<std::string> find_wildcard_matches(const std::string &prefix, const std::string &func, std::istream &symbol_name_stream);
-  virtual std::set<std::string> find_wildcard_matches(const std::string &prefix, const std::string &attach_point, const std::string &file_name);
+  static void sort_by_key(
+      std::vector<SizedType> key_args,
+      std::vector<std::pair<std::vector<uint8_t>,
+      std::vector<uint8_t>>> &values_by_key);
+  std::set<std::string> find_wildcard_matches(
+      const ast::AttachPoint &attach_point) const;
+  std::set<std::string> find_wildcard_matches(
+      const std::string &prefix,
+      const std::string &func,
+      std::istream &symbol_stream) const;
+  virtual std::unique_ptr<std::istream> get_symbols_from_file(const std::string &path) const;
+  virtual std::unique_ptr<std::istream> get_symbols_from_usdt(
+      int pid,
+      const std::string &target) const;
 
 protected:
   std::vector<Probe> probes_;