work in progress: not working correctly yet

1 files changed, 78 insertions, 1 deletions

diff --git a/instrumentation/afl-llvm-pass.so.cc b/instrumentation/afl-llvm-pass.so.cc
index 70480ff9..6c898c48 100644
--- a/instrumentation/afl-llvm-pass.so.cc
+++ b/instrumentation/afl-llvm-pass.so.cc
@@ -409,8 +409,14 @@ bool AFLCoverage::runOnModule(Module &M) {
 
     if (F.size() < function_minimum_size) continue;
 
+    unsigned extra_increment_BB = 0;
     for (auto &BB : F) {
 
+      if (extra_increment_BB) {
+          // increment BB
+          --extra_increment_BB;
+          continue;
+      }
       BasicBlock::iterator IP = BB.getFirstInsertionPt();
       IRBuilder<>          IRB(&(*IP));
 
@@ -628,7 +634,78 @@ bool AFLCoverage::runOnModule(Module &M) {
       /* Update bitmap */
 
 #if 1 /* Atomic */
-      IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One, llvm::AtomicOrdering::Monotonic);
+#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
+           */
+
+          // C: unsigned char old = atomic_load_explicit(MapPtrIdx, memory_order_relaxed);
+          LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
+          Counter->setAlignment(llvm::Align());
+          Counter->setAtomic(llvm::AtomicOrdering::Monotonic);
+          Counter->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
+
+          // 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++;
+          BasicBlock * end_bb = BB.splitBasicBlock(it);
+
+          extra_increment_BB = 2;
+          // 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();
+
+          auto saved = IRB.saveIP();
+          IRB.SetInsertPoint(do_while_bb, do_while_bb->getFirstInsertionPt());
+
+          PHINode * PN = IRB.CreatePHI(Int8Ty, 2);
+
+          auto * Cmp = IRB.CreateICmpEQ(Counter, ConstantInt::get(Int8Ty, -1));
+
+          Value *Incr = IRB.CreateAdd(Counter, One);
+
+          auto * Select = IRB.CreateSelect(Cmp, One, Incr);
+
+          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));
+
+          Value * Success = IRB.CreateExtractValue(CmpXchg, ArrayRef<unsigned>({1}));
+          Value * OldVal = IRB.CreateExtractValue(CmpXchg, ArrayRef<unsigned>({0}));
+
+          PN->addIncoming(Counter, &BB);
+          PN->addIncoming(OldVal, do_while_bb);
+
+//          term = do_while_bb->getTerminator();
+
+//          BranchInst::Create(/*true*/end_bb, /*false*/do_while_bb, Success, do_while_bb);
+          IRB.CreateCondBr(Success, end_bb, do_while_bb);
+//          BranchInst::Create(end_bb, do_while_bb);
+//          term->eraseFromParent();
+          IRB.restoreIP(saved);
+
+      } else {
+          IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One, llvm::AtomicOrdering::Monotonic);
+      }
 
 #else
       LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);