optimin: run showmap with an entire input directory

Rather than executing with one seed at a time. This required disabling
the target binary test (for now).

1 files changed, 179 insertions, 116 deletions

diff --git a/utils/optimin/src/OptiMin.cpp b/utils/optimin/src/OptiMin.cpp
index 276f6b0f..242eeea9 100644
--- a/utils/optimin/src/OptiMin.cpp
+++ b/utils/optimin/src/OptiMin.cpp
@@ -64,7 +64,10 @@ using AFLTuple = std::pair<AFLTupleID, /* Frequency */ unsigned>;
 /// Coverage for a given seed file
 using AFLCoverageVector = std::vector<AFLTuple>;
 
-/// Maps seed file paths to a weight
+/// Map seed file paths to its coverage vector
+using AFLCoverageMap = StringMap<AFLCoverageVector>;
+
+/// Map seed file paths to a weight
 using WeightsMap = StringMap<Weight>;
 
 /// A seed identifier in the MaxSAT solver
@@ -88,6 +91,9 @@ using MaxSATCoverageMap = DenseMap<AFLTupleID, MaxSATSeedSet>;
 // `afl-showmap.c`
 static constexpr uint32_t MAX_EDGE_FREQ = 8;
 
+// The maximum number of failures allowed when parsing a weights file
+static constexpr unsigned MAX_WEIGHT_FAILURES = 5;
+
 static sys::TimePoint<>     StartTime, EndTime;
 static std::chrono::seconds Duration;
 
@@ -98,30 +104,30 @@ static bool        SkipBinCheck = false;
 
 static const auto ErrMsg = [] {
 
-  return WithColor(errs(), HighlightColor::Error) << "[-] ";
+  return WithColor(errs(), raw_ostream::RED, /*Bold=*/true) << "[-] ";
 
 };
 
 static const auto WarnMsg = [] {
 
-  return WithColor(errs(), HighlightColor::Warning) << "[-] ";
+  return WithColor(errs(), raw_ostream::MAGENTA, /*Bold=*/true) << "[-] ";
 
 };
 
 static const auto SuccMsg = [] {
 
-  return WithColor(outs(), HighlightColor::String) << "[+] ";
+  return WithColor(outs(), raw_ostream::GREEN, /*Bold=*/true) << "[+] ";
 
 };
 
 static const auto StatMsg = [] {
 
-  return WithColor(outs(), HighlightColor::Remark) << "[*] ";
+  return WithColor(outs(), raw_ostream::BLUE, /*Bold=*/true) << "[*] ";
 
 };
 
-static cl::opt<std::string> CorpusDir("i", cl::desc("Input directory"),
-                                      cl::value_desc("dir"), cl::Required);
+static cl::opt<std::string> InputDir("i", cl::desc("Input directory"),
+                                     cl::value_desc("dir"), cl::Required);
 static cl::opt<std::string> OutputDir("o", cl::desc("Output directory"),
                                       cl::value_desc("dir"), cl::Required);
 
