diff options
| -rw-r--r-- | instrumentation/SanitizerCoverageLTO.so.cc | 140 |
1 files changed, 137 insertions, 3 deletions
diff --git a/instrumentation/SanitizerCoverageLTO.so.cc b/instrumentation/SanitizerCoverageLTO.so.cc index 7a3d8c4d..bc4df34e 100644 --- a/instrumentation/SanitizerCoverageLTO.so.cc +++ b/instrumentation/SanitizerCoverageLTO.so.cc @@ -235,6 +235,8 @@ class ModuleSanitizerCoverage { uint32_t autodictionary = 1; uint32_t inst = 0; uint32_t afl_global_id = 0; + uint32_t unhandled = 0; + uint32_t select_cnt = 0; uint64_t map_addr = 0; const char * skip_nozero = NULL; const char * use_threadsafe_counters = nullptr; @@ -1150,9 +1152,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 (on average %llu " - "collisions would be in afl-gcc/vanilla AFL) (%s mode).", - inst, calculateCollisions(inst), modeline); + OKF("Instrumented %u locations (%u selects) without collisions (%llu " + "collisions have been avoided) (%s mode).", + inst, select_cnt, calculateCollisions(inst), modeline); } @@ -1274,6 +1276,7 @@ void ModuleSanitizerCoverage::instrumentFunction( const DominatorTree * DT = DTCallback(F); const PostDominatorTree *PDT = PDTCallback(F); bool IsLeafFunc = true; + uint32_t skip_next = 0, local_selects = 0; for (auto &BB : F) { @@ -1291,6 +1294,137 @@ void ModuleSanitizerCoverage::instrumentFunction( Value *val = ConstantInt::get(Int32Ty, ++afl_global_id); callInst->setOperand(1, val); + ++inst; + + } + + SelectInst *selectInst = nullptr; + + if (!skip_next && (selectInst = dyn_cast<SelectInst>(&IN))) { + + uint32_t vector_cnt = 0; + Value * condition = selectInst->getCondition(); + Value * result; + auto t = condition->getType(); + IRBuilder<> IRB(selectInst->getNextNode()); + + ++select_cnt; + + if (t->getTypeID() == llvm::Type::IntegerTyID) { + + Value *val1 = ConstantInt::get(Int32Ty, ++afl_global_id); + Value *val2 = ConstantInt::get(Int32Ty, ++afl_global_id); + result = IRB.CreateSelect(condition, val1, val2); + inst += 2; + + } else if (t->getTypeID() == llvm::Type::FixedVectorTyID) { + + FixedVectorType *tt = dyn_cast<FixedVectorType>(t); + if (tt) { + + uint32_t elements = tt->getElementCount().getFixedValue(); + vector_cnt = elements; + inst += vector_cnt * 2; + if (elements) { + + FixedVectorType *GuardPtr1 = + FixedVectorType::get(Int32Ty, elements); + FixedVectorType *GuardPtr2 = + FixedVectorType::get(Int32Ty, elements); + Value *x, *y; + + Value *val1 = ConstantInt::get(Int32Ty, ++afl_global_id); + Value *val2 = ConstantInt::get(Int32Ty, ++afl_global_id); + x = IRB.CreateInsertElement(GuardPtr1, val1, (uint64_t)0); + y = IRB.CreateInsertElement(GuardPtr2, val2, (uint64_t)0); + + for (uint64_t i = 1; i < elements; i++) { + + val1 = ConstantInt::get(Int32Ty, ++afl_global_id); + val2 = ConstantInt::get(Int32Ty, ++afl_global_id); + x = IRB.CreateInsertElement(GuardPtr1, val1, i); + y = IRB.CreateInsertElement(GuardPtr2, val2, i); + + } + + /* + std::string errMsg; + raw_string_ostream os(errMsg); + x->print(os); + fprintf(stderr, "X: %s\n", os.str().c_str()); + */ + result = IRB.CreateSelect(condition, x, y); + + } + + } + + } else { + + unhandled++; + continue; + + } + + local_selects++; + uint32_t vector_cur = 0; + /* Load SHM pointer */ + LoadInst *MapPtr = IRB.CreateLoad(AFLMapPtr); + + while (1) { + + /* Get CurLoc */ + Value *MapPtrIdx = nullptr; + + /* Load counter for CurLoc */ + if (!vector_cnt) { + + MapPtrIdx = IRB.CreateGEP(MapPtr, result); + + } else { + + auto element = IRB.CreateExtractElement(result, vector_cur++); + MapPtrIdx = IRB.CreateGEP(MapPtr, element); + + } + + 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); + + } + + if (!vector_cnt || vector_cnt == vector_cur) { break; } + + } + + skip_next = 1; + + } else { + + skip_next = 0; } |