[libFuzzer] add 8-bit counters to trace-pc-guard handler

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

lib/Fuzzer/FuzzerInternal.h

@@ -360,9 +360,18 @@ class TracePC {
  public:
   void HandleTrace(uint8_t *guard, uintptr_t PC);
   void HandleInit(uint8_t *start, uint8_t *stop);
-  size_t GetTotalCoverage();
+  size_t GetTotalCoverage() { return TotalCoverage; }
+  void SetUseCounters(bool UC) { UseCounters = UC; }
+  size_t UpdateCounterMap(ValueBitMap *Map);
+  void FinalizeTrace();
+
 private:
+  bool UseCounters = false;
   size_t TotalCoverage = 0;
+  size_t TotalCounterBits = 0;
+
+  uint8_t *Start, *Stop;
+  ValueBitMap CounterMap;
 };
 
 extern TracePC TPC;
@@ -380,6 +389,7 @@ public:
       CounterBitmapBits = 0;
       CounterBitmap.clear();
       VPMap.Reset();
+      TPCMap.Reset();
       VPMapBits = 0;
     }
 
@@ -390,6 +400,7 @@ public:
     // Precalculated number of bits in CounterBitmap.
     size_t CounterBitmapBits;
     std::vector<uint8_t> CounterBitmap;
+    ValueBitMap TPCMap;
     ValueBitMap VPMap;
     size_t VPMapBits;
   };

lib/Fuzzer/FuzzerLoop.cpp

@@ -77,6 +77,8 @@ void Fuzzer::PrepareCounters(Fuzzer::Coverage *C) {
 bool Fuzzer::RecordMaxCoverage(Fuzzer::Coverage *C) {
   bool Res = false;
 
+  TPC.FinalizeTrace();
+
   uint64_t NewBlockCoverage =
       EF->__sanitizer_get_total_unique_coverage() + TPC.GetTotalCoverage();
   if (NewBlockCoverage > C->BlockCoverage) {
@@ -97,11 +99,13 @@ bool Fuzzer::RecordMaxCoverage(Fuzzer::Coverage *C) {
   if (Options.UseCounters) {
     uint64_t CounterDelta =
         EF->__sanitizer_update_counter_bitset_and_clear_counters(
-            C->CounterBitmap.data());
+            C->CounterBitmap.data()) +
+        TPC.UpdateCounterMap(&C->TPCMap);
     if (CounterDelta > 0) {
       Res = true;
       C->CounterBitmapBits += CounterDelta;
     }
+
   }
 
   size_t NewVPMapBits = VPMapMergeFromCurrent(C->VPMap);
@@ -158,6 +162,7 @@ Fuzzer::Fuzzer(UserCallback CB, MutationDispatcher &MD, FuzzingOptions Options)
   IsMyThread = true;
   if (Options.DetectLeaks && EF->__sanitizer_install_malloc_and_free_hooks)
     EF->__sanitizer_install_malloc_and_free_hooks(MallocHook, FreeHook);
+  TPC.SetUseCounters(Options.UseCounters);
 
   if (Options.PrintNewCovPcs) {
     PcBufferLen = 1 << 24;

lib/Fuzzer/FuzzerTracePC.cpp

@@ -18,25 +18,60 @@ namespace fuzzer {
 
 TracePC TPC;
 
-void TracePC::HandleTrace(uint8_t *guard, uintptr_t PC) {
-  *guard = 0xff;
+void TracePC::HandleTrace(uint8_t *Guard, uintptr_t PC) {
+  if (UseCounters) {
+    uintptr_t GV = *Guard;
+    if (GV == 0)
+      TotalCoverage++;
+    if (GV < 255)
+      GV++;
+    *Guard = GV;
+  } else {
+    *Guard = 0xff;
     TotalCoverage++;
   }
-void TracePC::HandleInit(uint8_t *start, uint8_t *stop) {
-  Printf("INFO: guards: [%p,%p)\n", start, stop);
 }
-size_t TracePC::GetTotalCoverage() { return TotalCoverage; }
+
+void TracePC::HandleInit(uint8_t *Start, uint8_t *Stop) {
+  // TODO: this handles only one DSO/binary.
+  this->Start = Start;
+  this->Stop = Stop;
+}
+
+void TracePC::FinalizeTrace() {
+  if (UseCounters && TotalCoverage) {
+    for (uint8_t *X = Start; X < Stop; X++) {
+      uint8_t Value = *X;
+      size_t Idx = X - Start;
+      if (Value >= 2) {
+        unsigned Bit = 31 - __builtin_clz(Value);
+        assert(Bit < 8);
+        CounterMap.AddValue(Idx * 8 + Bit);
+      }
+      *X = 1;
+    }
+  }
+}
+
+size_t TracePC::UpdateCounterMap(ValueBitMap *Map) {
+  if (!TotalCoverage) return 0;
+  size_t NewTotalCounterBits = Map->MergeFrom(CounterMap);
+  size_t Delta = NewTotalCounterBits - TotalCounterBits;
+  TotalCounterBits = NewTotalCounterBits;
+  return Delta;
+}
 
 } // namespace fuzzer
 
 extern "C" {
 __attribute__((visibility("default")))
-void __sanitizer_cov_trace_pc_guard(uint8_t *guard) {
+void __sanitizer_cov_trace_pc_guard(uint8_t *Guard) {
   uintptr_t PC = (uintptr_t)__builtin_return_address(0);
-  fuzzer::TPC.HandleTrace(guard, PC);
+  fuzzer::TPC.HandleTrace(Guard, PC);
 }
 
 __attribute__((visibility("default")))
-void __sanitizer_cov_trace_pc_guard_init(uint8_t *start, uint8_t *stop) {
+void __sanitizer_cov_trace_pc_guard_init(uint8_t *Start, uint8_t *Stop) {
+  fuzzer::TPC.HandleInit(Start, Stop);
 }
 }

lib/Fuzzer/test/fuzzer.test

@@ -24,7 +24,13 @@ NULL_DEREF_ON_EMPTY: stat::number_of_executed_units:
 
 #not LLVMFuzzer-FullCoverageSetTest -timeout=15 -seed=1 -mutate_depth=2 -use_full_coverage_set=1 2>&1 | FileCheck %s
 
-RUN: not LLVMFuzzer-CounterTest -use_counters=1 -max_len=6 -seed=1 -timeout=15 2>&1 | FileCheck %s
+RUN: not LLVMFuzzer-CounterTest         -use_counters=1 -max_len=6 -seed=1 -timeout=15 2>&1 | FileCheck %s --check-prefix=COUNTERS
+RUN: not LLVMFuzzer-CounterTest-TracePC -use_counters=1 -max_len=6 -seed=1 -timeout=15 2>&1 | FileCheck %s --check-prefix=COUNTERS
+
+COUNTERS: INITED {{.*}} bits:
+COUNTERS: NEW {{.*}} bits: {{[1-9]*}}
+COUNTERS: NEW {{.*}} bits: {{[1-9]*}}
+COUNTERS: BINGO
 
 RUN: not LLVMFuzzer-CallerCalleeTest                     -cross_over=0 -max_len=6 -seed=1 -timeout=15 2>&1 | FileCheck %s
 # This one is flaky, may actually find the goal even w/o use_indir_calls.

lib/Fuzzer/test/trace-pc/CMakeLists.txt

@@ -5,6 +5,7 @@ set(CMAKE_CXX_FLAGS
 
 set(TracePCTests
   SimpleTest
+  CounterTest
   )
 
 foreach(Test ${TracePCTests})