[libFuzzer] allow merge to resume after being preempted

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@317767 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 61 insertions, 25 deletions

diff --git a/lib/fuzzer/FuzzerFlags.def b/lib/fuzzer/FuzzerFlags.def
index db8c9eda0..9b2115779 100644
--- a/lib/fuzzer/FuzzerFlags.def
+++ b/lib/fuzzer/FuzzerFlags.def
@@ -41,6 +41,8 @@ FUZZER_FLAG_INT(merge, 0, "If 1, the 2-nd, 3-rd, etc corpora will be "
 FUZZER_FLAG_STRING(merge_inner, "internal flag")
 FUZZER_FLAG_STRING(merge_control_file,
                    "Specify a control file used for the merge proccess. "
+                   "If a merge process gets killed it tries to leave this file "
+                   "in a state suitable for resuming the merge. "
                    "By default a temporary file will be used.")
 FUZZER_FLAG_STRING(save_coverage_summary, "Experimental:"
                    " save coverage summary to a given file."
diff --git a/lib/fuzzer/FuzzerMerge.cpp b/lib/fuzzer/FuzzerMerge.cpp
index e6e935dfc..89b4821d0 100644
--- a/lib/fuzzer/FuzzerMerge.cpp
+++ b/lib/fuzzer/FuzzerMerge.cpp
@@ -256,6 +256,22 @@ void Fuzzer::CrashResistantMergeInternalStep(const std::string &CFPath) {
   }
 }
 
+static void WriteNewControlFile(const std::string &CFPath,
+                                const Vector<SizedFile> &AllFiles,
+                                size_t NumFilesInFirstCorpus) {
+  RemoveFile(CFPath);
+  std::ofstream ControlFile(CFPath);
+  ControlFile << AllFiles.size() << "\n";
+  ControlFile << NumFilesInFirstCorpus << "\n";
+  for (auto &SF: AllFiles)
+    ControlFile << SF.File << "\n";
+  if (!ControlFile) {
+    Printf("MERGE-OUTER: failed to write to the control file: %s\n",
+           CFPath.c_str());
+    exit(1);
+  }
+}
+
 // Outer process. Does not call the target code and thus sohuld not fail.
 void Fuzzer::CrashResistantMerge(const Vector<std::string> &Args,
                                  const Vector<std::string> &Corpora,
@@ -266,46 +282,64 @@ void Fuzzer::CrashResistantMerge(const Vector<std::string> &Args,
     Printf("Merge requires two or more corpus dirs\n");
     return;
   }
-  Vector<SizedFile> AllFiles;
-  GetSizedFilesFromDir(Corpora[0], &AllFiles);
-  size_t NumFilesInFirstCorpus = AllFiles.size();
-  std::sort(AllFiles.begin(), AllFiles.end());
-  for (size_t i = 1; i < Corpora.size(); i++)
-    GetSizedFilesFromDir(Corpora[i], &AllFiles);
-  std::sort(AllFiles.begin() + NumFilesInFirstCorpus, AllFiles.end());
-  Printf("MERGE-OUTER: %zd files, %zd in the initial corpus\n",
-         AllFiles.size(), NumFilesInFirstCorpus);
   auto CFPath =
       MergeControlFilePathOrNull
           ? MergeControlFilePathOrNull
           : DirPlusFile(TmpDir(),
                         "libFuzzerTemp." + std::to_string(GetPid()) + ".txt");
-  // Write the control file.
-  RemoveFile(CFPath);
-  std::ofstream ControlFile(CFPath);
-  ControlFile << AllFiles.size() << "\n";
-  ControlFile << NumFilesInFirstCorpus << "\n";
-  for (auto &SF: AllFiles)
-    ControlFile << SF.File << "\n";
-  if (!ControlFile) {
-    Printf("MERGE-OUTER: failed to write to the control file: %s\n",
-           CFPath.c_str());
-    exit(1);
+
+  size_t NumAttempts = 0;
+  if (MergeControlFilePathOrNull && FileSize(MergeControlFilePathOrNull)) {
+    Printf("MERGE-OUTER: non-empty control file provided: '%s'\n",
+           MergeControlFilePathOrNull);
+    Merger M;
+    std::ifstream IF(MergeControlFilePathOrNull);
+    if (M.Parse(IF, /*ParseCoverage=*/false)) {
+      Printf("MERGE-OUTER: control file ok, %zd files total,"
+             " first not processed file %zd\n",
+             M.Files.size(), M.FirstNotProcessedFile);
+      if (!M.LastFailure.empty())
+        Printf("MERGE-OUTER: '%s' will be skipped as unlucky "
+               "(merge has stumbled on it the last time)\n",
+               M.LastFailure.c_str());
+      if (M.FirstNotProcessedFile >= M.Files.size()) {
+        Printf("MERGE-OUTER: nothing to do, merge has been completed before\n");
+        exit(0);
+      }
+
+      NumAttempts = M.Files.size() - M.FirstNotProcessedFile;
+    } else {
+      Printf("MERGE-OUTER: bad control file, will overwrite it\n");
+    }
+  }
+
+  if (!NumAttempts) {
+    // The supplied control file is empty or bad, create a fresh one.
+    Vector<SizedFile> AllFiles;
+    GetSizedFilesFromDir(Corpora[0], &AllFiles);
+    size_t NumFilesInFirstCorpus = AllFiles.size();
+    std::sort(AllFiles.begin(), AllFiles.end());
+    for (size_t i = 1; i < Corpora.size(); i++)
+      GetSizedFilesFromDir(Corpora[i], &AllFiles);
+    std::sort(AllFiles.begin() + NumFilesInFirstCorpus, AllFiles.end());
+    Printf("MERGE-OUTER: %zd files, %zd in the initial corpus\n",
+           AllFiles.size(), NumFilesInFirstCorpus);
+    WriteNewControlFile(CFPath, AllFiles, NumFilesInFirstCorpus);
+    NumAttempts = AllFiles.size();
   }
-  ControlFile.close();
 
-  // Execute the inner process untill it passes.
+  // Execute the inner process until it passes.
   // Every inner process should execute at least one input.
   auto BaseCmd = SplitBefore("-ignore_remaining_args=1",
                              CloneArgsWithoutX(Args, "merge"));
   bool Success = false;
-  for (size_t i = 1; i <= AllFiles.size(); i++) {
-    Printf("MERGE-OUTER: attempt %zd\n", i);
+  for (size_t Attempt = 1; Attempt <= NumAttempts; Attempt++) {
+    Printf("MERGE-OUTER: attempt %zd\n", Attempt);
     auto ExitCode =
         ExecuteCommand(BaseCmd.first + " -merge_control_file=" + CFPath +
                        " -merge_inner=1 " + BaseCmd.second);
     if (!ExitCode) {
-      Printf("MERGE-OUTER: succesfull in %zd attempt(s)\n", i);
+      Printf("MERGE-OUTER: succesfull in %zd attempt(s)\n", Attempt);
       Success = true;
       break;
     }