[libFuzzer] move the implementation of the fork mode into a separate file

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

6 files changed, 165 insertions, 93 deletions

diff --git a/lib/fuzzer/CMakeLists.txt b/lib/fuzzer/CMakeLists.txt
index 0a696de9b..c81411339 100644
--- a/lib/fuzzer/CMakeLists.txt
+++ b/lib/fuzzer/CMakeLists.txt
@@ -6,6 +6,7 @@ set(LIBFUZZER_SOURCES
   FuzzerExtFunctionsWeak.cpp
   FuzzerExtFunctionsWindows.cpp
   FuzzerExtraCounters.cpp
+  FuzzerFork.cpp
   FuzzerIO.cpp
   FuzzerIOPosix.cpp
   FuzzerIOWindows.cpp
@@ -32,6 +33,7 @@ set(LIBFUZZER_HEADERS
   FuzzerExtFunctions.def
   FuzzerExtFunctions.h
   FuzzerFlags.def
+  FuzzerFork.h
   FuzzerIO.h
   FuzzerInterface.h
   FuzzerInternal.h
diff --git a/lib/fuzzer/FuzzerDriver.cpp b/lib/fuzzer/FuzzerDriver.cpp
index dc67512b6..00dae3fda 100644
--- a/lib/fuzzer/FuzzerDriver.cpp
+++ b/lib/fuzzer/FuzzerDriver.cpp
@@ -10,13 +10,14 @@
 
 #include "FuzzerCommand.h"
 #include "FuzzerCorpus.h"
+#include "FuzzerFork.h"
 #include "FuzzerIO.h"
 #include "FuzzerInterface.h"
 #include "FuzzerInternal.h"
+#include "FuzzerMerge.h"
 #include "FuzzerMutate.h"
 #include "FuzzerRandom.h"
 #include "FuzzerTracePC.h"
-#include "FuzzerMerge.h"
 #include <algorithm>
 #include <atomic>
 #include <chrono>
@@ -306,11 +307,6 @@ static std::string GetDedupTokenFromFile(const std::string &Path) {
   return S.substr(Beg, End - Beg);
 }
 
-static std::string TempPath(const char *Extension) {
-  return DirPlusFile(TmpDir(),
-                     "libFuzzerTemp." + std::to_string(GetPid()) + Extension);
-}
-
 int CleanseCrashInput(const Vector<std::string> &Args,
                        const FuzzingOptions &Options) {
   if (Inputs->size() != 1 || !Flags.exact_artifact_path) {
@@ -471,91 +467,6 @@ int MinimizeCrashInputInternalStep(Fuzzer *F, InputCorpus *Corpus) {
   return 0;
 }
 
-// This is just a skeleton of an experimental -fork=1 feature.
-void FuzzWithFork(Fuzzer *F, const FuzzingOptions &Options,
-                  const Vector<std::string> &Args,
-                  const Vector<std::string> &Corpora) {
-  Printf("INFO: -fork=1: doing fuzzing in a separate process in order to "
-         "be more resistant to crashes, timeouts, and OOMs\n");
-  auto Rand = F->GetMD().GetRand();
-
-  Vector<SizedFile> Corpus;
-  for (auto &Dir : Corpora)
-    GetSizedFilesFromDir(Dir, &Corpus);
-  std::sort(Corpus.begin(), Corpus.end());
-  auto CFPath = TempPath(".fork");
-  auto LogPath = TempPath(".log");
-
-  Vector<std::string> Files;
-  Set<uint32_t> Features;
-  if (!Corpus.empty()) {
-    CrashResistantMerge(Args, {}, Corpus, &Files, {}, &Features, CFPath, false);
-    RemoveFile(CFPath);
-  }
-  auto TempDir = TempPath("Dir");
-  MkDir(TempDir);
-  Printf("INFO: -fork=1: %zd seeds, starting to fuzz; scratch: %s\n",
-         Files.size(), TempDir.c_str());
-
-  Command BaseCmd(Args);
-  BaseCmd.removeFlag("fork");
-  for (auto &C : Corpora) // Remove all corpora from the args.
-    BaseCmd.removeArgument(C);
-  if (!BaseCmd.hasFlag("max_total_time"))
-    BaseCmd.addFlag("max_total_time", "60");
-  BaseCmd.addArgument(TempDir);
-  int ExitCode = 0;
-  for (size_t i = 0; i < 1000000; i++) {
-    // TODO: take new files from disk e.g. those generated by another process.
-    Command Cmd(BaseCmd);
-    if (Files.size() >= 2)
-      Cmd.addFlag("seed_inputs",
-                  Files[Rand.SkewTowardsLast(Files.size())] + "," +
-                      Files[Rand.SkewTowardsLast(Files.size())]);
-    Cmd.setOutputFile(LogPath);
-    Cmd.combineOutAndErr();
-    RmFilesInDir(TempDir);
-    ExitCode = ExecuteCommand(Cmd);
-    // Printf("done [%d] %s\n", ExitCode, Cmd.toString().c_str());
-    if (ExitCode == Options.InterruptExitCode)
-      break;
-    Vector<SizedFile> TempFiles;
-    Vector<std::string>FilesToAdd;
-    Set<uint32_t> NewFeatures;
-    GetSizedFilesFromDir(TempDir, &TempFiles);
-    CrashResistantMerge(Args, {}, TempFiles, &FilesToAdd, Features,
-                        &NewFeatures, CFPath, false);
-    RemoveFile(CFPath);
-    for (auto &Path : FilesToAdd) {
-      auto NewPath = F->WriteToOutputCorpus(FileToVector(Path, Options.MaxLen));
-      if (!NewPath.empty())
-        Files.push_back(NewPath);
-    }
-    Features.insert(NewFeatures.begin(), NewFeatures.end());
-    Printf("INFO: temp_files: %zd files_added: %zd newft: %zd ft: %zd\n",
-           TempFiles.size(), FilesToAdd.size(), NewFeatures.size(),
-           Features.size());
-    // Continue if our crash is one of the ignorred ones.
-    if (Options.IgnoreTimeouts && ExitCode == Options.TimeoutExitCode)
-      continue;
-    if (Options.IgnoreOOMs && ExitCode == Options.OOMExitCode)
-      continue;
-    // And exit if we don't ignore this crash.
-    if (ExitCode != 0) {
-      Printf("INFO: log from the inner process:\n%s",
-             FileToString(LogPath).c_str());
-      break;
-    }
-  }
-
-  RmFilesInDir(TempDir);
-  RmDir(TempDir);
-
-  // Use the exit code from the last child process.
-  Printf("Fork: exiting: %d\n", ExitCode);
-  exit(ExitCode);
-}
-
 void Merge(Fuzzer *F, FuzzingOptions &Options, const Vector<std::string> &Args,
            const Vector<std::string> &Corpora, const char *CFPathOrNull) {
   if (Corpora.size() < 2) {
@@ -814,7 +725,7 @@ int FuzzerDriver(int *argc, char ***argv, UserCallback Callback) {
   }
 
   if (Flags.fork)
-    FuzzWithFork(F, Options, Args, *Inputs);
+    FuzzWithFork(F->GetMD().GetRand(), Options, Args, *Inputs);
 
   if (Flags.merge)
     Merge(F, Options, Args, *Inputs, Flags.merge_control_file);
diff --git a/lib/fuzzer/FuzzerFork.cpp b/lib/fuzzer/FuzzerFork.cpp
new file mode 100644
index 000000000..bb0ea2a3c
--- /dev/null
+++ b/lib/fuzzer/FuzzerFork.cpp
@@ -0,0 +1,127 @@
+//===- FuzzerFork.cpp - run fuzzing in separate subprocesses --------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+// Spawn an orchestrate separate fuzzing processes.
+//===----------------------------------------------------------------------===//
+
+#include "FuzzerCommand.h"
+#include "FuzzerFork.h"
+#include "FuzzerIO.h"
+#include "FuzzerMerge.h"
+#include "FuzzerSHA1.h"
+#include "FuzzerUtil.h"
+
+namespace fuzzer {
+
+// This is just a skeleton of an experimental -fork=1 feature.
+void FuzzWithFork(Random &Rand, const FuzzingOptions &Options,
+                  const Vector<std::string> &Args,
+                  const Vector<std::string> &CorpusDirs) {
+  Printf("INFO: -fork=1: doing fuzzing in a separate process in order to "
+         "be more resistant to crashes, timeouts, and OOMs\n");
+
+  Vector<SizedFile> SeedFiles;
+  for (auto &Dir : CorpusDirs)
+    GetSizedFilesFromDir(Dir, &SeedFiles);
+  std::sort(SeedFiles.begin(), SeedFiles.end());
+  auto CFPath = TempPath(".fork");
+  auto LogPath = TempPath(".log");
+  auto TempDir = TempPath(".scratch_dir");
+  std::string MainCorpusDir;
+  if (CorpusDirs.empty())
+    MkDir(MainCorpusDir = TempPath(".corpus_dir"));
+  else
+    MainCorpusDir = CorpusDirs[0];
+  MkDir(TempDir);
+
+  Vector<std::string> Files;
+  Set<uint32_t> Features;
+  if (!SeedFiles.empty()) {
+    CrashResistantMerge(Args, {}, SeedFiles, &Files, {}, &Features, CFPath,
+                        false);
+    RemoveFile(CFPath);
+  }
+  Printf("INFO: -fork=1: %zd seeds, starting to fuzz; scratch: %s\n",
+         Files.size(), TempDir.c_str());
+
+  Command BaseCmd(Args);
+  BaseCmd.removeFlag("fork");
+  for (auto &C : CorpusDirs) // Remove all corpora from the args.
+    BaseCmd.removeArgument(C);
+  BaseCmd.addArgument(TempDir);
+  BaseCmd.addFlag("len_control", "0");  // len_control is bad for short runs.
+  BaseCmd.addFlag("reload", "0");  // working in an isolated dir, no reload.
+  int ExitCode = 0;
+  int max_total_time = 1;
+  for (size_t i = 0; i < 1000000; i++) {
+    // TODO: take new files from disk e.g. those generated by another process.
+    Command Cmd(BaseCmd);
+    if (size_t CorpusSubsetSize = std::min(Files.size(), (size_t)10)) {
+      std::string Seeds;
+      for (size_t i = 0; i < CorpusSubsetSize; i++) {
+        if (i) Seeds += ",";
+        Seeds += Files[Rand.SkewTowardsLast(Files.size())];
+      }
+      Cmd.addFlag("seed_inputs", Seeds);
+    }
+    if (Options.MaxTotalTimeSec > max_total_time)
+      max_total_time++;
+    if (!Cmd.hasFlag("max_total_time"))
+      Cmd.addFlag("max_total_time", std::to_string(max_total_time));
+    Cmd.setOutputFile(LogPath);
+    Cmd.combineOutAndErr();
+    RmFilesInDir(TempDir);
+    ExitCode = ExecuteCommand(Cmd);
+    // Printf("done [%d] %s\n", ExitCode, Cmd.toString().c_str());
+    if (ExitCode == Options.InterruptExitCode)
+      break;
+    Vector<SizedFile> TempFiles;
+    Vector<std::string>FilesToAdd;
+    Set<uint32_t> NewFeatures;
+    GetSizedFilesFromDir(TempDir, &TempFiles);
+    if (!TempFiles.empty())
+      CrashResistantMerge(Args, {}, TempFiles, &FilesToAdd, Features,
+                          &NewFeatures, CFPath, false);
+    RemoveFile(CFPath);
+    for (auto &Path : FilesToAdd) {
+      auto U = FileToVector(Path);
+      auto NewPath = DirPlusFile(MainCorpusDir, Hash(U));
+      WriteToFile(U, NewPath);
+      Files.push_back(NewPath);
+    }
+    Features.insert(NewFeatures.begin(), NewFeatures.end());
+    Printf("INFO: temp_files: %zd files_added: %zd newft: %zd ft: %zd\n",
+           TempFiles.size(), FilesToAdd.size(), NewFeatures.size(),
+           Features.size());
+    // Continue if our crash is one of the ignorred ones.
+    if (Options.IgnoreTimeouts && ExitCode == Options.TimeoutExitCode)
+      continue;
+    if (Options.IgnoreOOMs && ExitCode == Options.OOMExitCode)
+      continue;
+    // And exit if we don't ignore this crash.
+    if (ExitCode != 0) {
+      Printf("INFO: log from the inner process:\n%s",
+             FileToString(LogPath).c_str());
+      break;
+    }
+  }
+
+  RmFilesInDir(TempDir);
+  RmDir(TempDir);
+
+  if (CorpusDirs.empty()) {
+    RmFilesInDir(MainCorpusDir);
+    RmDir(MainCorpusDir);
+  }
+
+  // Use the exit code from the last child process.
+  Printf("Fork: exiting: %d\n", ExitCode);
+  exit(ExitCode);
+}
+
+} // namespace fuzzer
+
diff --git a/lib/fuzzer/FuzzerFork.h b/lib/fuzzer/FuzzerFork.h
new file mode 100644
index 000000000..d9e8de515
--- /dev/null
+++ b/lib/fuzzer/FuzzerFork.h
@@ -0,0 +1,24 @@
+//===- FuzzerFork.h - run fuzzing in sub-processes --------------*- C++ -* ===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_FUZZER_FORK_H
+#define LLVM_FUZZER_FORK_H
+
+#include "FuzzerDefs.h"
+#include "FuzzerOptions.h"
+#include "FuzzerRandom.h"
+
+#include <string>
+
+namespace fuzzer {
+void FuzzWithFork(Random &Rand, const FuzzingOptions &Options,
+                  const Vector<std::string> &Args,
+                  const Vector<std::string> &CorpusDirs);
+} // namespace fuzzer
+
+#endif // LLVM_FUZZER_FORK_H
diff --git a/lib/fuzzer/FuzzerIO.cpp b/lib/fuzzer/FuzzerIO.cpp
index ba4153bfc..91e1d2080 100644
--- a/lib/fuzzer/FuzzerIO.cpp
+++ b/lib/fuzzer/FuzzerIO.cpp
@@ -8,9 +8,10 @@
 // IO functions.
 //===----------------------------------------------------------------------===//
 
-#include "FuzzerIO.h"
 #include "FuzzerDefs.h"
 #include "FuzzerExtFunctions.h"
+#include "FuzzerIO.h"
+#include "FuzzerUtil.h"
 #include <algorithm>
 #include <cstdarg>
 #include <fstream>
@@ -141,4 +142,9 @@ void RmFilesInDir(const std::string &Path) {
     RemoveFile(F);
 }
 
+std::string TempPath(const char *Extension) {
+  return DirPlusFile(TmpDir(),
+                     "libFuzzerTemp." + std::to_string(GetPid()) + Extension);
+}
+
 }  // namespace fuzzer
diff --git a/lib/fuzzer/FuzzerIO.h b/lib/fuzzer/FuzzerIO.h
index a78d50592..9d849bed6 100644
--- a/lib/fuzzer/FuzzerIO.h
+++ b/lib/fuzzer/FuzzerIO.h
@@ -39,6 +39,8 @@ std::string DirName(const std::string &FileName);
 // Returns path to a TmpDir.
 std::string TmpDir();
 
+std::string TempPath(const char *Extension);
+
 bool IsInterestingCoverageFile(const std::string &FileName);
 
 void DupAndCloseStderr();