support new env var AFL_LLVM_THREADSAFE_INST to enable atomic counters.

add new test case for that.

3 files changed, 183 insertions, 127 deletions

diff --git a/include/envs.h b/include/envs.h
index ebe98257..e6f6d7c9 100644
--- a/include/envs.h
+++ b/include/envs.h
@@ -114,6 +114,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_LLVM_LTO_STARTID",
diff --git a/instrumentation/afl-llvm-pass.so.cc b/instrumentation/afl-llvm-pass.so.cc
index 53e076ff..3b1119fc 100644
--- a/instrumentation/afl-llvm-pass.so.cc
+++ b/instrumentation/afl-llvm-pass.so.cc
@@ -86,6 +86,7 @@ class AFLCoverage : public ModulePass {
   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;
 
 };
 
@@ -182,6 +183,19 @@ 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) {
+      SAYF(cCYA "afl-llvm-pass" VERSION cRST " using threadsafe instrumentation\n");
+    }
+    else
+    {
+      SAYF(cCYA "afl-llvm-pass" VERSION cRST " using non-threadsafe instrumentation\n");
+    }
+
+  }
 
   unsigned PrevLocSize = 0;
   unsigned PrevCallerSize = 0;
@@ -628,57 +642,63 @@ bool AFLCoverage::runOnModule(Module &M) {
 
       /* Update bitmap */
 
-#if 1 /* 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
+      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
           // register MapPtrIdx in a todo list
           todo.push_back(MapPtrIdx);
 
-      } else {
-          IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One, llvm::AtomicOrdering::Monotonic);
+        }
+        else
+        {
+          IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One,
+                              llvm::AtomicOrdering::Monotonic);
+        }
       }
+      else
+      {
 
-#else
-      LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
-      Counter->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
+        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 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) {
+  #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
-         */
+  #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);
+          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));
 
-#endif /* non atomic case */
+      } /* 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) */
@@ -735,99 +755,104 @@ bool AFLCoverage::runOnModule(Module &M) {
 
     }
 
-#if 1 /*Atomic NeverZero */
-    // handle the todo list
-    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));
+    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
+         */
 
-          // 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}));
+        /* 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));
 
-          // 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);
+        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);
+      }
 
-     }
-#endif
+    }
 
   }
 
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"