@@ -164,6 +170,7 @@ static void GetWeights(const MemoryBuffer &MB, WeightsMap &Weights) {
   SmallVector<StringRef, 0> Lines;
   MB.getBuffer().trim().split(Lines, '\n');
 
+  unsigned FailureCount = 0;
   unsigned Weight = 0;
 
   for (const auto &Line : Lines) {
@@ -176,7 +183,13 @@ static void GetWeights(const MemoryBuffer &MB, WeightsMap &Weights) {
 
     } else {
 
-      WarnMsg() << "Failed to read weight for `" << Seed << "`. Skipping...\n";
+      if (FailureCount >= MAX_WEIGHT_FAILURES) {
+        ErrMsg() << "Too many failures. Aborting\n";
+        std::exit(1);
+      }
+
+      WarnMsg() << "Failed to read weight for '" << Seed << "'. Skipping...\n";
+      FailureCount++;
 
     }
 
@@ -184,23 +197,52 @@ static void GetWeights(const MemoryBuffer &MB, WeightsMap &Weights) {
 
 }
 
-static Error getAFLCoverage(const StringRef Seed, AFLCoverageVector &Cov,
-                            bool BinCheck = false) {
+static std::error_code readCov(const StringRef Trace, AFLCoverageVector &Cov) {
+
+  const auto CovOrErr = MemoryBuffer::getFile(Trace);
+  if (const auto EC = CovOrErr.getError()) return EC;
+
+  SmallVector<StringRef, 0> Lines;
+  CovOrErr.get()->getBuffer().trim().split(Lines, '\n');
+
+  AFLTupleID Edge = 0;
+  unsigned   Freq = 0;
+
+  for (const auto &Line : Lines) {
+
+    const auto &[EdgeStr, FreqStr] = Line.split(':');
+
+    to_integer(EdgeStr, Edge, 10);
+    to_integer(FreqStr, Freq, 10);
+    Cov.push_back({Edge, Freq});
+
+  }
+
+  return std::error_code();
+
+}
+
+static Error runShowmap(AFLCoverageMap &CovMap) {
 
   Optional<StringRef> Redirects[] = {None, None, None};
 
-  SmallString<32> TracePath{OutputDir};
-  StringRef TraceName = BinCheck ? ".run_test" : sys::path::filename(Seed);
-  sys::path::append(TracePath, ".traces", TraceName);
+  SmallString<32> TraceDir{OutputDir};
+  sys::path::append(TraceDir, ".traces");
+
+  SmallString<32> StdinFile{TraceDir};
+  sys::path::append(StdinFile, ".cur_input");
 
   // Prepare afl-showmap arguments
   SmallVector<StringRef, 12> AFLShowmapArgs{
-      AFLShowmapPath, "-m", MemLimit, "-t", Timeout, "-q", "-o", TracePath};
+      AFLShowmapPath, "-m", MemLimit, "-t", Timeout,
+      "-q",           "-i", InputDir, "-o", TraceDir};
 
-  if (TargetArgsHasAtAt)
-    AFLShowmapArgs.append({"-A", Seed});
-  else
-    Redirects[/* stdin */ 0] = Seed;
+  if (TargetArgsHasAtAt) {
+
+    AFLShowmapArgs.append({"-A", StdinFile});
+    Redirects[/* stdin */ 0] = "/dev/null";
+
+  }
 
   if (FridaMode) AFLShowmapArgs.push_back("-O");
   if (QemuMode) AFLShowmapArgs.push_back("-Q");
@@ -214,39 +256,48 @@ static Error getAFLCoverage(const StringRef Seed, AFLCoverageVector &Cov,
                                      /*env=*/None, Redirects);
   if (RC && !CrashMode) {
 
-    ErrMsg() << "Exit code " << RC << " != 0 received from afl-showmap";
-    sys::fs::remove(TracePath);
+    ErrMsg() << "Exit code " << RC << " != 0 received from afl-showmap\n";
+    sys::fs::remove_directories(TraceDir);
     return createStringError(inconvertibleErrorCode(), "afl-showmap failed");
 
   }
 
   // Parse afl-showmap output
-  const auto CovOrErr = MemoryBuffer::getFile(TracePath);
-  if (const auto EC = CovOrErr.getError()) {
+  AFLCoverageVector    Cov;
+  std::error_code      EC;
+  sys::fs::file_status Status;
 
-    sys::fs::remove(TracePath);
-    return createStringError(EC, "Failed to read afl-showmap output file `%s`",
-                             TracePath.c_str());
+  for (sys::fs::recursive_directory_iterator Dir(TraceDir, EC), DirEnd;
+       Dir != DirEnd && !EC; Dir.increment(EC)) {
 
-  }
+    if (EC) return errorCodeToError(EC);
 
-  SmallVector<StringRef, 0> Lines;
-  CovOrErr.get()->getBuffer().trim().split(Lines, '\n');
+    const auto &Path = Dir->path();
+    if ((EC = sys::fs::status(Path, Status))) return errorCodeToError(EC);
 
-  AFLTupleID Edge = 0;
-  unsigned   Freq = 0;
+    switch (Status.type()) {
 
-  for (const auto &Line : Lines) {
+      case sys::fs::file_type::regular_file:
+      case sys::fs::file_type::symlink_file:
+      case sys::fs::file_type::type_unknown:
+        Cov.clear();
+        if ((EC = readCov(Path, Cov))) {
 
-    const auto &[EdgeStr, FreqStr] = Line.split(':');
+          sys::fs::remove(Path);
+          return errorCodeToError(EC);
 
-    to_integer(EdgeStr, Edge, 10);
-    to_integer(FreqStr, Freq, 10);
-    Cov.push_back({Edge, Freq});
+        }
+
+        CovMap.try_emplace(sys::path::filename(Path), Cov);
+      default:
+        /* Ignore */
+        break;
+
+    }
 
   }
 
-  if (!KeepTraces || BinCheck) sys::fs::remove(TracePath);
+  if (!KeepTraces) sys::fs::remove_directories(TraceDir);
   return Error::success();
 
 }
@@ -326,14 +377,14 @@ int main(int argc, char *argv[]) {
 
   if (WeightsFile != "") {
 
-    StatMsg() << "Reading weights from `" << WeightsFile << "`... ";
+    StatMsg() << "Reading weights from '" << WeightsFile << "'... ";
     StartTimer(/*ShowProgBar=*/false);
 
     const auto WeightsOrErr = MemoryBuffer::getFile(WeightsFile);
     if ((EC = WeightsOrErr.getError())) {
 
-      ErrMsg() << "Failed to read weights from `" << WeightsFile
-               << "`: " << EC.message() << '\n';
+      ErrMsg() << "Failed to read weights from '" << WeightsFile
+               << "': " << EC.message() << '\n';
       return 1;
 
     }
@@ -345,47 +396,33 @@ int main(int argc, char *argv[]) {
   }
 
   // ------------------------------------------------------------------------ //
-  // Setup output directory
+  // Traverse input directory
+  //
+  // Find the seed files inside this directory (and subdirectories).
   // ------------------------------------------------------------------------ //
 
-  SmallString<32> TraceDir{OutputDir};
-  sys::path::append(TraceDir, ".traces");
-
-  if ((EC = sys::fs::remove_directories(TraceDir))) {
-
-    ErrMsg() << "Failed to remove existing trace directory in `" << OutputDir
-             << "`: " << EC.message() << '\n';
-    return 1;
-
-  }
+  StatMsg() << "Locating seeds in '" << InputDir << "'... ";
+  StartTimer(/*ShowProgBar=*/false);
 
-  if ((EC = sys::fs::create_directories(TraceDir))) {
+  bool IsDirResult;
+  if ((EC = sys::fs::is_directory(InputDir, IsDirResult))) {
 
-    ErrMsg() << "Failed to create output directory `" << OutputDir
-             << "`: " << EC.message() << '\n';
+    ErrMsg() << "Invalid input directory '" << InputDir
+             << "': " << EC.message() << '\n';
     return 1;
 
   }
 
-  // ------------------------------------------------------------------------ //
-  // Traverse corpus directory
-  //
-  // Find the seed files inside this directory.
-  // ------------------------------------------------------------------------ //
-
-  StatMsg() << "Locating seeds in `" << CorpusDir << "`... ";
-  StartTimer(/*ShowProgBar=*/false);
-
-  std::vector<std::string> SeedFiles;
-  sys::fs::file_status     Status;
+  sys::fs::file_status Status;
+  StringMap<std::string> SeedFiles;
 
-  for (sys::fs::recursive_directory_iterator Dir(CorpusDir, EC), DirEnd;
+  for (sys::fs::recursive_directory_iterator Dir(InputDir, EC), DirEnd;
        Dir != DirEnd && !EC; Dir.increment(EC)) {
 
     if (EC) {
 
-      ErrMsg() << "Failed to traverse corpus directory `" << CorpusDir
-               << "`: " << EC.message() << '\n';
+      ErrMsg() << "Failed to traverse input directory '" << InputDir
+               << "': " << EC.message() << '\n';
       return 1;
 
     }
@@ -393,21 +430,22 @@ int main(int argc, char *argv[]) {
     const auto &Path = Dir->path();
     if ((EC = sys::fs::status(Path, Status))) {
 
-      WarnMsg() << "Failed to access seed file `" << Path
-                << "`: " << EC.message() << ". Skipping...\n";
-      continue;
+      ErrMsg() << "Failed to access '" << Path << "': " << EC.message()
+               << '\n';
+      return 1;
 
     }
 
     switch (Status.type()) {
 
-      case sys::fs::file_type::regular_file:
-      case sys::fs::file_type::symlink_file:
-      case sys::fs::file_type::type_unknown:
-        SeedFiles.push_back(Path);
-      default:
-        /* Ignore */
-        break;
+    case sys::fs::file_type::regular_file:
+    case sys::fs::file_type::symlink_file:
+    case sys::fs::file_type::type_unknown:
+      SeedFiles.try_emplace(sys::path::filename(Path),
+                            sys::path::parent_path(Path));
+    default:
+      /* Ignore */
+      break;
 
     }
 
@@ -415,24 +453,55 @@ int main(int argc, char *argv[]) {
 
   EndTimer(/*ShowProgBar=*/false);
 
+  if (SeedFiles.empty()) {
+
+    ErrMsg() << "Failed to find any seed files in '" << InputDir << "'\n";
+    return 1;
+
+  }
+
+  // ------------------------------------------------------------------------ //
+  // Setup output directory
+  // ------------------------------------------------------------------------ //
+
+  SmallString<32> TraceDir{OutputDir};
+  sys::path::append(TraceDir, ".traces");
+
+  if ((EC = sys::fs::remove_directories(TraceDir))) {
+
+    ErrMsg() << "Failed to remove existing trace directory in '" << OutputDir
+             << "': " << EC.message() << '\n';
+    return 1;
+
+  }
+
+  if ((EC = sys::fs::create_directories(TraceDir))) {
+
+    ErrMsg() << "Failed to create output directory '" << OutputDir
+             << "': " << EC.message() << '\n';
+    return 1;
+
+  }
+
   // ------------------------------------------------------------------------ //
   // Test the target binary
   // ------------------------------------------------------------------------ //
 
   AFLCoverageVector Cov;
 
-  if (!SkipBinCheck && SeedFiles.size() > 0) {
+  if (!SkipBinCheck) {
 
     StatMsg() << "Testing the target binary... ";
     StartTimer(/*ShowProgBar=*/false);
 
-    if (auto Err = getAFLCoverage(SeedFiles.front(), Cov, /*BinCheck=*/true)) {
-
-      ErrMsg()
-          << "No instrumentation output detected (perhaps crash or timeout)";
-      return 1;
-
-    }
+    //    if (auto Err = runShowmap(TestSeed, Cov, /*BinCheck=*/true)) {
+    //
+    //      ErrMsg()
+    //          << "No instrumentation output detected (perhaps crash or
+    //          timeout)";
+    //      return 1;
+    //
+    //    }
 
     EndTimer(/*ShowProgBar=*/false);
     SuccMsg() << "OK, " << Cov.size() << " tuples recorded\n";
@@ -447,35 +516,28 @@ int main(int argc, char *argv[]) {
   // then store this coverage information in the appropriate data structures.
   // ------------------------------------------------------------------------ //
 
-  size_t       SeedCount = 0;
-  const size_t NumSeeds = SeedFiles.size();
-
-  if (!ShowProgBar)
-    StatMsg() << "Generating coverage for " << NumSeeds << " seeds... ";
-  StartTimer(ShowProgBar);
+  StatMsg() << "Running afl-showmap on " << SeedFiles.size() << " seeds...\n";
+  StartTimer(/*ShowProgBar=*/false);
 
-  EvalMaxSAT        Solver(/*nbMinimizeThread=*/0);
+  AFLCoverageMap    CovMap;
   MaxSATSeeds       SeedVars;
   MaxSATCoverageMap SeedCoverage;
+  EvalMaxSAT        Solver(/*nbMinimizeThread=*/0);
 
-  for (const auto &SeedFile : SeedFiles) {
-
-    // Execute seed
-    Cov.clear();
-    if (auto Err = getAFLCoverage(SeedFile, Cov)) {
+  if (auto Err = runShowmap(CovMap)) {
 
-      ErrMsg() << "Failed to get coverage for seed `" << SeedFile
-               << "`: " << Err << '\n';
-      return 1;
+    ErrMsg() << "Failed to generate coverage: " << Err << '\n';
+    return 1;
 
-    }
+  }
 
+  for (const auto &SeedCov : CovMap) {
     // Create a variable to represent the seed
     const SeedID Var = Solver.newVar();
-    SeedVars.push_back({Var, SeedFile});
+    SeedVars.emplace_back(Var, SeedCov.first());
 
     // Record the set of seeds that cover a particular edge
-    for (const auto &[Edge, Freq] : Cov) {
+    for (auto &[Edge, Freq] : SeedCov.second) {
 
       if (EdgesOnly) {
 
@@ -492,12 +554,10 @@ int main(int argc, char *argv[]) {
 
     }
 
-    if ((++SeedCount % 10 == 0) && ShowProgBar)
-      ProgBar.update(SeedCount * 100 / NumSeeds, "Generating seed coverage");
-
   }
 
-  EndTimer(ShowProgBar);
+  SuccMsg() << "afl-showmap completed in ";
+  EndTimer(/*ShowProgBar=*/false);
 
   // ------------------------------------------------------------------------ //
   // Set the hard and soft constraints in the solver
@@ -506,7 +566,7 @@ int main(int argc, char *argv[]) {
   if (!ShowProgBar) StatMsg() << "Generating constraints... ";
   StartTimer(ShowProgBar);
 
-  SeedCount = 0;
+  size_t SeedCount = 0;
 
   // Ensure that at least one seed is selected that covers a particular edge
   // (hard constraint)
@@ -552,7 +612,7 @@ int main(int argc, char *argv[]) {
   // ------------------------------------------------------------------------ //
 
   SmallVector<StringRef, 64> Solution;
-  SmallString<32>            OutputSeed;
+  SmallString<32>            InputSeed, OutputSeed;
 
   if (Solved) {
 
@@ -561,28 +621,31 @@ int main(int argc, char *argv[]) {
 
   } else {
 
-    ErrMsg() << "Failed to find an optimal solution for `" << CorpusDir
-             << "`\n";
+    ErrMsg() << "Failed to find an optimal solution for '" << InputDir << "'\n";
     return 1;
 
   }
 
   SuccMsg() << "Minimized corpus size: " << Solution.size() << " seeds\n";
 
-  if (!ShowProgBar) StatMsg() << "Copying to `" << OutputDir << "`... ";
+  if (!ShowProgBar) StatMsg() << "Copying to '" << OutputDir << "'... ";
   StartTimer(ShowProgBar);
 
   SeedCount = 0;
 
   for (const auto &Seed : Solution) {
 
+    InputSeed = SeedFiles[Seed];
+    sys::path::append(InputSeed, Seed);
+
     OutputSeed = OutputDir;
-    sys::path::append(OutputSeed, sys::path::filename(Seed));
+    sys::path::append(OutputSeed, Seed);
 
-    if ((EC = sys::fs::copy_file(Seed, OutputSeed))) {
+    if ((EC = sys::fs::copy_file(InputSeed, OutputSeed))) {
 
-      WarnMsg() << "Failed to copy `" << Seed << "` to `" << OutputDir
-                << "`: " << EC.message() << '\n';
+      ErrMsg() << "Failed to copy '" << Seed << "' to '" << OutputDir
+               << "': " << EC.message() << '\n';
+      return 1;
 
     }