[libFuzzer] add -features_dir= flag to dump unique input features on disk

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

9 files changed, 53 insertions, 6 deletions

diff --git a/lib/fuzzer/FuzzerCorpus.h b/lib/fuzzer/FuzzerCorpus.h
index 23a9a89bf..6a95ef3a8 100644
--- a/lib/fuzzer/FuzzerCorpus.h
+++ b/lib/fuzzer/FuzzerCorpus.h
@@ -85,9 +85,10 @@ class InputCorpus {
 
   bool empty() const { return Inputs.empty(); }
   const Unit &operator[] (size_t Idx) const { return Inputs[Idx]->U; }
-  void AddToCorpus(const Unit &U, size_t NumFeatures, bool MayDeleteFile,
-                   bool HasFocusFunction, const Vector<uint32_t> &FeatureSet,
-                   const DataFlowTrace &DFT, const InputInfo *BaseII) {
+  InputInfo *AddToCorpus(const Unit &U, size_t NumFeatures, bool MayDeleteFile,
+                         bool HasFocusFunction,
+                         const Vector<uint32_t> &FeatureSet,
+                         const DataFlowTrace &DFT, const InputInfo *BaseII) {
     assert(!U.empty());
     if (FeatureDebug)
       Printf("ADD_TO_CORPUS %zd NF %zd\n", Inputs.size(), NumFeatures);
@@ -113,6 +114,7 @@ class InputCorpus {
     UpdateCorpusDistribution();
     PrintCorpus();
     // ValidateFeatureSet();
+    return &II;
   }
 
   // Debug-only
diff --git a/lib/fuzzer/FuzzerDriver.cpp b/lib/fuzzer/FuzzerDriver.cpp
index 9c99d5ffb..a51ac9310 100644
--- a/lib/fuzzer/FuzzerDriver.cpp
+++ b/lib/fuzzer/FuzzerDriver.cpp
@@ -659,6 +659,8 @@ int FuzzerDriver(int *argc, char ***argv, UserCallback Callback) {
     Options.FocusFunction = Flags.focus_function;
   if (Flags.data_flow_trace)
     Options.DataFlowTrace = Flags.data_flow_trace;
+  if (Flags.features_dir)
+    Options.FeaturesDir = Flags.features_dir;
   Options.LazyCounters = Flags.lazy_counters;
 
   unsigned Seed = Flags.seed;
diff --git a/lib/fuzzer/FuzzerFlags.def b/lib/fuzzer/FuzzerFlags.def
index 79eb75c3f..d3da3e5cb 100644
--- a/lib/fuzzer/FuzzerFlags.def
+++ b/lib/fuzzer/FuzzerFlags.def
@@ -67,6 +67,10 @@ FUZZER_FLAG_INT(cleanse_crash, 0, "If 1, tries to cleanse the provided"
   " Use with -exact_artifact_path to specify the output."
   )
 FUZZER_FLAG_INT(minimize_crash_internal_step, 0, "internal flag")
+FUZZER_FLAG_STRING(features_dir, "internal flag. Used to dump feature sets on disk."
+  "Every time a new input is added to the corpus, a corresponding file in the features_dir"
+  " is created containing the unique features of that input."
+  " Features are stored in binary format.")
 FUZZER_FLAG_INT(use_counters, 1, "Use coverage counters")
 FUZZER_FLAG_INT(use_memmem, 1,
                 "Use hints from intercepting memmem, strstr, etc")
diff --git a/lib/fuzzer/FuzzerIO.cpp b/lib/fuzzer/FuzzerIO.cpp
index 423ab4bef..a8140b601 100644
--- a/lib/fuzzer/FuzzerIO.cpp
+++ b/lib/fuzzer/FuzzerIO.cpp
@@ -61,10 +61,14 @@ void CopyFileToErr(const std::string &Path) {
 }
 
 void WriteToFile(const Unit &U, const std::string &Path) {
+  WriteToFile(U.data(), U.size(), Path);
+}
+
+void WriteToFile(const uint8_t *Data, size_t Size, const std::string &Path) {
   // Use raw C interface because this function may be called from a sig handler.
   FILE *Out = fopen(Path.c_str(), "wb");
   if (!Out) return;
-  fwrite(U.data(), sizeof(U[0]), U.size(), Out);
+  fwrite(Data, sizeof(Data[0]), Size, Out);
   fclose(Out);
 }
 
diff --git a/lib/fuzzer/FuzzerIO.h b/lib/fuzzer/FuzzerIO.h
index 9d4e7650a..cbfafa5c3 100644
--- a/lib/fuzzer/FuzzerIO.h
+++ b/lib/fuzzer/FuzzerIO.h
@@ -24,6 +24,7 @@ std::string FileToString(const std::string &Path);
 
 void CopyFileToErr(const std::string &Path);
 
+void WriteToFile(const uint8_t *Data, size_t Size, const std::string &Path);
 void WriteToFile(const Unit &U, const std::string &Path);
 
 void ReadDirToVectorOfUnits(const char *Path, Vector<Unit> *V,
@@ -89,6 +90,7 @@ int CloseFile(int Fd);
 int DuplicateFile(int Fd);
 
 void RemoveFile(const std::string &Path);
+void RenameFile(const std::string &OldPath, const std::string &NewPath);
 
 void DiscardOutput(int Fd);
 
diff --git a/lib/fuzzer/FuzzerIOPosix.cpp b/lib/fuzzer/FuzzerIOPosix.cpp
index 953848a89..cfd69bbc8 100644
--- a/lib/fuzzer/FuzzerIOPosix.cpp
+++ b/lib/fuzzer/FuzzerIOPosix.cpp
@@ -120,6 +120,10 @@ void RemoveFile(const std::string &Path) {
   unlink(Path.c_str());
 }
 
+void RenameFile(const std::string &OldPath, const std::string &NewPath) {
+  rename(OldPath.c_str(), NewPath.c_str());
+}
+
 void DiscardOutput(int Fd) {
   FILE* Temp = fopen("/dev/null", "w");
   if (!Temp)
diff --git a/lib/fuzzer/FuzzerIOWindows.cpp b/lib/fuzzer/FuzzerIOWindows.cpp
index b8c200b81..510afebef 100644
--- a/lib/fuzzer/FuzzerIOWindows.cpp
+++ b/lib/fuzzer/FuzzerIOWindows.cpp
@@ -219,6 +219,10 @@ void RemoveFile(const std::string &Path) {
   _unlink(Path.c_str());
 }
 
+void RenameFile(const std::string &OldPath, const std::string &NewPath) {
+  rename(OldPath.c_str(), NewPath.c_str());
+}
+
 void DiscardOutput(int Fd) {
   FILE* Temp = fopen("nul", "w");
   if (!Temp)
diff --git a/lib/fuzzer/FuzzerLoop.cpp b/lib/fuzzer/FuzzerLoop.cpp
index 40461c2e1..b68d759cf 100644
--- a/lib/fuzzer/FuzzerLoop.cpp
+++ b/lib/fuzzer/FuzzerLoop.cpp
@@ -445,6 +445,24 @@ void Fuzzer::PrintPulseAndReportSlowInput(const uint8_t *Data, size_t Size) {
   }
 }
 
+static void WriteFeatureSetToFile(const std::string &FeaturesDir,
+                                  const uint8_t Sha1[],
+                                  const Vector<uint32_t> &FeatureSet) {
+  if (FeaturesDir.empty() || FeatureSet.empty()) return;
+  WriteToFile(reinterpret_cast<const uint8_t *>(FeatureSet.data()),
+              FeatureSet.size() * sizeof(FeatureSet[0]),
+              DirPlusFile(FeaturesDir, Sha1ToString(Sha1)));
+  Printf("Features: %s\n", Sha1ToString(Sha1).c_str());
+}
+
+static void RenameFeatureSetFile(const std::string &FeaturesDir,
+                                 const std::string &OldFile,
+                                 const std::string &NewFile) {
+  if (FeaturesDir.empty()) return;
+  RenameFile(DirPlusFile(FeaturesDir, OldFile),
+             DirPlusFile(FeaturesDir, NewFile));
+}
+
 bool Fuzzer::RunOne(const uint8_t *Data, size_t Size, bool MayDeleteFile,
                     InputInfo *II, bool *FoundUniqFeatures) {
   if (!Size)
@@ -469,15 +487,21 @@ bool Fuzzer::RunOne(const uint8_t *Data, size_t Size, bool MayDeleteFile,
   size_t NumNewFeatures = Corpus.NumFeatureUpdates() - NumUpdatesBefore;
   if (NumNewFeatures) {
     TPC.UpdateObservedPCs();
-    Corpus.AddToCorpus({Data, Data + Size}, NumNewFeatures, MayDeleteFile,
-                       TPC.ObservedFocusFunction(), UniqFeatureSetTmp, DFT, II);
+    auto NewII = Corpus.AddToCorpus({Data, Data + Size}, NumNewFeatures,
+                                    MayDeleteFile, TPC.ObservedFocusFunction(),
+                                    UniqFeatureSetTmp, DFT, II);
+    WriteFeatureSetToFile(Options.FeaturesDir, NewII->Sha1,
+                          NewII->UniqFeatureSet);
     return true;
   }
   if (II && FoundUniqFeaturesOfII &&
       II->DataFlowTraceForFocusFunction.empty() &&
       FoundUniqFeaturesOfII == II->UniqFeatureSet.size() &&
       II->U.size() > Size) {
+    auto OldFeaturesFile = Sha1ToString(II->Sha1);
     Corpus.Replace(II, {Data, Data + Size});
+    RenameFeatureSetFile(Options.FeaturesDir, OldFeaturesFile,
+                         Sha1ToString(II->Sha1));
     return true;
   }
   return false;
diff --git a/lib/fuzzer/FuzzerOptions.h b/lib/fuzzer/FuzzerOptions.h
index 7a607e8ee..d48439daa 100644
--- a/lib/fuzzer/FuzzerOptions.h
+++ b/lib/fuzzer/FuzzerOptions.h
@@ -51,6 +51,7 @@ struct FuzzingOptions {
   std::string ExitOnItem;
   std::string FocusFunction;
   std::string DataFlowTrace;
+  std::string FeaturesDir;
   bool SaveArtifacts = true;
   bool PrintNEW = true; // Print a status line when new units are found;
   bool PrintNewCovPcs = false;