1 files changed, 25 insertions, 18 deletions

diff --git a/llvm/lib/Fuzzer/FuzzerTraceState.cpp b/llvm/lib/Fuzzer/FuzzerTraceState.cpp
index 153b2e63aa3..1948b282c94 100644
--- a/llvm/lib/Fuzzer/FuzzerTraceState.cpp
+++ b/llvm/lib/Fuzzer/FuzzerTraceState.cpp
@@ -172,6 +172,10 @@ struct TraceBasedMutation {
 
 const size_t TraceBasedMutation::kMaxSize;
 
+// Declared as static globals for faster checks inside the hooks.
+static bool RecordingTraces = false;
+static bool RecordingMemcmp = false;
+
 class TraceState {
  public:
    TraceState(UserSuppliedFuzzer &USF, const Fuzzer::FuzzingOptions &Options,
@@ -206,15 +210,17 @@ class TraceState {
                           const uint8_t *DesiredData, size_t DataSize);
 
   void StartTraceRecording() {
-    if (!Options.UseTraces) return;
-    RecordingTraces = true;
+    if (!Options.UseTraces && !Options.UseMemcmp) return;
+    RecordingTraces = Options.UseTraces;
+    RecordingMemcmp = Options.UseMemcmp;
     NumMutations = 0;
     USF.GetMD().ClearAutoDictionary();
   }
 
   void StopTraceRecording() {
-    if (!RecordingTraces) return;
+    if (!RecordingTraces && !RecordingMemcmp) return;
     RecordingTraces = false;
+    RecordingMemcmp = false;
     for (size_t i = 0; i < NumMutations; i++) {
       auto &M = Mutations[i];
       Unit U(M.Data, M.Data + M.Size);
@@ -260,7 +266,6 @@ class TraceState {
     Signed >>= 16;
     return Signed == 0 || Signed == -1L;
   }
-  bool RecordingTraces = false;
   static const size_t kMaxMutations = 1 << 16;
   size_t NumMutations;
   TraceBasedMutation Mutations[kMaxMutations];
@@ -320,7 +325,7 @@ void TraceState::DFSanMemcmpCallback(size_t CmpSize, const uint8_t *Data1,
                                      dfsan_label L2) {
 
   assert(ReallyHaveDFSan());
-  if (!RecordingTraces || !IsMyThread) return;
+  if (!RecordingMemcmp || !IsMyThread) return;
   if (L1 == 0 && L2 == 0)
     return;  // Not actionable.
   if (L1 != 0 && L2 != 0)
@@ -417,7 +422,7 @@ void TraceState::TraceCmpCallback(uintptr_t PC, size_t CmpSize, size_t CmpType,
 
 void TraceState::TraceMemcmpCallback(size_t CmpSize, const uint8_t *Data1,
                                      const uint8_t *Data2) {
-  if (!RecordingTraces || !IsMyThread) return;
+  if (!RecordingMemcmp || !IsMyThread) return;
   CmpSize = std::min(CmpSize, TraceBasedMutation::kMaxSize);
   int Added2 = TryToAddDesiredData(Data1, Data2, CmpSize);
   int Added1 = TryToAddDesiredData(Data2, Data1, CmpSize);
@@ -462,14 +467,14 @@ void Fuzzer::StopTraceRecording() {
 }
 
 void Fuzzer::AssignTaintLabels(uint8_t *Data, size_t Size) {
-  if (!Options.UseTraces) return;
+  if (!Options.UseTraces && !Options.UseMemcmp) return;
   if (!ReallyHaveDFSan()) return;
   for (size_t i = 0; i < Size; i++)
     dfsan_set_label(i + 1, &Data[i], 1);
 }
 
 void Fuzzer::InitializeTraceState() {
-  if (!Options.UseTraces) return;
+  if (!Options.UseTraces && !Options.UseMemcmp) return;
   TS = new TraceState(USF, Options, &CurrentUnitData, &CurrentUnitSize);
   if (ReallyHaveDFSan()) {
     for (size_t i = 0; i < static_cast<size_t>(Options.MaxLen); i++) {
@@ -492,12 +497,14 @@ static size_t InternalStrnlen(const char *S, size_t MaxLen) {
 }  // namespace fuzzer
 
 using fuzzer::TS;
+using fuzzer::RecordingTraces;
+using fuzzer::RecordingMemcmp;
 
 extern "C" {
 void __dfsw___sanitizer_cov_trace_cmp(uint64_t SizeAndType, uint64_t Arg1,
                                       uint64_t Arg2, dfsan_label L0,
                                       dfsan_label L1, dfsan_label L2) {
-  if (!TS) return;
+  if (!RecordingTraces) return;
   assert(L0 == 0);
   uintptr_t PC = reinterpret_cast<uintptr_t>(__builtin_return_address(0));
   uint64_t CmpSize = (SizeAndType >> 32) / 8;
@@ -507,7 +514,7 @@ void __dfsw___sanitizer_cov_trace_cmp(uint64_t SizeAndType, uint64_t Arg1,
 
 void __dfsw___sanitizer_cov_trace_switch(uint64_t Val, uint64_t *Cases,
                                          dfsan_label L1, dfsan_label L2) {
-  if (!TS) return;
+  if (!RecordingTraces) return;
   uintptr_t PC = reinterpret_cast<uintptr_t>(__builtin_return_address(0));
   TS->DFSanSwitchCallback(PC, Cases[1], Val, Cases[0], Cases+2, L1);
 }
@@ -515,7 +522,7 @@ void __dfsw___sanitizer_cov_trace_switch(uint64_t Val, uint64_t *Cases,
 void dfsan_weak_hook_memcmp(void *caller_pc, const void *s1, const void *s2,
                             size_t n, dfsan_label s1_label,
                             dfsan_label s2_label, dfsan_label n_label) {
-  if (!TS) return;
+  if (!RecordingMemcmp) return;
   dfsan_label L1 = dfsan_read_label(s1, n);
   dfsan_label L2 = dfsan_read_label(s2, n);
   TS->DFSanMemcmpCallback(n, reinterpret_cast<const uint8_t *>(s1),
@@ -525,7 +532,7 @@ void dfsan_weak_hook_memcmp(void *caller_pc, const void *s1, const void *s2,
 void dfsan_weak_hook_strncmp(void *caller_pc, const char *s1, const char *s2,
                              size_t n, dfsan_label s1_label,
                              dfsan_label s2_label, dfsan_label n_label) {
-  if (!TS) return;
+  if (!RecordingMemcmp) return;
   n = std::min(n, fuzzer::InternalStrnlen(s1, n));
   n = std::min(n, fuzzer::InternalStrnlen(s2, n));
   dfsan_label L1 = dfsan_read_label(s1, n);
@@ -536,7 +543,7 @@ void dfsan_weak_hook_strncmp(void *caller_pc, const char *s1, const char *s2,
 
 void dfsan_weak_hook_strcmp(void *caller_pc, const char *s1, const char *s2,
                             dfsan_label s1_label, dfsan_label s2_label) {
-  if (!TS) return;
+  if (!RecordingMemcmp) return;
   size_t Len1 = strlen(s1);
   size_t Len2 = strlen(s2);
   size_t N = std::min(Len1, Len2);
@@ -555,7 +562,7 @@ void dfsan_weak_hook_strcmp(void *caller_pc, const char *s1, const char *s2,
 #if LLVM_FUZZER_DEFINES_SANITIZER_WEAK_HOOOKS
 void __sanitizer_weak_hook_memcmp(void *caller_pc, const void *s1,
                                   const void *s2, size_t n, int result) {
-  if (!TS) return;
+  if (!RecordingMemcmp) return;
   if (result == 0) return;  // No reason to mutate.
   if (n <= 1) return;  // Not interesting.
   TS->TraceMemcmpCallback(n, reinterpret_cast<const uint8_t *>(s1),
@@ -564,7 +571,7 @@ void __sanitizer_weak_hook_memcmp(void *caller_pc, const void *s1,
 
 void __sanitizer_weak_hook_strncmp(void *caller_pc, const char *s1,
                                    const char *s2, size_t n, int result) {
-  if (!TS) return;
+  if (!RecordingMemcmp) return;
   if (result == 0) return;  // No reason to mutate.
   size_t Len1 = fuzzer::InternalStrnlen(s1, n);
   size_t Len2 = fuzzer::InternalStrnlen(s2, n);
@@ -577,7 +584,7 @@ void __sanitizer_weak_hook_strncmp(void *caller_pc, const char *s1,
 
 void __sanitizer_weak_hook_strcmp(void *caller_pc, const char *s1,
                                    const char *s2, int result) {
-  if (!TS) return;
+  if (!RecordingMemcmp) return;
   if (result == 0) return;  // No reason to mutate.
   size_t Len1 = strlen(s1);
   size_t Len2 = strlen(s2);
@@ -592,7 +599,7 @@ void __sanitizer_weak_hook_strcmp(void *caller_pc, const char *s1,
 __attribute__((visibility("default")))
 void __sanitizer_cov_trace_cmp(uint64_t SizeAndType, uint64_t Arg1,
                                uint64_t Arg2) {
-  if (!TS) return;
+  if (!RecordingTraces) return;
   uintptr_t PC = reinterpret_cast<uintptr_t>(__builtin_return_address(0));
   uint64_t CmpSize = (SizeAndType >> 32) / 8;
   uint64_t Type = (SizeAndType << 32) >> 32;
@@ -601,7 +608,7 @@ void __sanitizer_cov_trace_cmp(uint64_t SizeAndType, uint64_t Arg1,
 
 __attribute__((visibility("default")))
 void __sanitizer_cov_trace_switch(uint64_t Val, uint64_t *Cases) {
-  if (!TS) return;
+  if (!RecordingTraces) return;
   uintptr_t PC = reinterpret_cast<uintptr_t>(__builtin_return_address(0));
   TS->TraceSwitchCallback(PC, Cases[1], Val, Cases[0], Cases + 2);
 }