[libFuzzer] make the corpus elements aware of their data flow traces

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

5 files changed, 41 insertions, 5 deletions

diff --git a/lib/fuzzer/FuzzerCorpus.h b/lib/fuzzer/FuzzerCorpus.h
index b4ba64b4b..c1603e1ad 100644
--- a/lib/fuzzer/FuzzerCorpus.h
+++ b/lib/fuzzer/FuzzerCorpus.h
@@ -12,6 +12,7 @@
 #ifndef LLVM_FUZZER_CORPUS
 #define LLVM_FUZZER_CORPUS
 
+#include "FuzzerDataFlowTrace.h"
 #include "FuzzerDefs.h"
 #include "FuzzerIO.h"
 #include "FuzzerRandom.h"
@@ -37,6 +38,7 @@ struct InputInfo {
   bool Reduced = false;
   bool HasFocusFunction = false;
   Vector<uint32_t> UniqFeatureSet;
+  Vector<bool> DataFlowTraceForFocusFunction;
 };
 
 class InputCorpus {
@@ -76,10 +78,17 @@ class InputCorpus {
     });
   }
 
+  size_t NumInputsWithDataFlowTrace() {
+    return std::count_if(Inputs.begin(), Inputs.end(), [](const InputInfo *II) {
+      return !II->DataFlowTraceForFocusFunction.empty();
+    });
+  }
+
   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) {
+                   bool HasFocusFunction, const Vector<uint32_t> &FeatureSet,
+                   const DataFlowTrace &DFT) {
     assert(!U.empty());
     if (FeatureDebug)
       Printf("ADD_TO_CORPUS %zd NF %zd\n", Inputs.size(), NumFeatures);
@@ -92,7 +101,11 @@ class InputCorpus {
     II.HasFocusFunction = HasFocusFunction;
     std::sort(II.UniqFeatureSet.begin(), II.UniqFeatureSet.end());
     ComputeSHA1(U.data(), U.size(), II.Sha1);
-    Hashes.insert(Sha1ToString(II.Sha1));
+    auto Sha1Str = Sha1ToString(II.Sha1);
+    Hashes.insert(Sha1Str);
+    if (HasFocusFunction)
+      if (auto V = DFT.Get(Sha1Str))
+        II.DataFlowTraceForFocusFunction = *V;
     UpdateCorpusDistribution();
     PrintCorpus();
     // ValidateFeatureSet();
diff --git a/lib/fuzzer/FuzzerDataFlowTrace.cpp b/lib/fuzzer/FuzzerDataFlowTrace.cpp
index 69efd6f38..114034c5b 100644
--- a/lib/fuzzer/FuzzerDataFlowTrace.cpp
+++ b/lib/fuzzer/FuzzerDataFlowTrace.cpp
@@ -73,6 +73,7 @@ void DataFlowTrace::Init(const std::string &DirPath,
             ParseError("the trace should contain only 0 or 1");
           V[I] = Beg[I] == '1';
         }
+        Traces[Name] = V;
         // Print just a few small traces.
         if (NumTracesWithFocusFunction <= 3 && Len <= 16)
           Printf("%s => |%s|\n", Name.c_str(), L.c_str() + SpacePos + 1);
diff --git a/lib/fuzzer/FuzzerDataFlowTrace.h b/lib/fuzzer/FuzzerDataFlowTrace.h
index 2b7b71fdb..1511430c3 100644
--- a/lib/fuzzer/FuzzerDataFlowTrace.h
+++ b/lib/fuzzer/FuzzerDataFlowTrace.h
@@ -31,9 +31,25 @@
 
 #include "FuzzerDefs.h"
 
+#include <unordered_map>
+#include <vector>
+#include <string>
+
 namespace fuzzer {
-struct DataFlowTrace {
+class DataFlowTrace {
+ public:
   void Init(const std::string &DirPath, const std::string &FocusFunction);
+  void Clear() { Traces.clear(); }
+  const Vector<bool> *Get(const std::string &InputSha1) const {
+    auto It = Traces.find(InputSha1);
+    if (It != Traces.end())
+      return &It->second;
+    return nullptr;
+  }
+
+ private:
+  // Input's sha1 => DFT for the FocusFunction.
+  std::unordered_map<std::string, Vector<bool> > Traces;
 };
 }  // namespace fuzzer
 
diff --git a/lib/fuzzer/FuzzerLoop.cpp b/lib/fuzzer/FuzzerLoop.cpp
index a195d21d3..346f90e6e 100644
--- a/lib/fuzzer/FuzzerLoop.cpp
+++ b/lib/fuzzer/FuzzerLoop.cpp
@@ -469,10 +469,11 @@ bool Fuzzer::RunOne(const uint8_t *Data, size_t Size, bool MayDeleteFile,
     TPC.UpdateObservedPCs();
     Corpus.AddToCorpus({Data, Data + Size}, NumNewFeatures, MayDeleteFile,
                        TPC.ObservedFocusFunction(),
-                       UniqFeatureSetTmp);
+                       UniqFeatureSetTmp, DFT);
     return true;
   }
   if (II && FoundUniqFeaturesOfII &&
+      II->DataFlowTraceForFocusFunction.empty() &&
       FoundUniqFeaturesOfII == II->UniqFeatureSet.size() &&
       II->U.size() > Size) {
     Corpus.Replace(II, {Data, Data + Size});
@@ -739,6 +740,9 @@ void Fuzzer::ReadAndExecuteSeedCorpora(const Vector<std::string> &CorpusDirs) {
   if (!Options.FocusFunction.empty())
     Printf("INFO: %zd/%zd inputs touch the focus function\n",
            Corpus.NumInputsThatTouchFocusFunction(), Corpus.size());
+  if (!Options.DataFlowTrace.empty())
+    Printf("INFO: %zd/%zd inputs have the Data Flow Trace\n",
+           Corpus.NumInputsWithDataFlowTrace(), Corpus.size());
 
   if (Corpus.empty() && Options.MaxNumberOfRuns) {
     Printf("ERROR: no interesting inputs were found. "
@@ -749,6 +753,7 @@ void Fuzzer::ReadAndExecuteSeedCorpora(const Vector<std::string> &CorpusDirs) {
 
 void Fuzzer::Loop(const Vector<std::string> &CorpusDirs) {
   ReadAndExecuteSeedCorpora(CorpusDirs);
+  DFT.Clear();  // No need for DFT any more.
   TPC.SetPrintNewPCs(Options.PrintNewCovPcs);
   TPC.SetPrintNewFuncs(Options.PrintNewCovFuncs);
   system_clock::time_point LastCorpusReload = system_clock::now();
diff --git a/lib/fuzzer/tests/FuzzerUnittest.cpp b/lib/fuzzer/tests/FuzzerUnittest.cpp
index 0b8673876..1b3a0934a 100644
--- a/lib/fuzzer/tests/FuzzerUnittest.cpp
+++ b/lib/fuzzer/tests/FuzzerUnittest.cpp
@@ -582,12 +582,13 @@ TEST(FuzzerUtil, Base64) {
 }
 
 TEST(Corpus, Distribution) {
+  DataFlowTrace DFT;
   Random Rand(0);
   std::unique_ptr<InputCorpus> C(new InputCorpus(""));
   size_t N = 10;
   size_t TriesPerUnit = 1<<16;
   for (size_t i = 0; i < N; i++)
-    C->AddToCorpus(Unit{ static_cast<uint8_t>(i) }, 1, false, false, {});
+    C->AddToCorpus(Unit{ static_cast<uint8_t>(i) }, 1, false, false, {}, DFT);
 
   Vector<size_t> Hist(N);
   for (size_t i = 0; i < N * TriesPerUnit; i++) {