[libFuzzer] add an experimental flag -focus_function: libFuzzer will try to focus on inputs that trigger that function

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

10 files changed, 96 insertions, 2 deletions

diff --git a/lib/fuzzer/FuzzerCorpus.h b/lib/fuzzer/FuzzerCorpus.h
index 2da929835..05d527ba0 100644
--- a/lib/fuzzer/FuzzerCorpus.h
+++ b/lib/fuzzer/FuzzerCorpus.h
@@ -35,6 +35,7 @@ struct InputInfo {
   size_t NumSuccessfullMutations = 0;
   bool MayDeleteFile = false;
   bool Reduced = false;
+  bool HasFocusFunction = false;
   Vector<uint32_t> UniqFeatureSet;
   float FeatureFrequencyScore = 1.0;
 };
@@ -70,10 +71,17 @@ class InputCorpus {
         Res = std::max(Res, II->U.size());
     return Res;
   }
+
+  size_t NumInputsThatTouchFocusFunction() {
+    return std::count_if(Inputs.begin(), Inputs.end(), [](const InputInfo *II) {
+      return II->HasFocusFunction;
+    });
+  }
+
   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,
-                   const Vector<uint32_t> &FeatureSet) {
+                   bool HasFocusFunction, const Vector<uint32_t> &FeatureSet) {
     assert(!U.empty());
     if (FeatureDebug)
       Printf("ADD_TO_CORPUS %zd NF %zd\n", Inputs.size(), NumFeatures);
@@ -83,6 +91,7 @@ class InputCorpus {
     II.NumFeatures = NumFeatures;
     II.MayDeleteFile = MayDeleteFile;
     II.UniqFeatureSet = FeatureSet;
+    II.HasFocusFunction = HasFocusFunction;
     std::sort(II.UniqFeatureSet.begin(), II.UniqFeatureSet.end());
     ComputeSHA1(U.data(), U.size(), II.Sha1);
     Hashes.insert(Sha1ToString(II.Sha1));
@@ -265,6 +274,7 @@ private:
     for (size_t i = 0; i < N; i++)
       Weights[i] = Inputs[i]->NumFeatures
                        ? (i + 1) * Inputs[i]->FeatureFrequencyScore
+                       * (Inputs[i]->HasFocusFunction ? 1000 : 1)
                        : 0.;
     if (FeatureDebug) {
       for (size_t i = 0; i < N; i++)
diff --git a/lib/fuzzer/FuzzerDriver.cpp b/lib/fuzzer/FuzzerDriver.cpp
index 4f14373b3..26e5548a9 100644
--- a/lib/fuzzer/FuzzerDriver.cpp
+++ b/lib/fuzzer/FuzzerDriver.cpp
@@ -620,6 +620,8 @@ int FuzzerDriver(int *argc, char ***argv, UserCallback Callback) {
     Options.ExitOnSrcPos = Flags.exit_on_src_pos;
   if (Flags.exit_on_item)
     Options.ExitOnItem = Flags.exit_on_item;
+  if (Flags.focus_function)
+    Options.FocusFunction = Flags.focus_function;
 
   unsigned Seed = Flags.seed;
   // Initialize Seed.
diff --git a/lib/fuzzer/FuzzerFlags.def b/lib/fuzzer/FuzzerFlags.def
index 64bd36281..1ff3fd95e 100644
--- a/lib/fuzzer/FuzzerFlags.def
+++ b/lib/fuzzer/FuzzerFlags.def
@@ -143,6 +143,8 @@ FUZZER_FLAG_STRING(exit_on_item, "Exit if an item with a given sha1 sum"
 FUZZER_FLAG_INT(ignore_remaining_args, 0, "If 1, ignore all arguments passed "
                 "after this one. Useful for fuzzers that need to do their own "
                 "argument parsing.")
+FUZZER_FLAG_STRING(focus_function, "Experimental. "
+     "Fuzzing will focus on inputs that trigger calls to this function")
 
 FUZZER_DEPRECATED_FLAG(run_equivalence_server)
 FUZZER_DEPRECATED_FLAG(use_equivalence_server)
diff --git a/lib/fuzzer/FuzzerLoop.cpp b/lib/fuzzer/FuzzerLoop.cpp
index dfa6cf38b..4bf5c7802 100644
--- a/lib/fuzzer/FuzzerLoop.cpp
+++ b/lib/fuzzer/FuzzerLoop.cpp
@@ -159,6 +159,7 @@ Fuzzer::Fuzzer(UserCallback CB, InputCorpus &Corpus, MutationDispatcher &MD,
   AllocateCurrentUnitData();
   CurrentUnitSize = 0;
   memset(BaseSha1, 0, sizeof(BaseSha1));
+  TPC.SetFocusFunction(Options.FocusFunction);
 }
 
 Fuzzer::~Fuzzer() {}
@@ -333,6 +334,8 @@ void Fuzzer::PrintStats(const char *Where, const char *End, size_t Units) {
       else
         Printf("/%zdMb", N >> 20);
     }
+    if (size_t FF = Corpus.NumInputsThatTouchFocusFunction())
+      Printf(" focus: %zd", FF);
   }
   if (TmpMaxMutationLen)
     Printf(" lim: %zd", TmpMaxMutationLen);
@@ -464,6 +467,7 @@ bool Fuzzer::RunOne(const uint8_t *Data, size_t Size, bool MayDeleteFile,
   if (NumNewFeatures) {
     TPC.UpdateObservedPCs();
     Corpus.AddToCorpus({Data, Data + Size}, NumNewFeatures, MayDeleteFile,
+                       TPC.ObservedFocusFunction(),
                        UniqFeatureSetTmp);
     return true;
   }
@@ -733,6 +737,10 @@ void Fuzzer::ReadAndExecuteSeedCorpora(const Vector<std::string> &CorpusDirs) {
   }
 
   PrintStats("INITED");
+  if (!Options.FocusFunction.empty())
+    Printf("INFO: %zd/%zd inputs touch the focus function\n",
+           Corpus.NumInputsThatTouchFocusFunction(), Corpus.size());
+
   if (Corpus.empty()) {
     Printf("ERROR: no interesting inputs were found. "
            "Is the code instrumented for coverage? Exiting.\n");
diff --git a/lib/fuzzer/FuzzerOptions.h b/lib/fuzzer/FuzzerOptions.h
index 4fd4e8d35..0c51d9e39 100644
--- a/lib/fuzzer/FuzzerOptions.h
+++ b/lib/fuzzer/FuzzerOptions.h
@@ -45,6 +45,7 @@ struct FuzzingOptions {
   std::string ExactArtifactPath;
   std::string ExitOnSrcPos;
   std::string ExitOnItem;
+  std::string FocusFunction;
   bool SaveArtifacts = true;
   bool PrintNEW = true; // Print a status line when new units are found;
   bool PrintNewCovPcs = false;
diff --git a/lib/fuzzer/FuzzerTracePC.cpp b/lib/fuzzer/FuzzerTracePC.cpp
index 20230d496..ed62cdcf4 100644
--- a/lib/fuzzer/FuzzerTracePC.cpp
+++ b/lib/fuzzer/FuzzerTracePC.cpp
@@ -229,6 +229,39 @@ void TracePC::IterateCoveredFunctions(CallBack CB) {
   }
 }
 
+void TracePC::SetFocusFunction(const std::string &FuncName) {
+  // This function should be called once.
+  assert(FocusFunction.first > NumModulesWithInline8bitCounters);
+  if (FuncName.empty())
+    return;
+  for (size_t M = 0; M < NumModulesWithInline8bitCounters; M++) {
+    auto &PCTE = ModulePCTable[M];
+    size_t N = PCTE.Stop - PCTE.Start;
+    for (size_t I = 0; I < N; I++) {
+      if (!(PCTE.Start[I].PCFlags & 1)) continue;  // not a function entry.
+      auto Name = DescribePC("%F", GetNextInstructionPc(PCTE.Start[I].PC));
+      if (Name[0] == 'i' && Name[1] == 'n' && Name[2] == ' ')
+        Name = Name.substr(3, std::string::npos);
+      if (FuncName != Name) continue;
+      Printf("INFO: Focus function is set to '%s'\n", Name.c_str());
+      FocusFunction = {M, I};
+      return;
+    }
+  }
+}
+
+bool TracePC::ObservedFocusFunction() {
+  size_t I = FocusFunction.first;
+  size_t J = FocusFunction.second;
+  if (I >= NumModulesWithInline8bitCounters)
+    return false;
+  auto &MC = ModuleCounters[I];
+  size_t Size = MC.Stop - MC.Start;
+  if (J >= Size)
+    return false;
+  return MC.Start[J] != 0;
+}
+
 void TracePC::PrintCoverage() {
   if (!EF->__sanitizer_symbolize_pc ||
       !EF->__sanitizer_get_module_and_offset_for_pc) {
diff --git a/lib/fuzzer/FuzzerTracePC.h b/lib/fuzzer/FuzzerTracePC.h
index d68da76b2..e1db5127c 100644
--- a/lib/fuzzer/FuzzerTracePC.h
+++ b/lib/fuzzer/FuzzerTracePC.h
@@ -131,6 +131,9 @@ class TracePC {
       CB(PC);
   }
 
+  void SetFocusFunction(const std::string &FuncName);
+  bool ObservedFocusFunction();
+
 private:
   bool UseCounters = false;
   bool UseValueProfile = false;
@@ -163,6 +166,9 @@ private:
   Set<uintptr_t> ObservedPCs;
   Set<uintptr_t> ObservedFuncs;
 
+  std::pair<size_t, size_t> FocusFunction = {-1, -1};  // Module and PC IDs.
+
+
   ValueBitMap ValueProfileMap;
   uintptr_t InitialStack;
 };
diff --git a/lib/fuzzer/tests/FuzzerUnittest.cpp b/lib/fuzzer/tests/FuzzerUnittest.cpp
index c795eddc6..a38a45344 100644
--- a/lib/fuzzer/tests/FuzzerUnittest.cpp
+++ b/lib/fuzzer/tests/FuzzerUnittest.cpp
@@ -579,7 +579,7 @@ TEST(Corpus, Distribution) {
   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, {});
+    C->AddToCorpus(Unit{ static_cast<uint8_t>(i) }, 1, false, false, {});
 
   Vector<size_t> Hist(N);
   for (size_t i = 0; i < N * TriesPerUnit; i++) {
diff --git a/test/fuzzer/OnlySomeBytesTest.cpp b/test/fuzzer/OnlySomeBytesTest.cpp
index 05793f0ab..3873b710b 100644
--- a/test/fuzzer/OnlySomeBytesTest.cpp
+++ b/test/fuzzer/OnlySomeBytesTest.cpp
@@ -12,6 +12,7 @@
 const size_t N = 2048;
 typedef const uint8_t *IN;
 
+extern "C" {
 __attribute__((noinline)) void bad() {
   fprintf(stderr, "BINGO\n");
   abort();
@@ -27,6 +28,8 @@ __attribute__((noinline)) void fC(IN in) { if (in[2] == 'C') f0(in); }
 __attribute__((noinline)) void fB(IN in) { if (in[1] == 'B') fC(in); }
 __attribute__((noinline)) void fA(IN in) { if (in[0] == 'A') fB(in); }
 
+} // extern "C"
+
 extern "C" int LLVMFuzzerTestOneInput(const uint8_t *Data, size_t Size) {
   if (Size < N) return 0;
   fA((IN)Data);
diff --git a/test/fuzzer/target-function.test b/test/fuzzer/target-function.test
new file mode 100644
index 000000000..f39e1ac1b
--- /dev/null
+++ b/test/fuzzer/target-function.test
@@ -0,0 +1,29 @@
+# Tests -focus_function
+#
+# TODO: don't require linux.
+# REQUIRES: linux
+#
+RUN: %cpp_compiler %S/OnlySomeBytesTest.cpp -o %t-exe
+
+RUN: %t-exe -runs=100 2>&1 | FileCheck %s --check-prefix=FOCUS_NONE
+FOCUS_NONE-NOT: INFO: Focus function is set to
+FOCUS_NONE-NOT: INFO: {{.*}} inputs touch the focus function
+
+RUN: %t-exe -runs=100 -focus_function=WRONG 2>&1 | FileCheck %s --check-prefix=FOCUS_WRONG
+FOCUS_WRONG-NOT: INFO: Focus function is set to
+FOCUS_WRONG: INFO: 0/1 inputs touch the focus function
+
+RUN: %t-exe -runs=100 -focus_function=f0 2>&1 | FileCheck %s --check-prefix=FOCUS_F0
+FOCUS_F0: INFO: Focus function is set to 'f0'
+FOCUS_F0: INFO: 0/1 inputs touch the focus function
+
+RUN: rm -rf %t-corpus
+RUN: mkdir %t-corpus
+# ABC triggers the focus function, others don't.
+RUN: echo ABC$(for((i=0;i<2048;i++)); do echo -n x; done) > %t-corpus/ABC
+RUN: echo AXY$(for((i=0;i<2048;i++)); do echo -n x; done) > %t-corpus/AXY
+RUN: echo ABX$(for((i=0;i<2048;i++)); do echo -n x; done) > %t-corpus/ABX
+
+RUN: %t-exe -runs=10000 -focus_function=f0 %t-corpus 2>&1 | FileCheck %s --check-prefix=CORPUS_1_3
+CORPUS_1_3: INFO: 1/3 inputs touch the focus function
+CORPUS_1_3: DONE {{.*}} focus: