about summary refs log tree commit diff
diff options
context:
space:
mode:
authorhexcoder <hexcoder-@users.noreply.github.com>2021-05-31 20:06:35 +0200
committerGitHub <noreply@github.com>2021-05-31 20:06:35 +0200
commit21c8b225719835d8e4fecdb8bd6f42e61ea0600b (patch)
tree46e844356f7cf88c08f9f9907caa11656a24f416
parent1a2da67ed0505c9ac0aa1048ba3d607f3c1aa639 (diff)
parent97a1f89881878db9bd6b4cd666b3447a63818dcf (diff)
downloadafl++-21c8b225719835d8e4fecdb8bd6f42e61ea0600b.tar.gz
Merge pull request #948 from AFLplusplus/going_atomic
Going atomic
-rw-r--r--docs/Changelog.md1
-rw-r--r--docs/env_variables.md5
-rw-r--r--include/envs.h1
-rw-r--r--instrumentation/README.llvm.md4
-rw-r--r--instrumentation/README.neverzero.md14
-rw-r--r--instrumentation/SanitizerCoverageLTO.so.cc36
-rw-r--r--instrumentation/SanitizerCoveragePCGUARD.so.cc33
-rw-r--r--instrumentation/afl-llvm-lto-instrumentation.so.cc32
-rw-r--r--instrumentation/afl-llvm-pass.so.cc200
-rw-r--r--src/afl-cc.c1
-rwxr-xr-xtest/test-llvm.sh30
11 files changed, 292 insertions, 65 deletions
diff --git a/docs/Changelog.md b/docs/Changelog.md
index 298a3998..d8ffe498 100644
--- a/docs/Changelog.md
+++ b/docs/Changelog.md
@@ -58,6 +58,7 @@ sending a mail to <afl-users+subscribe@googlegroups.com>.
     MacOS shared memory
   - updated the grammar custom mutator to the newest version
   - add -d (add dead fuzzer stats) to afl-whatsup
+  - add thread safe counters for LLVM CLASSIC (set AFL_LLVM_THREADSAFE_INST)
   - added AFL_PRINT_FILENAMES to afl-showmap/cmin to print the
     current filename
   - afl-showmap/cmin will now process queue items in alphabetical order
diff --git a/docs/env_variables.md b/docs/env_variables.md
index 7bbc0fdd..b4b866ab 100644
--- a/docs/env_variables.md
+++ b/docs/env_variables.md
@@ -231,6 +231,11 @@ Then there are a few specific features that are only available in instrumentatio
 
   See [instrumentation/README.instrument_list.md](../instrumentation/README.instrument_list.md) for more information.
 
+### Thread safe instrumentation counters (in mode LLVM CLASSIC)
+   - Setting `AFL_LLVM_THREADSAFE_INST` will inject code that implements thread safe counters.
+     The overhead is a bit higher compared to the older non-thread safe case. 
+     `AFL_LLVM_NOT_ZERO` and `AFL_LLVM_SKIP_NEVERZERO` are supported (see below). 
+
 ### NOT_ZERO
 
    - Setting `AFL_LLVM_NOT_ZERO=1` during compilation will use counters
diff --git a/include/envs.h b/include/envs.h
index 08b3284a..15116fc1 100644
--- a/include/envs.h
+++ b/include/envs.h
@@ -126,6 +126,7 @@ static char *afl_environment_variables[] = {
     "AFL_NGRAM_SIZE",
     "AFL_LLVM_NOT_ZERO",
     "AFL_LLVM_INSTRUMENT_FILE",
+    "AFL_LLVM_THREADSAFE_INST",
     "AFL_LLVM_SKIP_NEVERZERO",
     "AFL_NO_AFFINITY",
     "AFL_TRY_AFFINITY",
diff --git a/instrumentation/README.llvm.md b/instrumentation/README.llvm.md
index cfe537d5..02722588 100644
--- a/instrumentation/README.llvm.md
+++ b/instrumentation/README.llvm.md
@@ -144,6 +144,10 @@ is not optimal and was only fixed in llvm 9.
 You can set this with AFL_LLVM_NOT_ZERO=1
 See [README.neverzero.md](README.neverzero.md)
 
+Support for thread safe counters has been added for mode LLVM CLASSIC.
+Activate it with `AFL_LLVM_THREADSAFE_INST=1`. The tradeoff is better precision in 
+multi threaded apps for a slightly higher instrumentation overhead.
+
 ## 4) Snapshot feature
 
 To speed up fuzzing you can use a linux loadable kernel module which enables
diff --git a/instrumentation/README.neverzero.md b/instrumentation/README.neverzero.md
index 49104e00..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
@@ -33,3 +34,8 @@ AFL_LLVM_SKIP_NEVERZERO=1
 ```
 If the target does not have extensive loops or functions that are called
 a lot then this can give a small performance boost.
+
+Please note that the default counter implementations are not thread safe!
+
+Support for thread safe counters in mode LLVM CLASSIC can be activated with setting
+`AFL_LLVM_THREADSAFE_INST=1`.
\ No newline at end of file
diff --git a/instrumentation/SanitizerCoverageLTO.so.cc b/instrumentation/SanitizerCoverageLTO.so.cc
index 2f4337eb..58969e18 100644
--- a/instrumentation/SanitizerCoverageLTO.so.cc
+++ b/instrumentation/SanitizerCoverageLTO.so.cc
@@ -236,7 +236,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;
@@ -437,6 +438,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)
@@ -1208,7 +1210,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.
@@ -1495,23 +1497,33 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
     }
 
     /* Update bitmap */
+    if (use_threadsafe_counters) { /* Atomic */
 
-    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) {
+
+        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));
+      }
+
+      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 8878d3b1..dbddad0a 100644
--- a/instrumentation/SanitizerCoveragePCGUARD.so.cc
+++ b/instrumentation/SanitizerCoveragePCGUARD.so.cc
@@ -86,7 +86,8 @@ const char SanCovPCsSectionName[] = "sancov_pcs";
 
 const char SanCovLowestStackName[] = "__sancov_lowest_stack";
 
-static char *skip_nozero;
+static const char *skip_nozero;
+static const char *use_threadsafe_counters;
 
 namespace {
 
@@ -386,6 +387,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);
@@ -1068,21 +1070,32 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
     /* Load counter for CurLoc */
 
     Value *   MapPtrIdx = IRB.CreateGEP(MapPtr, CurLoc);
-    LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
 
-    /* Update bitmap */
+    if (use_threadsafe_counters) {
 
-    Value *Incr = IRB.CreateAdd(Counter, One);
+      IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One,
+                          llvm::AtomicOrdering::Monotonic);
 
-    if (skip_nozero == NULL) {
+    }
+    else
+    {
 
-      auto cf = IRB.CreateICmpEQ(Incr, Zero);
-      auto carry = IRB.CreateZExt(cf, Int8Ty);
-      Incr = IRB.CreateAdd(Incr, carry);
+      LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
+      /* Update bitmap */
 
-    }
+      Value *Incr = IRB.CreateAdd(Counter, One);
+
+      if (skip_nozero == NULL) {
+
+        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);
+
+    }
 
     // done :)
 
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 :)
 
diff --git a/instrumentation/afl-llvm-pass.so.cc b/instrumentation/afl-llvm-pass.so.cc
index 0f773aba..fe9e2e40 100644
--- a/instrumentation/afl-llvm-pass.so.cc
+++ b/instrumentation/afl-llvm-pass.so.cc
@@ -85,7 +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;
+  const char *   ctx_str = NULL, *caller_str = NULL, *skip_nozero = NULL;
+  const char *   use_threadsafe_counters = nullptr;
 
 };
 
@@ -182,6 +183,26 @@ bool AFLCoverage::runOnModule(Module &M) {
   char *neverZero_counters_str = getenv("AFL_LLVM_NOT_ZERO");
 #endif
   skip_nozero = getenv("AFL_LLVM_SKIP_NEVERZERO");
+  use_threadsafe_counters = getenv("AFL_LLVM_THREADSAFE_INST");
+
+  if ((isatty(2) && !getenv("AFL_QUIET")) || !!getenv("AFL_DEBUG")) {
+
+    if (use_threadsafe_counters) {
+      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-thread safe instrumentation\n");
+    }
+
+  }
 
   unsigned PrevLocSize = 0;
   unsigned PrevCallerSize = 0;
@@ -388,7 +409,6 @@ bool AFLCoverage::runOnModule(Module &M) {
 #endif
 
   // other constants we need
-  ConstantInt *Zero = ConstantInt::get(Int8Ty, 0);
   ConstantInt *One = ConstantInt::get(Int8Ty, 1);
 
   Value *   PrevCtx = NULL;     // CTX sensitive coverage
@@ -410,6 +430,7 @@ bool AFLCoverage::runOnModule(Module &M) {
 
     if (F.size() < function_minimum_size) continue;
 
+    std::list<Value *> todo;
     for (auto &BB : F) {
 
       BasicBlock::iterator IP = BB.getFirstInsertionPt();
@@ -628,37 +649,63 @@ bool AFLCoverage::runOnModule(Module &M) {
 
       /* Update bitmap */
 
-      LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
-      Counter->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
 
-      Value *Incr = IRB.CreateAdd(Counter, One);
+      if (use_threadsafe_counters) {/* Atomic */
 
-#if LLVM_VERSION_MAJOR < 9
-      if (neverZero_counters_str !=
-          NULL) {  // with llvm 9 we make this the default as the bug in llvm is
-                   // then fixed
-#else
-      if (!skip_nozero) {
-
-#endif
-        /* hexcoder: Realize a counter that skips zero during overflow.
-         * Once this counter reaches its maximum value, it next increments to 1
-         *
-         * Instead of
-         * Counter + 1 -> Counter
-         * we inject now this
-         * Counter + 1 -> {Counter, OverflowFlag}
-         * Counter + OverflowFlag -> Counter
-         */
+  #if LLVM_VERSION_MAJOR < 9
+        if (neverZero_counters_str !=
+            NULL) {  // with llvm 9 we make this the default as the bug in llvm is then fixed
+  #else
+        if (!skip_nozero) {
 
-        auto cf = IRB.CreateICmpEQ(Incr, Zero);
-        auto carry = IRB.CreateZExt(cf, Int8Ty);
-        Incr = IRB.CreateAdd(Incr, carry);
+  #endif
+          // register MapPtrIdx in a todo list
+          todo.push_back(MapPtrIdx);
 
+        }
+        else
+        {
+          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);
+
+  #if LLVM_VERSION_MAJOR < 9
+        if (neverZero_counters_str !=
+            NULL) {  // with llvm 9 we make this the default as the bug in llvm is
+                     // then fixed
+  #else
+        if (!skip_nozero) {
+
+  #endif
+          /* hexcoder: Realize a counter that skips zero during overflow.
+           * Once this counter reaches its maximum value, it next increments to 1
+           *
+           * Instead of
+           * Counter + 1 -> Counter
+           * we inject now this
+           * Counter + 1 -> {Counter, OverflowFlag}
+           * Counter + OverflowFlag -> Counter
+           */
+
+          ConstantInt *Zero = ConstantInt::get(Int8Ty, 0);
+          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));
+      } /* non atomic case */
 
       /* Update prev_loc history vector (by placing cur_loc at the head of the
          vector and shuffle the other elements back by one) */
@@ -715,6 +762,105 @@ bool AFLCoverage::runOnModule(Module &M) {
 
     }
 
+    if (use_threadsafe_counters) { /*Atomic NeverZero */
+      // handle the list of registered blocks to instrument
+      for (auto val : todo) {
+        /* hexcoder: Realize a thread-safe counter that skips zero during overflow. Once this counter reaches its maximum value, it next increments to 1
+             *
+             * Instead of
+             * Counter + 1 -> Counter
+             * we inject now this
+             * Counter + 1 -> {Counter, OverflowFlag}
+             * Counter + OverflowFlag -> Counter
+         */
+
+        /* equivalent c code looks like this
+         * Thanks to
+         https://preshing.com/20150402/you-can-do-any-kind-of-atomic-read-modify-write-operation/
+
+            int old = atomic_load_explicit(&Counter, memory_order_relaxed);
+            int new;
+            do {
+                 if (old == 255) {
+                   new = 1;
+                 } else {
+                   new = old + 1;
+                 }
+            } while (!atomic_compare_exchange_weak_explicit(&Counter, &old, new,
+         memory_order_relaxed, memory_order_relaxed));
+
+         */
+
+        Value *              MapPtrIdx = val;
+        Instruction *        MapPtrIdxInst = cast<Instruction>(val);
+        BasicBlock::iterator it0(&(*MapPtrIdxInst));
+        ++it0;
+        IRBuilder<> IRB(&(*it0));
+
+        // load the old counter value atomically
+        LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
+        Counter->setAlignment(llvm::Align());
+        Counter->setAtomic(llvm::AtomicOrdering::Monotonic);
+        Counter->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
+
+        BasicBlock *BB = IRB.GetInsertBlock();
+        // insert a basic block with the corpus of a do while loop
+        // the calculation may need to repeat, if atomic compare_exchange is not successful
+
+        BasicBlock::iterator it(*Counter);
+        it++;  // split after load counter
+        BasicBlock *end_bb = BB->splitBasicBlock(it);
+        end_bb->setName("injected");
+
+        // insert the block before the second half of the split
+        BasicBlock *do_while_bb =
+            BasicBlock::Create(C, "injected", end_bb->getParent(), end_bb);
+
+        // set terminator of BB from target end_bb to target do_while_bb
+        auto term = BB->getTerminator();
+        BranchInst::Create(do_while_bb, BB);
+        term->eraseFromParent();
+
+        // continue to fill instructions into the do_while loop
+        IRB.SetInsertPoint(do_while_bb, do_while_bb->getFirstInsertionPt());
+
+        PHINode *PN = IRB.CreatePHI(Int8Ty, 2);
+
+        // compare with maximum value 0xff
+        auto *Cmp = IRB.CreateICmpEQ(Counter, ConstantInt::get(Int8Ty, -1));
+
+        // increment the counter
+        Value *Incr = IRB.CreateAdd(Counter, One);
+
+        // select the counter value or 1
+        auto *Select = IRB.CreateSelect(Cmp, One, Incr);
+
+        // try to save back the new counter value
+        auto *CmpXchg = IRB.CreateAtomicCmpXchg(
+            MapPtrIdx, PN, Select, llvm::AtomicOrdering::Monotonic,
+            llvm::AtomicOrdering::Monotonic);
+        CmpXchg->setAlignment(llvm::Align());
+        CmpXchg->setWeak(true);
+        CmpXchg->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
+
+        // get the result of trying to update the Counter
+        Value *Success =
+            IRB.CreateExtractValue(CmpXchg, ArrayRef<unsigned>({1}));
+        // get the (possibly updated) value of Counter
+        Value *OldVal =
+            IRB.CreateExtractValue(CmpXchg, ArrayRef<unsigned>({0}));
+
+        // initially we use Counter
+        PN->addIncoming(Counter, BB);
+        // on retry, we use the updated value
+        PN->addIncoming(OldVal, do_while_bb);
+
+        // if the cmpXchg was not successful, retry
+        IRB.CreateCondBr(Success, end_bb, do_while_bb);
+      }
+
+    }
+
   }
 
   /*
diff --git a/src/afl-cc.c b/src/afl-cc.c
index 8af8e7b0..6be6e165 100644
--- a/src/afl-cc.c
+++ b/src/afl-cc.c
@@ -1777,6 +1777,7 @@ int main(int argc, char **argv, char **envp) {
         SAYF(
             "\nLLVM/LTO/afl-clang-fast/afl-clang-lto specific environment "
             "variables:\n"
+            "  AFL_LLVM_THREADSAFE_INST: instrument with thread safe counters\n"
 
             COUNTER_BEHAVIOUR
 
diff --git a/test/test-llvm.sh b/test/test-llvm.sh
index 06d0a0f8..1152cc4e 100755
--- a/test/test-llvm.sh
+++ b/test/test-llvm.sh
@@ -43,6 +43,36 @@ test -e ../afl-clang-fast -a -e ../split-switches-pass.so && {
     $ECHO "$RED[!] llvm_mode failed"
     CODE=1
   }
+  AFL_LLVM_INSTRUMENT=CLASSIC AFL_LLVM_THREADSAFE_INST=1 ../afl-clang-fast -o test-instr.ts ../test-instr.c > /dev/null 2>&1
+  test -e test-instr.ts && {
+    $ECHO "$GREEN[+] llvm_mode threadsafe compilation succeeded"
+    echo 0 | AFL_QUIET=1 ../afl-showmap -m ${MEM_LIMIT} -o test-instr.ts.0 -r -- ./test-instr.ts > /dev/null 2>&1
+    AFL_QUIET=1 ../afl-showmap -m ${MEM_LIMIT} -o test-instr.ts.1 -r -- ./test-instr.ts < /dev/null > /dev/null 2>&1
+    test -e test-instr.ts.0 -a -e test-instr.ts.1 && {
+      diff test-instr.ts.0 test-instr.ts.1 > /dev/null 2>&1 && {
+        $ECHO "$RED[!] llvm_mode threadsafe instrumentation should be different on different input but is not"
+        CODE=1
+      } || {
+        $ECHO "$GREEN[+] llvm_mode threadsafe instrumentation present and working correctly"
+        TUPLES=`echo 0|AFL_QUIET=1 ../afl-showmap -m ${MEM_LIMIT} -o /dev/null -- ./test-instr.ts 2>&1 | grep Captur | awk '{print$3}'`
+        test "$TUPLES" -gt 2 -a "$TUPLES" -lt 8 && {
+          $ECHO "$GREEN[+] llvm_mode run reported $TUPLES threadsafe instrumented locations which is fine"
+        } || {
+          $ECHO "$RED[!] llvm_mode threadsafe instrumentation produces weird numbers: $TUPLES"
+          CODE=1
+        }
+        test "$TUPLES" -lt 3 && SKIP=1
+        true
+      }
+    } || {
+      $ECHO "$RED[!] llvm_mode threadsafe instrumentation failed"
+      CODE=1
+    }
+    rm -f test-instr.ts.0 test-instr.ts.1
+  } || {
+    $ECHO "$RED[!] llvm_mode (threadsafe) failed"
+    CODE=1
+  }
   ../afl-clang-fast -DTEST_SHARED_OBJECT=1 -z defs -fPIC -shared -o test-instr.so ../test-instr.c > /dev/null 2>&1
   test -e test-instr.so && {
     $ECHO "$GREEN[+] llvm_mode shared object with -z defs compilation succeeded"