add support for AFL_LLVM_THREADSAFE_INST to other LLVM passes

5 files changed, 81 insertions, 56 deletions

diff --git a/instrumentation/README.neverzero.md b/instrumentation/README.neverzero.md
index 06334eab..9bcae324 100644
--- a/instrumentation/README.neverzero.md
+++ b/instrumentation/README.neverzero.md
@@ -16,11 +16,12 @@ at a very little cost (one instruction per edge).
 (The alternative of saturated counters has been tested also and proved to be
 inferior in terms of path discovery.)
 
-This is implemented in afl-gcc and afl-gcc-fast, however for llvm_mode this is optional if
-the llvm version is below 9 - as there is a perfomance bug that is only fixed
-in version 9 and onwards.
+This is implemented in afl-gcc and afl-gcc-fast, however for llvm_mode this is
+optional if multithread safe counters are selected or the llvm version is below
+9 - as there are severe performance costs in these cases.
 
-If you want to enable this for llvm versions below 9 then set
+If you want to enable this for llvm versions below 9 or thread safe counters
+then set
 
 ```
 export AFL_LLVM_NOT_ZERO=1
diff --git a/instrumentation/SanitizerCoverageLTO.so.cc b/instrumentation/SanitizerCoverageLTO.so.cc
index cd6b1939..f5af32d2 100644
--- a/instrumentation/SanitizerCoverageLTO.so.cc
+++ b/instrumentation/SanitizerCoverageLTO.so.cc
@@ -237,7 +237,8 @@ class ModuleSanitizerCoverage {
   uint32_t                         inst = 0;
   uint32_t                         afl_global_id = 0;
   uint64_t                         map_addr = 0;
-  char *                           skip_nozero = NULL;
+  const char *                     skip_nozero = NULL;
+  const char *                     use_threadsafe_counters = nullptr;
   std::vector<BasicBlock *>        BlockList;
   DenseMap<Value *, std::string *> valueMap;
   std::vector<std::string>         dictionary;
@@ -438,6 +439,7 @@ bool ModuleSanitizerCoverage::instrumentModule(
     be_quiet = 1;
 
   skip_nozero = getenv("AFL_LLVM_SKIP_NEVERZERO");
+  use_threadsafe_counters = getenv("AFL_LLVM_THREADSAFE_INST");
 
   if ((ptr = getenv("AFL_LLVM_LTO_STARTID")) != NULL)
     if ((afl_global_id = atoi(ptr)) < 0)
@@ -1209,7 +1211,7 @@ void ModuleSanitizerCoverage::instrumentFunction(
     return;  // Should not instrument sanitizer init functions.
   if (F.getName().startswith("__sanitizer_"))
     return;  // Don't instrument __sanitizer_* callbacks.
-  // Don't touch available_externally functions, their actual body is elewhere.
+  // Don't touch available_externally functions, their actual body is elsewhere.
   if (F.getLinkage() == GlobalValue::AvailableExternallyLinkage) return;
   // Don't instrument MSVC CRT configuration helpers. They may run before normal
   // initialization.
@@ -1496,27 +1498,33 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
     }
 
     /* Update bitmap */
-#if 1 /* Atomic */
-    IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One, 
-        llvm::AtomicOrdering::Monotonic);
+    if (use_threadsafe_counters) { /* Atomic */
 
-#else
-    LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
-    Counter->setMetadata(Mo->getMDKindID("nosanitize"), MDNode::get(*Ct, None));
+      IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One,
+                          llvm::AtomicOrdering::Monotonic);
 
-    Value *Incr = IRB.CreateAdd(Counter, One);
+    }
+    else
+    {
 
-    if (skip_nozero == NULL) {
+      LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
+      Counter->setMetadata(Mo->getMDKindID("nosanitize"),
+                           MDNode::get(*Ct, None));
 
-      auto cf = IRB.CreateICmpEQ(Incr, Zero);
-      auto carry = IRB.CreateZExt(cf, Int8Tyi);
-      Incr = IRB.CreateAdd(Incr, carry);
+      Value *Incr = IRB.CreateAdd(Counter, One);
 
-    }
+      if (skip_nozero == NULL) {
 
-    IRB.CreateStore(Incr, MapPtrIdx)
-        ->setMetadata(Mo->getMDKindID("nosanitize"), MDNode::get(*Ct, None));
-#endif
+        auto cf = IRB.CreateICmpEQ(Incr, Zero);
+        auto carry = IRB.CreateZExt(cf, Int8Tyi);
+        Incr = IRB.CreateAdd(Incr, carry);
+
+      }
+
+      IRB.CreateStore(Incr, MapPtrIdx)
+          ->setMetadata(Mo->getMDKindID("nosanitize"), MDNode::get(*Ct, None));
+
+    }
     // done :)
 
     inst++;
diff --git a/instrumentation/SanitizerCoveragePCGUARD.so.cc b/instrumentation/SanitizerCoveragePCGUARD.so.cc
index dd2e1459..e1e922be 100644
--- a/instrumentation/SanitizerCoveragePCGUARD.so.cc
+++ b/instrumentation/SanitizerCoveragePCGUARD.so.cc
@@ -96,7 +96,8 @@ static const char *const SanCovPCsSectionName = "sancov_pcs";
 
 static const char *const SanCovLowestStackName = "__sancov_lowest_stack";
 
-static char *skip_nozero;
+static const char *skip_nozero;
+static const char *use_threadsafe_counters;
 
 namespace {
 
@@ -396,6 +397,7 @@ bool ModuleSanitizerCoverage::instrumentModule(
     be_quiet = 1;
 
   skip_nozero = getenv("AFL_LLVM_SKIP_NEVERZERO");
+  use_threadsafe_counters = getenv("AFL_LLVM_THREADSAFE_INST");
 
   initInstrumentList();
   scanForDangerousFunctions(&M);
@@ -1081,27 +1083,32 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
 
     Value *   MapPtrIdx = IRB.CreateGEP(MapPtr, CurLoc);
 
-#if 1 /* Atomic */
-    IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One,
-        llvm::AtomicOrdering::Monotonic);
+    if (use_threadsafe_counters) {
 
-#else
-    LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
+      IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One,
+                          llvm::AtomicOrdering::Monotonic);
+
+    }
+    else
+    {
 
-    /* Update bitmap */
+      LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
+      /* Update bitmap */
 
-    Value *Incr = IRB.CreateAdd(Counter, One);
+      Value *Incr = IRB.CreateAdd(Counter, One);
 
-    if (skip_nozero == NULL) {
+      if (skip_nozero == NULL) {
 
-      auto cf = IRB.CreateICmpEQ(Incr, Zero);
-      auto carry = IRB.CreateZExt(cf, Int8Ty);
-      Incr = IRB.CreateAdd(Incr, carry);
+        auto cf = IRB.CreateICmpEQ(Incr, Zero);
+        auto carry = IRB.CreateZExt(cf, Int8Ty);
+        Incr = IRB.CreateAdd(Incr, carry);
+
+      }
+
+      IRB.CreateStore(Incr, MapPtrIdx);
 
     }
 
-    IRB.CreateStore(Incr, MapPtrIdx);
-#endif
     // done :)
 
     //    IRB.CreateCall(SanCovTracePCGuard, Offset)->setCannotMerge();
diff --git a/instrumentation/afl-llvm-lto-instrumentation.so.cc b/instrumentation/afl-llvm-lto-instrumentation.so.cc
index 5ed13ff0..10cfa579 100644
--- a/instrumentation/afl-llvm-lto-instrumentation.so.cc
+++ b/instrumentation/afl-llvm-lto-instrumentation.so.cc
@@ -93,7 +93,8 @@ class AFLLTOPass : public ModulePass {
   uint32_t function_minimum_size = 1;
   uint32_t inst_blocks = 0, inst_funcs = 0, total_instr = 0;
   unsigned long long int map_addr = 0x10000;
-  char *   skip_nozero = NULL;
+  const char *skip_nozero = NULL;
+  const char *use_threadsafe_counters = nullptr;
 
 };
 
@@ -131,6 +132,8 @@ bool AFLLTOPass::runOnModule(Module &M) {
 
     be_quiet = 1;
 
+  use_threadsafe_counters = getenv("AFL_LLVM_THREADSAFE_INST");
+
   if ((ptr = getenv("AFL_LLVM_DOCUMENT_IDS")) != NULL) {
 
     if ((documentFile = fopen(ptr, "a")) == NULL)
@@ -839,29 +842,28 @@ bool AFLLTOPass::runOnModule(Module &M) {
 
           /* Update bitmap */
 
-#if 1 /* Atomic */
-          IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One,
-              llvm::AtomicOrdering::Monotonic);
+          if (use_threadsafe_counters) {
+            IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One,
+                                llvm::AtomicOrdering::Monotonic);
+          } else {
+            LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
+            Counter->setMetadata(M.getMDKindID("nosanitize"),
+                                 MDNode::get(C, None));
 
-#else
-          LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
-          Counter->setMetadata(M.getMDKindID("nosanitize"),
-                               MDNode::get(C, None));
+            Value *Incr = IRB.CreateAdd(Counter, One);
 
-          Value *Incr = IRB.CreateAdd(Counter, One);
+            if (skip_nozero == NULL) {
 
-          if (skip_nozero == NULL) {
+              auto cf = IRB.CreateICmpEQ(Incr, Zero);
+              auto carry = IRB.CreateZExt(cf, Int8Ty);
+              Incr = IRB.CreateAdd(Incr, carry);
 
-            auto cf = IRB.CreateICmpEQ(Incr, Zero);
-            auto carry = IRB.CreateZExt(cf, Int8Ty);
-            Incr = IRB.CreateAdd(Incr, carry);
+            }
 
+            IRB.CreateStore(Incr, MapPtrIdx)
+                ->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
           }
 
-          IRB.CreateStore(Incr, MapPtrIdx)
-              ->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
-#endif
-
           // done :)
 
           inst_blocks++;
diff --git a/instrumentation/afl-llvm-pass.so.cc b/instrumentation/afl-llvm-pass.so.cc
index 3b1119fc..fe9e2e40 100644
--- a/instrumentation/afl-llvm-pass.so.cc
+++ b/instrumentation/afl-llvm-pass.so.cc
@@ -85,8 +85,8 @@ class AFLCoverage : public ModulePass {
   uint32_t ctx_k = 0;
   uint32_t map_size = MAP_SIZE;
   uint32_t function_minimum_size = 1;
-  char *   ctx_str = NULL, *caller_str = NULL, *skip_nozero = NULL;
-  char *   use_threadsafe_counters = nullptr;
+  const char *   ctx_str = NULL, *caller_str = NULL, *skip_nozero = NULL;
+  const char *   use_threadsafe_counters = nullptr;
 
 };
 
@@ -188,11 +188,18 @@ bool AFLCoverage::runOnModule(Module &M) {
   if ((isatty(2) && !getenv("AFL_QUIET")) || !!getenv("AFL_DEBUG")) {
 
     if (use_threadsafe_counters) {
-      SAYF(cCYA "afl-llvm-pass" VERSION cRST " using threadsafe instrumentation\n");
+      if (!getenv("AFL_LLVM_NOT_ZERO")) {
+        skip_nozero = "1";
+        SAYF(cCYA "afl-llvm-pass" VERSION cRST " using thread safe counters\n");
+      }
+      else {
+        SAYF(cCYA "afl-llvm-pass" VERSION cRST
+                  " using thread safe not-zero-counters\n");
+      }
     }
     else
     {
-      SAYF(cCYA "afl-llvm-pass" VERSION cRST " using non-threadsafe instrumentation\n");
+      SAYF(cCYA "afl-llvm-pass" VERSION cRST " using non-thread safe instrumentation\n");
     }
 
   }