From 265cf04f9c738f3be7a941d87eaee80d928be39e Mon Sep 17 00:00:00 2001
From: Kostya Serebryany <kcc@google.com>
Date: Wed, 28 Jan 2015 23:48:39 +0000
Subject: [PATCH] [fuzzer] add option -save_minimized_corpus

llvm-svn: 227395
---
 llvm/lib/Fuzzer/FuzzerFlags.def  |  3 +++
 llvm/lib/Fuzzer/FuzzerIO.cpp     |  9 +++++++--
 llvm/lib/Fuzzer/FuzzerInternal.h |  6 +++++-
 llvm/lib/Fuzzer/FuzzerLoop.cpp   | 11 ++++++++++-
 llvm/lib/Fuzzer/FuzzerMain.cpp   |  2 ++
 5 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Fuzzer/FuzzerFlags.def b/llvm/lib/Fuzzer/FuzzerFlags.def
index af5d81b..e165ebd 100644
--- a/llvm/lib/Fuzzer/FuzzerFlags.def
+++ b/llvm/lib/Fuzzer/FuzzerFlags.def
@@ -22,3 +22,6 @@ FUZZER_FLAG(int, exit_on_first, 0,
             "If 1, exit after the first new interesting input is found.")
 FUZZER_FLAG(int, timeout, -1, "Timeout in seconds (if positive).")
 FUZZER_FLAG(int, help, 0, "Print help.")
+FUZZER_FLAG(
+    int, save_minimized_corpus, 0,
+    "If 1, the minimized corpus is saved into the first input directory")
diff --git a/llvm/lib/Fuzzer/FuzzerIO.cpp b/llvm/lib/Fuzzer/FuzzerIO.cpp
index 172540c..4e0ac81 100644
--- a/llvm/lib/Fuzzer/FuzzerIO.cpp
+++ b/llvm/lib/Fuzzer/FuzzerIO.cpp
@@ -13,7 +13,7 @@
 #include <dirent.h>
 namespace fuzzer {
 
-std::vector<std::string> ListFilesInDir(const std::string &Dir) {
+static std::vector<std::string> ListFilesInDir(const std::string &Dir) {
   std::vector<std::string> V;
   DIR *D = opendir(Dir.c_str());
   if (!D) return V;
@@ -38,7 +38,12 @@ void WriteToFile(const Unit &U, const std::string &Path) {
 
 void ReadDirToVectorOfUnits(const char *Path, std::vector<Unit> *V) {
   for (auto &X : ListFilesInDir(Path))
-    V->push_back(FileToVector(std::string(Path) + "/" + X));
+    V->push_back(FileToVector(DirPlusFile(Path, X)));
+}
+
+std::string DirPlusFile(const std::string &DirPath,
+                        const std::string &FileName) {
+  return DirPath + "/" + FileName;
 }
 
 }  // namespace fuzzer
diff --git a/llvm/lib/Fuzzer/FuzzerInternal.h b/llvm/lib/Fuzzer/FuzzerInternal.h
index f5492182..c361ffb 100644
--- a/llvm/lib/Fuzzer/FuzzerInternal.h
+++ b/llvm/lib/Fuzzer/FuzzerInternal.h
@@ -20,9 +20,11 @@ typedef std::vector<uint8_t> Unit;
 using namespace std::chrono;
 
 Unit ReadFile(const char *Path);
-std::vector<std::string> ListFilesInDir(const std::string &Dir);
 void ReadDirToVectorOfUnits(const char *Path, std::vector<Unit> *V);
 void WriteToFile(const Unit &U, const std::string &Path);
+// Returns "Dir/FileName" or equivalent for the current OS.
+std::string DirPlusFile(const std::string &DirPath,
+                        const std::string &FileName);
 
 void Mutate(Unit *U, size_t MaxLen);
 
@@ -53,6 +55,8 @@ class Fuzzer {
   void ReadDir(const std::string &Path) {
     ReadDirToVectorOfUnits(Path.c_str(), &Corpus);
   }
+  // Save the current corpus to OutputCorpus.
+  void SaveCorpus();
 
   static void AlarmCallback();
 
diff --git a/llvm/lib/Fuzzer/FuzzerLoop.cpp b/llvm/lib/Fuzzer/FuzzerLoop.cpp
index b0bfce7..1f2193a 100644
--- a/llvm/lib/Fuzzer/FuzzerLoop.cpp
+++ b/llvm/lib/Fuzzer/FuzzerLoop.cpp
@@ -96,7 +96,7 @@ size_t Fuzzer::RunOne(const Unit &U) {
 
 void Fuzzer::WriteToOutputCorpus(const Unit &U) {
   if (Options.OutputCorpus.empty()) return;
-  std::string Path = Options.OutputCorpus + "/" + Hash(U);
+  std::string Path = DirPlusFile(Options.OutputCorpus, Hash(U));
   WriteToFile(U, Path);
   if (Options.Verbosity >= 2)
     std::cerr << "Written to " << Path << std::endl;
@@ -108,6 +108,15 @@ void Fuzzer::WriteToCrash(const Unit &U, const char *Prefix) {
   std::cerr << "CRASHED; file written to " << Path << std::endl;
 }
 
+void Fuzzer::SaveCorpus() {
+  if (Options.OutputCorpus.empty()) return;
+  for (const auto &U : Corpus)
+    WriteToFile(U, DirPlusFile(Options.OutputCorpus, Hash(U)));
+  if (Options.Verbosity)
+    std::cerr << "Written corpus of " << Corpus.size() << " files to "
+              << Options.OutputCorpus << "\n";
+}
+
 size_t Fuzzer::MutateAndTestOne(Unit *U) {
   size_t NewUnits = 0;
   for (size_t i = 0; i < Options.MutateDepth; i++) {
diff --git a/llvm/lib/Fuzzer/FuzzerMain.cpp b/llvm/lib/Fuzzer/FuzzerMain.cpp
index 9d4004e..6031fc8 100644
--- a/llvm/lib/Fuzzer/FuzzerMain.cpp
+++ b/llvm/lib/Fuzzer/FuzzerMain.cpp
@@ -139,6 +139,8 @@ int main(int argc, char **argv) {
   if (F.CorpusSize() == 0)
     F.AddToCorpus(Unit());  // Can't fuzz empty corpus, so add an empty input.
   F.ShuffleAndMinimize();
+  if (Flags.save_minimized_corpus)
+    F.SaveCorpus();
   F.Loop(Flags.iterations < 0 ? INT_MAX : Flags.iterations);
   if (Flags.verbosity)
     std::cerr << "Done\n";
-- 
2.7.4