From 78288d961ed0fac526404b7d3f198514beafe7d7 Mon Sep 17 00:00:00 2001
From: Egor Tensin <Egor.Tensin@gmail.com>
Date: Wed, 15 Jan 2020 12:45:39 +0300
Subject: enum_symbols: support symbol mask using --mask

---
 .appveyor.yml           |  3 ++-
 include/pdb/dbghelp.hpp |  3 +++
 src/dbghelp.cpp         | 34 ++++++++++++++++++++++++----------
 utils/enum_symbols.cpp  |  9 +++++++--
 4 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/.appveyor.yml b/.appveyor.yml
index 5eb6562..53958e2 100644
--- a/.appveyor.yml
+++ b/.appveyor.yml
@@ -37,7 +37,8 @@ after_build:
   - appveyor.exe PushArtifact "%APPVEYOR_PROJECT_NAME%-%PLATFORM%-%CONFIGURATION%.zip"
 
 test_script:
-  - '"%install_dir%\bin\enum_symbols" --pdb "%install_dir%\bin\lab_rat.pdb" --functions'
+  - '"%install_dir%\bin\enum_symbols" --pdb "%install_dir%\bin\lab_rat.pdb" --functions --mask "lab_rat!*"'
+  - '"%install_dir%\bin\enum_symbols" --pdb "%install_dir%\bin\lab_rat.pdb" --functions --mask "lab_rat*"'
 
 for:
 # Only build Release builds on master to speed things up:
diff --git a/include/pdb/dbghelp.hpp b/include/pdb/dbghelp.hpp
index edb095c..2809eac 100644
--- a/include/pdb/dbghelp.hpp
+++ b/include/pdb/dbghelp.hpp
@@ -29,7 +29,10 @@ public:
     void enum_modules(const OnModule&) const;
 
     typedef std::function<void(const SymbolInfo&)> OnSymbol;
+    static constexpr auto all_symbols = "*!*";
+    void enum_symbols(const ModuleInfo&, const std::string& mask, const OnSymbol&) const;
     void enum_symbols(const ModuleInfo&, const OnSymbol&) const;
+    void enum_symbols(const std::string& mask, const OnSymbol&) const;
     void enum_symbols(const OnSymbol&) const;
 
     SymbolInfo resolve_symbol(Address) const;
diff --git a/src/dbghelp.cpp b/src/dbghelp.cpp
index 8268673..399271e 100644
--- a/src/dbghelp.cpp
+++ b/src/dbghelp.cpp
@@ -73,6 +73,18 @@ BOOL CALLBACK enum_symbols_callback(SYMBOL_INFO* info, ULONG, VOID* raw_callback
     return TRUE;
 }
 
+void enum_symbols(HANDLE id,
+                  Address module_base,
+                  const std::string& mask,
+                  const DbgHelp::OnSymbol& callback) {
+    if (!SymEnumSymbols(id,
+                        module_base,
+                        mask.c_str(),
+                        &enum_symbols_callback,
+                        const_cast<DbgHelp::OnSymbol*>(&callback)))
+        throw error::windows(GetLastError());
+}
+
 } // namespace
 
 DbgHelp::DbgHelp() {
@@ -120,20 +132,22 @@ void DbgHelp::enum_modules(const OnModule& callback) const {
         throw error::windows(GetLastError());
 }
 
+void DbgHelp::enum_symbols(const ModuleInfo& module,
+                           const std::string& mask,
+                           const OnSymbol& callback) const {
+    pdb::enum_symbols(id, module.get_offline_base(), mask, callback);
+}
+
 void DbgHelp::enum_symbols(const ModuleInfo& module, const OnSymbol& callback) const {
-    if (!SymEnumSymbols(id,
-                        module.get_offline_base(),
-                        NULL,
-                        &enum_symbols_callback,
-                        const_cast<OnSymbol*>(&callback)))
-        throw error::windows(GetLastError());
+    enum_symbols(module, all_symbols, callback);
+}
+
+void DbgHelp::enum_symbols(const std::string& mask, const OnSymbol& callback) const {
+    pdb::enum_symbols(id, 0, mask, callback);
 }
 
 void DbgHelp::enum_symbols(const OnSymbol& callback) const {
-    static constexpr auto all_symbols = "*!*";
-    if (!SymEnumSymbols(
-            id, 0, all_symbols, &enum_symbols_callback, const_cast<OnSymbol*>(&callback)))
-        throw error::windows(GetLastError());
+    enum_symbols(all_symbols, callback);
 }
 
 SymbolInfo DbgHelp::resolve_symbol(Address offline) const {
diff --git a/utils/enum_symbols.cpp b/utils/enum_symbols.cpp
index 87cc56a..d7f4c37 100644
--- a/utils/enum_symbols.cpp
+++ b/utils/enum_symbols.cpp
@@ -27,10 +27,12 @@ public:
             "functions",
             po::value<pdb::symbol::Tag>(&tag)->implicit_value(function_tag)->zero_tokens(),
             "only list functions");
+        visible.add_options()(
+            "mask", po::value<std::string>(&symbol_mask)->value_name("MASK"), "symbol mask");
     }
 
     const char* get_short_description() const override {
-        return "[-h|--help] [--pdb PATH]... [--functions]";
+        return "[-h|--help] [--pdb PATH]... [--functions] [--mask MASK]";
     }
 
     std::vector<std::string> pdbs;
@@ -39,11 +41,14 @@ public:
 
     pdb::Symbol::Type get_type() const { return static_cast<pdb::Symbol::Type>(tag); }
 
+    std::string get_mask() const { return symbol_mask; }
+
 private:
     static constexpr auto reserved_tag = static_cast<pdb::symbol::Tag>(pdb::Symbol::Type::RESERVED);
     static constexpr auto function_tag = static_cast<pdb::symbol::Tag>(pdb::Symbol::Type::Function);
 
     pdb::symbol::Tag tag = reserved_tag;
+    std::string symbol_mask{pdb::DbgHelp::all_symbols};
 };
 
 constexpr pdb::symbol::Tag EnumSymbols::reserved_tag;
@@ -72,7 +77,7 @@ int main(int argc, char* argv[]) {
         for (const auto& pdb : settings.pdbs) {
             const auto id = dbghelp.load_pdb(pdb);
 
-            dbghelp.enum_symbols(id, [&](const pdb::SymbolInfo& symbol) {
+            dbghelp.enum_symbols(id, settings.get_mask(), [&](const pdb::SymbolInfo& symbol) {
                 if (!settings.type_specified() || settings.get_type() == symbol.get_type())
                     std::cout << symbol.get_name() << '\n';
             });
-- 
cgit v1.2.3