diff options
| -rw-r--r-- | instrumentation/SanitizerCoveragePCGUARD.so.cc | 67 |
1 files changed, 64 insertions, 3 deletions
diff --git a/instrumentation/SanitizerCoveragePCGUARD.so.cc b/instrumentation/SanitizerCoveragePCGUARD.so.cc index 013492f9..d6098897 100644 --- a/instrumentation/SanitizerCoveragePCGUARD.so.cc +++ b/instrumentation/SanitizerCoveragePCGUARD.so.cc @@ -203,7 +203,7 @@ class ModuleSanitizerCoverage { SanitizerCoverageOptions Options; - uint32_t instr = 0; + uint32_t instr = 0, selects = 0; GlobalVariable *AFLMapPtr = NULL; ConstantInt * One = NULL; ConstantInt * Zero = NULL; @@ -553,8 +553,9 @@ bool ModuleSanitizerCoverage::instrumentModule( getenv("AFL_USE_MSAN") ? ", MSAN" : "", getenv("AFL_USE_CFISAN") ? ", CFISAN" : "", getenv("AFL_USE_UBSAN") ? ", UBSAN" : ""); - OKF("Instrumented %u locations with no collisions (%s mode).", instr, - modeline); + OKF("Instrumented %u locations with no collisions (%s mode) and %u " + "selects.", + instr, modeline, selects); } @@ -836,6 +837,8 @@ bool ModuleSanitizerCoverage::InjectCoverage(Function & F, if (AllBlocks.empty()) return false; uint32_t special = 0; + uint32_t skip_next = 0; + for (auto &BB : F) { for (auto &IN : BB) { @@ -856,6 +859,64 @@ bool ModuleSanitizerCoverage::InjectCoverage(Function & F, } + SelectInst *selectInst = nullptr; + + if (!skip_next && (selectInst = dyn_cast<SelectInst>(&IN))) { + + selects++; + uint32_t id1 = 1 + instr + (uint32_t)AllBlocks.size() + special++; + uint32_t id2 = 1 + instr + (uint32_t)AllBlocks.size() + special++; + Value * val1 = ConstantInt::get(Int32Ty, id1); + Value * val2 = ConstantInt::get(Int32Ty, id2); + auto cond = selectInst->getCondition(); + IRBuilder<> IRB(selectInst->getNextNode()); + auto result = IRB.CreateSelect(cond, val1, val2); + + /* Get CurLoc */ + + /* Load SHM pointer */ + + LoadInst *MapPtr = IRB.CreateLoad(AFLMapPtr); + + /* Load counter for CurLoc */ + + Value *MapPtrIdx = IRB.CreateGEP(MapPtr, result); + + if (use_threadsafe_counters) { + + IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One, +#if LLVM_VERSION_MAJOR >= 13 + llvm::MaybeAlign(1), +#endif + llvm::AtomicOrdering::Monotonic); + + } else { + + 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); + + } + + skip_next = 1; + + } else { + + skip_next = 0; + + } + } } |