select support for LTO
MODIFY
instrumentation/SanitizerCoverageLTO.so.cc
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;
}

select support for LTO