1 files changed, 30 insertions, 17 deletions
diff --git a/instrumentation/afl-llvm-lto-instrumentation.so.cc b/instrumentation/afl-llvm-lto-instrumentation.so.cc
index 50306224..fe43fbe5 100644
--- a/instrumentation/afl-llvm-lto-instrumentation.so.cc
+++ b/instrumentation/afl-llvm-lto-instrumentation.so.cc
@@ -89,11 +89,12 @@ class AFLLTOPass : public ModulePass {
   bool runOnModule(Module &M) override;
 
  protected:
-  uint32_t afl_global_id = 1, autodictionary = 1;
-  uint32_t function_minimum_size = 1;
-  uint32_t inst_blocks = 0, inst_funcs = 0, total_instr = 0;
-  uint64_t map_addr = 0x10000;
-  char *   skip_nozero = NULL;
+  uint32_t               afl_global_id = 1, autodictionary = 1;
+  uint32_t               function_minimum_size = 1;
+  uint32_t               inst_blocks = 0, inst_funcs = 0, total_instr = 0;
+  unsigned long long int map_addr = 0x10000;
+  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)
@@ -176,7 +179,7 @@ bool AFLLTOPass::runOnModule(Module &M) {
 
   }
 
-  if (debug) { fprintf(stderr, "map address is 0x%lx\n", map_addr); }
+  if (debug) { fprintf(stderr, "map address is 0x%llx\n", map_addr); }
 
   /* Get/set the globals for the SHM region. */
 
@@ -839,22 +842,32 @@ bool AFLLTOPass::runOnModule(Module &M) {
 
           /* Update bitmap */
 
-          LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
-          Counter->setMetadata(M.getMDKindID("nosanitize"),
-                               MDNode::get(C, None));
+          if (use_threadsafe_counters) {
+
+            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(M.getMDKindID("nosanitize"),
+                                 MDNode::get(C, None));
 
-            auto cf = IRB.CreateICmpEQ(Incr, Zero);
-            auto carry = IRB.CreateZExt(cf, Int8Ty);
-            Incr = IRB.CreateAdd(Incr, carry);
+            Value *Incr = IRB.CreateAdd(Counter, One);
 
-          }
+            if (skip_nozero == NULL) {
 
-          IRB.CreateStore(Incr, MapPtrIdx)
-              ->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
+              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));
+
+          }
 
           // done :)