Merge pull request #948 from AFLplusplus/going_atomic

Going atomic

1 files changed, 20 insertions, 12 deletions

diff --git a/instrumentation/afl-llvm-lto-instrumentation.so.cc b/instrumentation/afl-llvm-lto-instrumentation.so.cc
index 68bd2fa5..b5fdb3d6 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,22 +842,27 @@ 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);
+          } 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));
+          }
 
           // done :)