diff options
| -rw-r--r-- | include/envs.h | 1 | ||||
| -rw-r--r-- | instrumentation/afl-llvm-pass.so.cc | 279 | ||||
| -rwxr-xr-x | test/test-llvm.sh | 30 |
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" |