[libFuzzer] reimplement DFT's collect_data_flow inside libFuzzer so that we don't need external python scripts

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

2 files changed, 146 insertions, 35 deletions

diff --git a/lib/fuzzer/FuzzerDataFlowTrace.cpp b/lib/fuzzer/FuzzerDataFlowTrace.cpp
index ac54fcf4b..bfcdf0061 100644
--- a/lib/fuzzer/FuzzerDataFlowTrace.cpp
+++ b/lib/fuzzer/FuzzerDataFlowTrace.cpp
@@ -9,14 +9,21 @@
 //===----------------------------------------------------------------------===//
 
 #include "FuzzerDataFlowTrace.h"
+
+#include "FuzzerCommand.h"
 #include "FuzzerIO.h"
 #include "FuzzerRandom.h"
+#include "FuzzerSHA1.h"
+#include "FuzzerUtil.h"
 
 #include <cstdlib>
 #include <fstream>
 #include <numeric>
+#include <queue>
 #include <sstream>
 #include <string>
+#include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 namespace fuzzer {
@@ -78,6 +85,7 @@ Vector<double> BlockCoverage::FunctionWeights(size_t NumFunctions) const {
   for (auto It : Functions) {
     auto FunctionID = It.first;
     auto Counters = It.second;
+    assert(FunctionID < NumFunctions);
     auto &Weight = Res[FunctionID];
     Weight = 1000.;  // this function is covered.
     Weight /= SmallestNonZeroCounter(Counters);
@@ -97,10 +105,57 @@ void DataFlowTrace::ReadCoverage(const std::string &DirPath) {
   }
 }
 
-void DataFlowTrace::Init(const std::string &DirPath,
+static void DFTStringAppendToVector(Vector<uint8_t> *DFT,
+                                    const std::string_view DFTString) {
+  assert(DFT->size() == DFTString.size());
+  for (size_t I = 0, Len = DFT->size(); I < Len; I++)
+    (*DFT)[I] = DFTString[I] == '1';
+}
+
+// converts a string of '0' and '1' into a Vector<uint8_t>
+static Vector<uint8_t> DFTStringToVector(const std::string_view DFTString) {
+  Vector<uint8_t> DFT(DFTString.size());
+  DFTStringAppendToVector(&DFT, DFTString);
+  return DFT;
+}
+
+static std::ostream &operator<<(std::ostream &OS, const Vector<uint8_t> &DFT) {
+  for (auto B : DFT)
+    OS << (B ? "1" : "0");
+  return OS;
+}
+
+static bool ParseError(const char *Err, const std::string &Line) {
+  Printf("DataFlowTrace: parse error: %s: Line: %s\n", Err, Line.c_str());
+  return false;
+};
+
+static bool ParseDFTLine(const std::string &Line, size_t *FunctionNum,
+                         std::string_view *DFTString) {
+  if (!Line.empty() && Line[0] != 'F')
+    return false; // Ignore coverage.
+  size_t SpacePos = Line.find(' ');
+  if (SpacePos == std::string::npos)
+    return ParseError("no space in the trace line", Line);
+  if (Line.empty() || Line[0] != 'F')
+    return ParseError("the trace line doesn't start with 'F'", Line);
+  *FunctionNum = std::atol(Line.c_str() + 1);
+  const char *Beg = Line.c_str() + SpacePos + 1;
+  const char *End = Line.c_str() + Line.size();
+  assert(Beg < End);
+  size_t Len = End - Beg;
+  for (size_t I = 0; I < Len; I++) {
+    if (Beg[I] != '0' && Beg[I] != '1')
+      return ParseError("the trace should contain only 0 or 1", Line);
+  }
+  *DFTString = Beg;
+  return true;
+}
+
+bool DataFlowTrace::Init(const std::string &DirPath,
                          std::string *FocusFunction,
                          Random &Rand) {
-  if (DirPath.empty()) return;
+  if (DirPath.empty()) return false;
   Printf("INFO: DataFlowTrace: reading from '%s'\n", DirPath.c_str());
   Vector<SizedFile> Files;
   GetSizedFilesFromDir(DirPath, &Files);
@@ -144,7 +199,7 @@ void DataFlowTrace::Init(const std::string &DirPath,
   }
 
   if (!NumFunctions || FocusFuncIdx == SIZE_MAX || Files.size() <= 1)
-    return;
+    return false;
 
   // Read traces.
   size_t NumTraceFiles = 0;
@@ -152,41 +207,23 @@ void DataFlowTrace::Init(const std::string &DirPath,
   for (auto &SF : Files) {
     auto Name = Basename(SF.File);
     if (Name == kFunctionsTxt) continue;
-    auto ParseError = [&](const char *Err) {
-      Printf("DataFlowTrace: parse error: %s\n  File: %s\n  Line: %s\n", Err,
-             Name.c_str(), L.c_str());
-    };
     NumTraceFiles++;
     // Printf("=== %s\n", Name.c_str());
     std::ifstream IF(SF.File);
     while (std::getline(IF, L, '\n')) {
-      if (!L.empty() && L[0] == 'C')
-        continue; // Ignore coverage.
-      size_t SpacePos = L.find(' ');
-      if (SpacePos == std::string::npos)
-        return ParseError("no space in the trace line");
-      if (L.empty() || L[0] != 'F')
-        return ParseError("the trace line doesn't start with 'F'");
-      size_t N = std::atol(L.c_str() + 1);
-      if (N >= NumFunctions)
-        return ParseError("N is greater than the number of functions");
-      if (N == FocusFuncIdx) {
+      size_t FunctionNum = 0;
+      std::string_view DFTString;
+      if (ParseDFTLine(L, &FunctionNum, &DFTString) &&
+          FunctionNum == FocusFuncIdx) {
         NumTracesWithFocusFunction++;
-        const char *Beg = L.c_str() + SpacePos + 1;
-        const char *End = L.c_str() + L.size();
-        assert(Beg < End);
-        size_t Len = End - Beg;
-        Vector<uint8_t> V(Len);
-        for (size_t I = 0; I < Len; I++) {
-          if (Beg[I] != '0' && Beg[I] != '1')
-            ParseError("the trace should contain only 0 or 1");
-          V[I] = Beg[I] == '1';
-        }
-        Traces[Name] = V;
+
+        if (FunctionNum >= NumFunctions)
+          return ParseError("N is greater than the number of functions", L);
+        Traces[Name] = DFTStringToVector(DFTString);
         // Print just a few small traces.
-        if (NumTracesWithFocusFunction <= 3 && Len <= 16)
-          Printf("%s => |%s|\n", Name.c_str(), L.c_str() + SpacePos + 1);
-        break;  // No need to parse the following lines.
+        if (NumTracesWithFocusFunction <= 3 && DFTString.size() <= 16)
+          Printf("%s => |%s|\n", Name.c_str(), std::string(DFTString).c_str());
+        break; // No need to parse the following lines.
       }
     }
   }
@@ -194,13 +231,87 @@ void DataFlowTrace::Init(const std::string &DirPath,
   Printf("INFO: DataFlowTrace: %zd trace files, %zd functions, "
          "%zd traces with focus function\n",
          NumTraceFiles, NumFunctions, NumTracesWithFocusFunction);
+  return true;
 }
 
 int CollectDataFlow(const std::string &DFTBinary, const std::string &DirPath,
                     const Vector<SizedFile> &CorporaFiles) {
-  Printf("INFO: collecting data flow for %zd files\n", CorporaFiles.size());
+  Printf("INFO: collecting data flow: bin: %s dir: %s files: %zd\n",
+         DFTBinary.c_str(), DirPath.c_str(), CorporaFiles.size());
+  MkDir(DirPath);
+  auto Temp = TempPath(".dft");
+  for (auto &F : CorporaFiles) {
+    // For every input F we need to collect the data flow and the coverage.
+    // Data flow collection may fail if we request too many DFSan tags at once.
+    // So, we start from requesting all tags in range [0,Size) and if that fails
+    // we then request tags in [0,Size/2) and [Size/2, Size), and so on.
+    // Function number => DFT.
+    std::unordered_map<size_t, Vector<uint8_t>> DFTMap;
+    std::unordered_set<std::string> Cov;
+    std::queue<std::pair<size_t, size_t>> Q;
+    Q.push({0, F.Size});
+    while (!Q.empty()) {
+      auto R = Q.front();
+      Printf("\n\n\n********* Trying: [%zd, %zd)\n", R.first, R.second);
+      Q.pop();
+      Command Cmd;
+      Cmd.addArgument(DFTBinary);
+      Cmd.addArgument(std::to_string(R.first));
+      Cmd.addArgument(std::to_string(R.second));
+      Cmd.addArgument(F.File);
+      Cmd.addArgument(Temp);
+      Printf("CMD: %s\n", Cmd.toString().c_str());
+      if (ExecuteCommand(Cmd)) {
+        // DFSan has failed, collect tags for two subsets.
+        if (R.second - R.first >= 2) {
+          size_t Mid = (R.second + R.first) / 2;
+          Q.push({R.first, Mid});
+          Q.push({Mid, R.second});
+        }
+      } else {
+        Printf("********* Success: [%zd, %zd)\n", R.first, R.second);
+        std::ifstream IF(Temp);
+        std::string L;
+        while (std::getline(IF, L, '\n')) {
+          // Data flow collection has succeeded.
+          // Merge the results with the other runs.
+          if (L.empty()) continue;
+          if (L[0] == 'C') {
+            // Take coverage lines as is, they will be the same in all attempts.
+            Cov.insert(L);
+          } else if (L[0] == 'F') {
+            size_t FunctionNum = 0;
+            std::string_view DFTString;
+            if (ParseDFTLine(L, &FunctionNum, &DFTString)) {
+              auto &DFT = DFTMap[FunctionNum];
+              if (DFT.empty()) {
+                // Haven't seen this function before, take DFT as is.
+                DFT = DFTStringToVector(DFTString);
+              } else if (DFT.size() == DFTString.size()) {
+                // Have seen this function already, merge DFTs.
+                DFTStringAppendToVector(&DFT, DFTString);
+              }
+            }
+          }
+        }
+      }
+    }
+    auto OutPath = DirPlusFile(DirPath, Hash(FileToVector(F.File)));
+    // Dump combined DFT to disk.
+    Printf("Producing DFT for %s\n", OutPath.c_str());
+    std::ofstream OF(OutPath);
+    for (auto &DFT: DFTMap)
+      OF << "F" << DFT.first << " " << DFT.second << std::endl;
+    for (auto &C : Cov)
+      OF << C << std::endl;
+  }
+  RemoveFile(Temp);
+  // Write functions.txt.
+  Command Cmd;
+  Cmd.addArgument(DFTBinary);
+  Cmd.setOutputFile(DirPlusFile(DirPath, "functions.txt"));
+  ExecuteCommand(Cmd);
   return 0;
 }
 
 }  // namespace fuzzer
-
diff --git a/lib/fuzzer/FuzzerDataFlowTrace.h b/lib/fuzzer/FuzzerDataFlowTrace.h
index 4b80de7da..cfb04ad3a 100644
--- a/lib/fuzzer/FuzzerDataFlowTrace.h
+++ b/lib/fuzzer/FuzzerDataFlowTrace.h
@@ -111,7 +111,7 @@ class BlockCoverage {
 class DataFlowTrace {
  public:
   void ReadCoverage(const std::string &DirPath);
-  void Init(const std::string &DirPath, std::string *FocusFunction,
+  bool Init(const std::string &DirPath, std::string *FocusFunction,
             Random &Rand);
   void Clear() { Traces.clear(); }
   const Vector<uint8_t> *Get(const std::string &InputSha1) const {