From be7a35769dbc850961818c405ef82189bda7225e Mon Sep 17 00:00:00 2001
From: Kostya Serebryany <kcc@google.com>
Date: Fri, 4 Aug 2017 23:13:58 +0000
Subject: [PATCH] [libFuzzer] print PCs using the in-binary PC-table instead of
 relying on PCs captured at run-time

llvm-svn: 310148
---
 llvm/lib/Fuzzer/FuzzerLoop.cpp    |  5 ++-
 llvm/lib/Fuzzer/FuzzerTracePC.cpp | 51 ++++++++++++++++++++++---------
 llvm/lib/Fuzzer/FuzzerTracePC.h   |  5 ++-
 3 files changed, 40 insertions(+), 21 deletions(-)

diff --git a/llvm/lib/Fuzzer/FuzzerLoop.cpp b/llvm/lib/Fuzzer/FuzzerLoop.cpp
index 682e4b7b05db..41fd213a6531 100644
--- a/llvm/lib/Fuzzer/FuzzerLoop.cpp
+++ b/llvm/lib/Fuzzer/FuzzerLoop.cpp
@@ -122,7 +122,6 @@ Fuzzer::Fuzzer(UserCallback CB, InputCorpus &Corpus, MutationDispatcher &MD,
     EF->__sanitizer_install_malloc_and_free_hooks(MallocHook, FreeHook);
   TPC.SetUseCounters(Options.UseCounters);
   TPC.SetUseValueProfile(Options.UseValueProfile);
-  TPC.SetPrintNewPCs(Options.PrintNewCovPcs);
 
   if (Options.Verbosity)
     TPC.PrintModuleInfo();
@@ -438,6 +437,7 @@ bool Fuzzer::RunOne(const uint8_t *Data, size_t Size, bool MayDeleteFile,
   PrintPulseAndReportSlowInput(Data, Size);
   size_t NumNewFeatures = Corpus.NumFeatureUpdates() - NumUpdatesBefore;
   if (NumNewFeatures) {
+    TPC.UpdateObservedPCs();
     Corpus.AddToCorpus({Data, Data + Size}, NumNewFeatures, MayDeleteFile,
                        UniqFeatureSetTmp);
     return true;
@@ -546,7 +546,6 @@ void Fuzzer::ReportNewCoverage(InputInfo *II, const Unit &U) {
                                          "NEW   ");
   WriteToOutputCorpus(U);
   NumberOfNewUnitsAdded++;
-  TPC.PrintNewPCs();
   CheckExitOnSrcPosOrItem();  // Check only after the unit is saved to corpus.
   LastCorpusUpdateRun = TotalNumberOfRuns;
   LastCorpusUpdateTime = system_clock::now();
@@ -626,7 +625,7 @@ void Fuzzer::MutateAndTestOne() {
 }
 
 void Fuzzer::Loop() {
-  TPC.InitializePrintNewPCs();
+  TPC.SetPrintNewPCs(Options.PrintNewCovPcs);
   system_clock::time_point LastCorpusReload = system_clock::now();
   if (Options.DoCrossOver)
     MD.SetCorpus(&Corpus);
diff --git a/llvm/lib/Fuzzer/FuzzerTracePC.cpp b/llvm/lib/Fuzzer/FuzzerTracePC.cpp
index 1cfd3f3ecff1..45deaa21a4ce 100644
--- a/llvm/lib/Fuzzer/FuzzerTracePC.cpp
+++ b/llvm/lib/Fuzzer/FuzzerTracePC.cpp
@@ -48,6 +48,8 @@ uintptr_t *TracePC::PCs() const {
 }
 
 size_t TracePC::GetTotalPCCoverage() {
+  if (ObservedPCs)
+    return ObservedPCs->size();
   size_t Res = 0;
   for (size_t i = 1, N = GetNumPCs(); i < N; i++)
     if (PCs()[i])
@@ -136,21 +138,40 @@ void TracePC::HandleCallerCallee(uintptr_t Caller, uintptr_t Callee) {
   ValueProfileMap.AddValueModPrime(Idx);
 }
 
-void TracePC::InitializePrintNewPCs() {
-  if (!DoPrintNewPCs) return;
-  assert(!PrintedPCs);
-  PrintedPCs = new std::set<uintptr_t>;
-  for (size_t i = 1; i < GetNumPCs(); i++)
-    if (PCs()[i])
-      PrintedPCs->insert(PCs()[i]);
-}
-
-void TracePC::PrintNewPCs() {
-  if (!DoPrintNewPCs) return;
-  assert(PrintedPCs);
-  for (size_t i = 1; i < GetNumPCs(); i++)
-    if (PCs()[i] && PrintedPCs->insert(PCs()[i]).second)
-      PrintPC("\tNEW_PC: %p %F %L\n", "\tNEW_PC: %p\n", PCs()[i]);
+void TracePC::UpdateObservedPCs() {
+  if (NumPCsInPCTables) {
+    auto Observe = [&](uintptr_t PC) {
+      bool Inserted = ObservedPCs->insert(PC).second;
+      if (Inserted && DoPrintNewPCs)
+        PrintPC("\tNEW_PC: %p %F %L\n", "\tNEW_PC: %p\n", PC + 1);
+    };
+
+    if (!ObservedPCs)
+      ObservedPCs = new std::set<uintptr_t>;
+
+    if (NumInline8bitCounters == NumPCsInPCTables) {
+      for (size_t i = 0; i < NumModulesWithInline8bitCounters; i++) {
+        uint8_t *Beg = ModuleCounters[i].Start;
+        size_t Size = ModuleCounters[i].Stop - Beg;
+        assert(Size ==
+               (size_t)(ModulePCTable[i].Stop - ModulePCTable[i].Start));
+        for (size_t j = 0; j < Size; j++)
+          if (Beg[j])
+            Observe(ModulePCTable[i].Start[j]);
+      }
+    } else if (NumGuards == NumPCsInPCTables) {
+      size_t GuardIdx = 1;
+      for (size_t i = 0; i < NumModules; i++) {
+        uint32_t *Beg = Modules[i].Start;
+        size_t Size = Modules[i].Stop - Beg;
+        assert(Size ==
+               (size_t)(ModulePCTable[i].Stop - ModulePCTable[i].Start));
+        for (size_t j = 0; j < Size; j++, GuardIdx++)
+          if (Counters()[GuardIdx])
+            Observe(ModulePCTable[i].Start[j]);
+      }
+    }
+  }
 }
 
 void TracePC::PrintCoverage() {
diff --git a/llvm/lib/Fuzzer/FuzzerTracePC.h b/llvm/lib/Fuzzer/FuzzerTracePC.h
index ba882b3af867..ad832d7b2d40 100644
--- a/llvm/lib/Fuzzer/FuzzerTracePC.h
+++ b/llvm/lib/Fuzzer/FuzzerTracePC.h
@@ -82,6 +82,7 @@ class TracePC {
   void SetUseCounters(bool UC) { UseCounters = UC; }
   void SetUseValueProfile(bool VP) { UseValueProfile = VP; }
   void SetPrintNewPCs(bool P) { DoPrintNewPCs = P; }
+  void UpdateObservedPCs();
   template <class Callback> void CollectFeatures(Callback CB) const;
 
   void ResetMaps() {
@@ -110,8 +111,6 @@ class TracePC {
   TableOfRecentCompares<Word, 32> TORCW;
   MemMemTable<1024> MMT;
 
-  void PrintNewPCs();
-  void InitializePrintNewPCs();
   size_t GetNumPCs() const {
     return NumGuards == 0 ? (1 << kTracePcBits) : Min(kNumPCs, NumGuards + 1);
   }
@@ -158,7 +157,7 @@ private:
   uint8_t *Counters() const;
   uintptr_t *PCs() const;
 
-  std::set<uintptr_t> *PrintedPCs;
+  std::set<uintptr_t> *ObservedPCs;
 
   ValueBitMap ValueProfileMap;
   uintptr_t InitialStack, LowestStack;  // Assume stack grows down.
-- 
2.34.1