From 4d493452a45655073d1b7b1dfe4ad04772b3c2b8 Mon Sep 17 00:00:00 2001 From: vanhauser-thc Date: Thu, 25 Jan 2024 17:00:53 +0100 Subject: [PATCH 1/9] tmp --- instrumentation/SanitizerCoverageLTO.so.cc | 55 +++++++++++++++++++++- src/afl-cc.c | 10 +++- 2 files changed, 61 insertions(+), 4 deletions(-) diff --git a/instrumentation/SanitizerCoverageLTO.so.cc b/instrumentation/SanitizerCoverageLTO.so.cc index 68423029..c74069e1 100644 --- a/instrumentation/SanitizerCoverageLTO.so.cc +++ b/instrumentation/SanitizerCoverageLTO.so.cc @@ -247,6 +247,7 @@ class ModuleSanitizerCoverageLTO uint32_t afl_global_id = 0; uint32_t unhandled = 0; uint32_t select_cnt = 0; + uint32_t instrument_ctx = 0; uint64_t map_addr = 0; const char *skip_nozero = NULL; const char *use_threadsafe_counters = nullptr; @@ -261,6 +262,7 @@ class ModuleSanitizerCoverageLTO LLVMContext *Ct = NULL; Module *Mo = NULL; GlobalVariable *AFLMapPtr = NULL; + GlobalVariable *AFLContext = NULL; Value *MapPtrFixed = NULL; std::ofstream dFile; size_t found = 0; @@ -420,11 +422,13 @@ bool ModuleSanitizerCoverageLTO::instrumentModule( setvbuf(stdout, NULL, _IONBF, 0); if (getenv("AFL_DEBUG")) { debug = 1; } if (getenv("AFL_LLVM_DICT2FILE_NO_MAIN")) { autodictionary_no_main = 1; } + if (getenv("AFL_LLVM_CALLER")) { instrument_ctx = 1; } if ((isatty(2) && !getenv("AFL_QUIET")) || debug) { SAYF(cCYA "afl-llvm-lto" VERSION cRST - " by Marc \"vanHauser\" Heuse \n"); + "%s by Marc \"vanHauser\" Heuse \n", + instrument_ctx ? " (CTX mode)" : ""); } else @@ -500,6 +504,10 @@ bool ModuleSanitizerCoverageLTO::instrumentModule( } + AFLContext = new GlobalVariable( + M, Int32Ty, false, GlobalValue::ExternalLinkage, 0, "__afl_prev_ctx", 0, + GlobalVariable::GeneralDynamicTLSModel, 0, false); + Zero = ConstantInt::get(Int8Tyi, 0); One = ConstantInt::get(Int8Tyi, 1); @@ -1284,7 +1292,50 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( const DominatorTree *DT = DTCallback(F); const PostDominatorTree *PDT = PDTCallback(F); bool IsLeafFunc = true; - uint32_t skip_next = 0; + uint32_t skip_next = 0, call_counter = 0; + Value *PrevCtx = NULL; + + MDNode *N = + MDNode::get(F.getContext(), MDString::get(F.getContext(), "nosanitize")); + + for (auto &BB : F) { + + if (/*F.size() > 1 &&*/ instrument_ctx && &BB == &F.getEntryBlock()) { + + // we insert a CTX value in all our callers: + LLVMContext &Context = F.getContext(); + IRBuilder<> Builder(Context); + for (auto *U : F.users()) { + + if (auto *CI = dyn_cast(U)) { + + fprintf(stderr, "Insert %s [%u] -> %s\n", + CI->getParent()->getParent()->getName().str().c_str(), + call_counter, F.getName().str().c_str()); + Builder.SetInsertPoint(CI); + StoreInst *StoreCtx = Builder.CreateStore( + ConstantInt::get(Type::getInt32Ty(Context), call_counter++), + AFLContext); + StoreCtx->setMetadata("nosanitize", N); + + } + + } + + // We read the CTX for this call + BasicBlock::iterator IP = BB.getFirstInsertionPt(); + IRBuilder<> IRB(&(*IP)); + LoadInst *PrevCtxLoad = IRB.CreateLoad( +#if LLVM_VERSION_MAJOR >= 14 + Builder.getInt32Ty(), +#endif + AFLContext); + PrevCtxLoad->setMetadata("nosanitize", N); + PrevCtx = PrevCtxLoad; + + } + + } for (auto &BB : F) { diff --git a/src/afl-cc.c b/src/afl-cc.c index 174b3783..4f6745ed 100644 --- a/src/afl-cc.c +++ b/src/afl-cc.c @@ -1103,12 +1103,18 @@ static void instrument_opt_mode_exclude(aflcc_state_t *aflcc) { } - if (aflcc->instrument_opt_mode && aflcc->compiler_mode != LLVM) + fprintf(stderr, "X %u %u\n", aflcc->compiler_mode, LTO); + + if (aflcc->instrument_opt_mode && aflcc->compiler_mode != LLVM && + !((aflcc->instrument_opt_mode & INSTRUMENT_OPT_CALLER) && + aflcc->compiler_mode == LTO)) FATAL("CTX, CALLER and NGRAM can only be used in LLVM mode"); if (aflcc->instrument_opt_mode && aflcc->instrument_opt_mode != INSTRUMENT_OPT_CODECOV && - aflcc->instrument_mode != INSTRUMENT_CLASSIC) + aflcc->instrument_mode != INSTRUMENT_CLASSIC && + !(aflcc->instrument_opt_mode & INSTRUMENT_OPT_CALLER && + aflcc->compiler_mode == LTO)) FATAL( "CALLER, CTX and NGRAM instrumentation options can only be used with " "the LLVM CLASSIC instrumentation mode."); From b0a912a83881052b3ce476459d8c8edfab59c2f9 Mon Sep 17 00:00:00 2001 From: vanhauser-thc Date: Fri, 26 Jan 2024 12:15:42 +0100 Subject: [PATCH 2/9] working ugly version --- instrumentation/SanitizerCoverageLTO.so.cc | 474 +++++++++++++++------ 1 file changed, 343 insertions(+), 131 deletions(-) diff --git a/instrumentation/SanitizerCoverageLTO.so.cc b/instrumentation/SanitizerCoverageLTO.so.cc index c74069e1..54cc1752 100644 --- a/instrumentation/SanitizerCoverageLTO.so.cc +++ b/instrumentation/SanitizerCoverageLTO.so.cc @@ -192,12 +192,14 @@ class ModuleSanitizerCoverageLTO PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM); private: - void instrumentFunction(Function &F, DomTreeCallback DTCallback, - PostDomTreeCallback PDTCallback); - void InjectCoverageForIndirectCalls(Function &F, - ArrayRef IndirCalls); - bool InjectCoverage(Function &F, ArrayRef AllBlocks, - bool IsLeafFunc = true); + void instrumentFunction(Function &F, DomTreeCallback DTCallback, + PostDomTreeCallback PDTCallback); + void InjectCoverageForIndirectCalls(Function &F, + ArrayRef IndirCalls); + bool InjectCoverage(Function &F, ArrayRef AllBlocks, + bool IsLeafFunc = true); + bool Fake_InjectCoverage(Function &F, ArrayRef AllBlocks, + bool IsLeafFunc = true); GlobalVariable *CreateFunctionLocalArrayInSection(size_t NumElements, Function &F, Type *Ty, const char *Section); @@ -248,6 +250,7 @@ class ModuleSanitizerCoverageLTO uint32_t unhandled = 0; uint32_t select_cnt = 0; uint32_t instrument_ctx = 0; + uint32_t extra_ctx_inst = 0; uint64_t map_addr = 0; const char *skip_nozero = NULL; const char *use_threadsafe_counters = nullptr; @@ -259,6 +262,7 @@ class ModuleSanitizerCoverageLTO IntegerType *Int64Tyi = NULL; ConstantInt *Zero = NULL; ConstantInt *One = NULL; + AllocaInst *CTX_add = NULL; LLVMContext *Ct = NULL; Module *Mo = NULL; GlobalVariable *AFLMapPtr = NULL; @@ -1164,9 +1168,10 @@ bool ModuleSanitizerCoverageLTO::instrumentModule( getenv("AFL_USE_TSAN") ? ", TSAN" : "", getenv("AFL_USE_CFISAN") ? ", CFISAN" : "", getenv("AFL_USE_UBSAN") ? ", UBSAN" : ""); - OKF("Instrumented %u locations (%u selects) without collisions (%llu " - "collisions have been avoided) (%s mode).", - inst, select_cnt, calculateCollisions(inst), modeline); + + OKF("Instrumented %u locations (%u selects) with %u extra map entries " + "for CTX (%s mode).", + inst, select_cnt, extra_ctx_inst, modeline); } @@ -1292,19 +1297,112 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( const DominatorTree *DT = DTCallback(F); const PostDominatorTree *PDT = PDTCallback(F); bool IsLeafFunc = true; - uint32_t skip_next = 0, call_counter = 0; - Value *PrevCtx = NULL; + uint32_t skip_next = 0; + uint32_t call_counter = 0; + uint32_t inst_save = inst; + uint32_t inst_in_this_func = 0; + LLVMContext &Context = F.getContext(); - MDNode *N = - MDNode::get(F.getContext(), MDString::get(F.getContext(), "nosanitize")); + MDNode *N = MDNode::get(Context, MDString::get(Context, "nosanitize")); + CTX_add = NULL; + + fprintf(stderr, "Function: %s\n", F.getName().str().c_str()); + + // Fake instrumentation so we can count how much there will be in this + // function for (auto &BB : F) { - if (/*F.size() > 1 &&*/ instrument_ctx && &BB == &F.getEntryBlock()) { + for (auto &IN : BB) { + + CallInst *callInst = nullptr; + + if ((callInst = dyn_cast(&IN))) { + + Function *Callee = callInst->getCalledFunction(); + if (!Callee) continue; + if (callInst->getCallingConv() != llvm::CallingConv::C) continue; + StringRef FuncName = Callee->getName(); + + if (FuncName.compare(StringRef("__afl_coverage_interesting"))) continue; + + ++inst; + + } + + SelectInst *selectInst = nullptr; + + if (!skip_next && (selectInst = dyn_cast(&IN)) && 1 == 0) { + + uint32_t vector_cnt = 0; + Value *condition = selectInst->getCondition(); + auto t = condition->getType(); + + if (t->getTypeID() == llvm::Type::IntegerTyID) { + + skip_next = 1; + inst += 2; + + } else + +#if LLVM_VERSION_MAJOR >= 14 + if (t->getTypeID() == llvm::Type::FixedVectorTyID) { + + FixedVectorType *tt = dyn_cast(t); + if (tt) { + + uint32_t elements = tt->getElementCount().getFixedValue(); + vector_cnt = elements; + inst += vector_cnt * 2; + + } + + } else + +#endif + { + + continue; + + } + + skip_next = 1; + + } else { + + skip_next = 0; + + } + + } + + if (shouldInstrumentBlock(F, &BB, DT, PDT, Options)) + BlocksToInstrument.push_back(&BB); + /* + if (Options.IndirectCalls) + for (auto &Inst : BB) { + + fprintf(stderr, "FUCK TODO\n"); + CallBase *CB = dyn_cast(&Inst); + if (CB && !CB->getCalledFunction()) IndirCalls.push_back(&Inst); + /// XXX TODO!!!! set 0 ! + + } + + */ + + } + + Fake_InjectCoverage(F, BlocksToInstrument, IsLeafFunc); + + // CTX init + + for (auto &BB : F) { + + if (instrument_ctx && &BB == &F.getEntryBlock()) { // we insert a CTX value in all our callers: - LLVMContext &Context = F.getContext(); - IRBuilder<> Builder(Context); + IRBuilder<> Builder(Context); for (auto *U : F.users()) { if (auto *CI = dyn_cast(U)) { @@ -1323,24 +1421,60 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( } // We read the CTX for this call + Value *CTX_offset; BasicBlock::iterator IP = BB.getFirstInsertionPt(); IRBuilder<> IRB(&(*IP)); LoadInst *PrevCtxLoad = IRB.CreateLoad( #if LLVM_VERSION_MAJOR >= 14 - Builder.getInt32Ty(), + IRB.getInt32Ty(), #endif AFLContext); + PrevCtxLoad->setMetadata("nosanitize", N); - PrevCtx = PrevCtxLoad; + + if (inst == inst_save || call_counter < 2) { + + fprintf(stderr, "%s: ZERO!\n", F.getName().str().c_str()); + CTX_offset = Zero; + + } else { + + fprintf(stderr, "%s: %u * %u\n", F.getName().str().c_str(), + inst - inst_save, call_counter); + CTX_offset = IRB.CreateMul( + ConstantInt::get(Type::getInt32Ty(Context), inst - inst_save), + PrevCtxLoad, "CTXmul", false, true); + + } + + CTX_add = IRB.CreateAlloca(Type::getInt32Ty(Context), nullptr, "CTX_add"); + auto nosan = IRB.CreateStore(CTX_offset, CTX_add); + nosan->setMetadata("nosanitize", N); } } + inst_in_this_func = inst - inst_save; + inst = inst_save; + + // Now the real instrumentation + + IsLeafFunc = true; + skip_next = 0; + for (auto &BB : F) { + // fprintf(stderr, "BB: %s\n", BB.getName().str().c_str()); + for (auto &IN : BB) { + /* + std::string Str; + llvm::raw_string_ostream RSO(Str); + IN.print(RSO); + errs() << RSO.str() << "\n"; + */ CallInst *callInst = nullptr; if ((callInst = dyn_cast(&IN))) { @@ -1363,8 +1497,16 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( if (FuncName.compare(StringRef("__afl_coverage_interesting"))) continue; - Value *val = ConstantInt::get(Int32Ty, ++afl_global_id); - callInst->setOperand(1, val); + IRBuilder<> Builder(Context); + Value *val = ConstantInt::get(Int32Ty, ++afl_global_id); + LoadInst *CTX_load = Builder.CreateLoad( +#if LLVM_VERSION_MAJOR >= 14 + Builder.getInt32Ty(), +#endif + CTX_add); + ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(CTX_load); + Value *val2 = Builder.CreateAdd(val, CTX_load); + callInst->setOperand(1, val2); ++inst; } @@ -1377,158 +1519,193 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( IN.print(os); fprintf(stderr, "X(%u): %s\n", skip_next, os.str().c_str()); */ - if (!skip_next && (selectInst = dyn_cast(&IN))) { + if ((selectInst = dyn_cast(&IN))) { - uint32_t vector_cnt = 0; - Value *condition = selectInst->getCondition(); - Value *result; - auto t = condition->getType(); - IRBuilder<> IRB(selectInst->getNextNode()); + if (!skip_next) { - ++select_cnt; + uint32_t vector_cnt = 0; + Value *condition = selectInst->getCondition(); + Value *result; + auto t = condition->getType(); + IRBuilder<> IRB(selectInst->getNextNode()); - if (t->getTypeID() == llvm::Type::IntegerTyID) { + // fprintf(stderr, "ADDING!!!\n"); + LoadInst *CTX_load = IRB.CreateLoad( +#if LLVM_VERSION_MAJOR >= 14 + IRB.getInt32Ty(), +#endif + CTX_add); + ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(CTX_load); - Value *val1 = ConstantInt::get(Int32Ty, ++afl_global_id); - Value *val2 = ConstantInt::get(Int32Ty, ++afl_global_id); - result = IRB.CreateSelect(condition, val1, val2); - skip_next = 1; - inst += 2; + ++select_cnt; - } else + if (t->getTypeID() == llvm::Type::IntegerTyID) { + + Value *val1 = ConstantInt::get(Int32Ty, ++afl_global_id); + Value *val11 = IRB.CreateAdd(val1, CTX_load); + Value *val2 = ConstantInt::get(Int32Ty, ++afl_global_id); + Value *val22 = IRB.CreateAdd(val2, CTX_load); + result = IRB.CreateSelect(condition, val11, val22); + skip_next = 1; + inst += 2; + + } else #if LLVM_VERSION_MAJOR >= 14 - if (t->getTypeID() == llvm::Type::FixedVectorTyID) { + if (t->getTypeID() == llvm::Type::FixedVectorTyID) { - FixedVectorType *tt = dyn_cast(t); - if (tt) { + FixedVectorType *tt = dyn_cast(t); + if (tt) { - uint32_t elements = tt->getElementCount().getFixedValue(); - vector_cnt = elements; - inst += vector_cnt * 2; - if (elements) { + 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; + 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); + Value *val1 = ConstantInt::get(Int32Ty, ++afl_global_id); + Value *val11 = IRB.CreateAdd(val1, CTX_add); + Value *val2 = ConstantInt::get(Int32Ty, ++afl_global_id); + Value *val22 = IRB.CreateAdd(val2, CTX_add); + x = IRB.CreateInsertElement(GuardPtr1, val11, (uint64_t)0); + y = IRB.CreateInsertElement(GuardPtr2, val22, (uint64_t)0); - for (uint64_t i = 1; i < elements; i++) { + 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); + val1 = ConstantInt::get(Int32Ty, ++afl_global_id); + val11 = IRB.CreateAdd(val1, CTX_add); + val2 = ConstantInt::get(Int32Ty, ++afl_global_id); + val11 = IRB.CreateAdd(val1, CTX_add); + x = IRB.CreateInsertElement(GuardPtr1, val11, i); + y = IRB.CreateInsertElement(GuardPtr2, val22, i); + + } + + result = IRB.CreateSelect(condition, x, y); + skip_next = 1; } - result = IRB.CreateSelect(condition, x, y); - skip_next = 1; + } + + } else + +#endif + { + + unhandled++; + continue; + + } + + uint32_t vector_cur = 0; + /* Load SHM pointer */ + LoadInst *MapPtr = + IRB.CreateLoad(PointerType::get(Int8Ty, 0), AFLMapPtr); + ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(MapPtr); + + while (1) { + + /* Get CurLoc */ + Value *MapPtrIdx = nullptr; + + /* Load counter for CurLoc */ + if (!vector_cnt) { + + MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, result); + + } else { + + auto element = IRB.CreateExtractElement(result, vector_cur++); + MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, element); } - } + if (use_threadsafe_counters) { - } else - -#endif - { - - unhandled++; - continue; - - } - - uint32_t vector_cur = 0; - /* Load SHM pointer */ - LoadInst *MapPtr = - IRB.CreateLoad(PointerType::get(Int8Ty, 0), AFLMapPtr); - ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(MapPtr); - - while (1) { - - /* Get CurLoc */ - Value *MapPtrIdx = nullptr; - - /* Load counter for CurLoc */ - if (!vector_cnt) { - - MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, result); - - } else { - - auto element = IRB.CreateExtractElement(result, vector_cur++); - MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, element); - - } - - if (use_threadsafe_counters) { - - IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One, + IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, + One, #if LLVM_VERSION_MAJOR >= 13 - llvm::MaybeAlign(1), + llvm::MaybeAlign(1), #endif - llvm::AtomicOrdering::Monotonic); + llvm::AtomicOrdering::Monotonic); - } else { + } else { - LoadInst *Counter = IRB.CreateLoad(IRB.getInt8Ty(), MapPtrIdx); - ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(Counter); + LoadInst *Counter = IRB.CreateLoad(IRB.getInt8Ty(), MapPtrIdx); + ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(Counter); - /* Update bitmap */ + /* Update bitmap */ - Value *Incr = IRB.CreateAdd(Counter, One); + Value *Incr = IRB.CreateAdd(Counter, One); - if (skip_nozero == NULL) { + if (skip_nozero == NULL) { - auto cf = IRB.CreateICmpEQ(Incr, Zero); - auto carry = IRB.CreateZExt(cf, Int8Ty); - Incr = IRB.CreateAdd(Incr, carry); + auto cf = IRB.CreateICmpEQ(Incr, Zero); + auto carry = IRB.CreateZExt(cf, Int8Ty); + Incr = IRB.CreateAdd(Incr, carry); + + } + + auto nosan = IRB.CreateStore(Incr, MapPtrIdx); + ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(nosan); } - auto nosan = IRB.CreateStore(Incr, MapPtrIdx); - ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(nosan); + if (!vector_cnt || vector_cnt == vector_cur) { break; } } - if (!vector_cnt || vector_cnt == vector_cur) { break; } + skip_next = 1; + + } else { + + skip_next = 0; } - skip_next = 1; - - } else { - - skip_next = 0; - } } - if (shouldInstrumentBlock(F, &BB, DT, PDT, Options)) - BlocksToInstrument.push_back(&BB); - for (auto &Inst : BB) { - - if (Options.IndirectCalls) { - - CallBase *CB = dyn_cast(&Inst); - if (CB && !CB->getCalledFunction()) IndirCalls.push_back(&Inst); - - } - - } + // if (shouldInstrumentBlock(F, &BB, DT, PDT, Options)) + // BlocksToInstrument.push_back(&BB); } InjectCoverage(F, BlocksToInstrument, IsLeafFunc); - InjectCoverageForIndirectCalls(F, IndirCalls); + // InjectCoverageForIndirectCalls(F, IndirCalls); + + if (inst_in_this_func && call_counter > 1) { + + extra_ctx_inst += inst_in_this_func * (call_counter - 1); + afl_global_id += inst_in_this_func * (call_counter - 1); + + } + + /* + fprintf(stderr, "FUNCTION: %s\n", F.getName().str().c_str()); + int n = 0; + for (auto &BB : F) { + + fprintf(stderr, "BB %d\n", n++); + for (auto &IN : BB) { + + std::string Str; + llvm::raw_string_ostream RSO(Str); + IN.print(RSO); + errs() << RSO.str() << "\n"; + + } + + } + + */ } @@ -1654,6 +1831,36 @@ bool ModuleSanitizerCoverageLTO::InjectCoverage( } +// Fake InjectCoverage to count the instrumentations +bool ModuleSanitizerCoverageLTO::Fake_InjectCoverage( + Function &F, ArrayRef AllBlocks, bool IsLeafFunc) { + + if (AllBlocks.empty()) return false; + + for (size_t i = 0, N = AllBlocks.size(); i < N; i++) { + + // AFL++ START + if (BlockList.size()) { + + int skip = 0; + for (uint32_t k = 0; k < BlockList.size(); k++) { + + if (AllBlocks[i] == BlockList[k]) { skip = 1; } + + } + + if (skip) continue; + + } + + inst++; // InjectCoverageAtBlock() + + } + + return true; + +} + // On every indirect call we call a run-time function // __sanitizer_cov_indir_call* with two parameters: // - callee address, @@ -1725,6 +1932,13 @@ void ModuleSanitizerCoverageLTO::InjectCoverageAtBlock(Function &F, /* Set the ID of the inserted basic block */ ConstantInt *CurLoc = ConstantInt::get(Int32Tyi, afl_global_id); + LoadInst *CTX_load = IRB.CreateLoad( +#if LLVM_VERSION_MAJOR >= 14 + IRB.getInt32Ty(), +#endif + CTX_add); + ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(CTX_load); + Value *offset = IRB.CreateAdd(CurLoc, CTX_load); /* Load SHM pointer */ @@ -1732,13 +1946,13 @@ void ModuleSanitizerCoverageLTO::InjectCoverageAtBlock(Function &F, if (map_addr) { - MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtrFixed, CurLoc); + MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtrFixed, offset); } else { LoadInst *MapPtr = IRB.CreateLoad(PointerType::get(Int8Ty, 0), AFLMapPtr); ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(MapPtr); - MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, CurLoc); + MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, offset); } @@ -1777,8 +1991,6 @@ void ModuleSanitizerCoverageLTO::InjectCoverageAtBlock(Function &F, // AFL++ END /* - XXXXXXXXXXXXXXXXXXX - auto GuardPtr = IRB.CreateIntToPtr( IRB.CreateAdd(IRB.CreatePointerCast(FunctionGuardArray, IntptrTy), ConstantInt::get(IntptrTy, Idx * 4)), From 44a7696169f52f6ef8b5c9a5a6de1167000e2138 Mon Sep 17 00:00:00 2001 From: vanhauser-thc Date: Fri, 26 Jan 2024 15:27:20 +0100 Subject: [PATCH 3/9] fixes --- instrumentation/SanitizerCoverageLTO.so.cc | 42 +++++++++++++++------- 1 file changed, 30 insertions(+), 12 deletions(-) diff --git a/instrumentation/SanitizerCoverageLTO.so.cc b/instrumentation/SanitizerCoverageLTO.so.cc index 54cc1752..b280e947 100644 --- a/instrumentation/SanitizerCoverageLTO.so.cc +++ b/instrumentation/SanitizerCoverageLTO.so.cc @@ -1453,6 +1453,24 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( } + // we have to set __afl_ctx 0 for all indirect calls + for (auto &IN : BB) { + + if (auto *Call = dyn_cast(&IN)) { + + if (Call->isIndirectCall()) { + + IRBuilder<> Builder(IN.getContext()); + Builder.SetInsertPoint(IN.getParent(), IN.getIterator()); + StoreInst *StoreCtx = Builder.CreateStore(Zero, AFLContext); + StoreCtx->setMetadata("nosanitize", N); + + } + + } + + } + } inst_in_this_func = inst - inst_save; @@ -1569,18 +1587,18 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( Value *x, *y; Value *val1 = ConstantInt::get(Int32Ty, ++afl_global_id); - Value *val11 = IRB.CreateAdd(val1, CTX_add); + Value *val11 = IRB.CreateAdd(val1, CTX_load); Value *val2 = ConstantInt::get(Int32Ty, ++afl_global_id); - Value *val22 = IRB.CreateAdd(val2, CTX_add); + Value *val22 = IRB.CreateAdd(val2, CTX_load); x = IRB.CreateInsertElement(GuardPtr1, val11, (uint64_t)0); y = IRB.CreateInsertElement(GuardPtr2, val22, (uint64_t)0); for (uint64_t i = 1; i < elements; i++) { val1 = ConstantInt::get(Int32Ty, ++afl_global_id); - val11 = IRB.CreateAdd(val1, CTX_add); + val11 = IRB.CreateAdd(val1, CTX_load); val2 = ConstantInt::get(Int32Ty, ++afl_global_id); - val11 = IRB.CreateAdd(val1, CTX_add); + val11 = IRB.CreateAdd(val1, CTX_load); x = IRB.CreateInsertElement(GuardPtr1, val11, i); y = IRB.CreateInsertElement(GuardPtr2, val22, i); @@ -1628,12 +1646,13 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( if (use_threadsafe_counters) { - IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, - One, + auto nosan = IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, + MapPtrIdx, One, #if LLVM_VERSION_MAJOR >= 13 - llvm::MaybeAlign(1), + llvm::MaybeAlign(1), #endif - llvm::AtomicOrdering::Monotonic); + llvm::AtomicOrdering::Monotonic); + ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(nosan); } else { @@ -1684,14 +1703,13 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( if (inst_in_this_func && call_counter > 1) { extra_ctx_inst += inst_in_this_func * (call_counter - 1); - afl_global_id += inst_in_this_func * (call_counter - 1); + afl_global_id += extra_ctx_inst; } /* - fprintf(stderr, "FUNCTION: %s\n", F.getName().str().c_str()); - int n = 0; - for (auto &BB : F) { + fprintf(stderr, "FUNCTION: %s [%u]\n", F.getName().str().c_str(), + extra_ctx_inst); int n = 0; for (auto &BB : F) { fprintf(stderr, "BB %d\n", n++); for (auto &IN : BB) { From d668010bedf5373e25ea12c24dbb477f54da91ba Mon Sep 17 00:00:00 2001 From: vanhauser-thc Date: Fri, 26 Jan 2024 16:44:31 +0100 Subject: [PATCH 4/9] fixes --- instrumentation/SanitizerCoverageLTO.so.cc | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/instrumentation/SanitizerCoverageLTO.so.cc b/instrumentation/SanitizerCoverageLTO.so.cc index b280e947..a3074ae2 100644 --- a/instrumentation/SanitizerCoverageLTO.so.cc +++ b/instrumentation/SanitizerCoverageLTO.so.cc @@ -249,7 +249,7 @@ class ModuleSanitizerCoverageLTO uint32_t afl_global_id = 0; uint32_t unhandled = 0; uint32_t select_cnt = 0; - uint32_t instrument_ctx = 0; + uint32_t instrument_ctx = 1; uint32_t extra_ctx_inst = 0; uint64_t map_addr = 0; const char *skip_nozero = NULL; @@ -1481,6 +1481,18 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( IsLeafFunc = true; skip_next = 0; + if (CTX_add == NULL) { + + auto BB = &F.getEntryBlock(); + fprintf(stderr, "NULL %s %p\n", F.getName().str().c_str(), BB); + if (!BB) { exit(-1); } + BasicBlock::iterator IP = BB->getFirstInsertionPt(); + IRBuilder<> IRB(&(*IP)); + CTX_add = IRB.CreateAlloca(Type::getInt32Ty(Context), nullptr, "CTX_add"); + auto nosan = IRB.CreateStore(Zero, CTX_add); + nosan->setMetadata("nosanitize", N); + } + for (auto &BB : F) { // fprintf(stderr, "BB: %s\n", BB.getName().str().c_str()); From ceb7e44e6fb614c3efd684b6a1ee71fa89a6fe24 Mon Sep 17 00:00:00 2001 From: vanhauser-thc Date: Sat, 27 Jan 2024 08:28:47 +0100 Subject: [PATCH 5/9] fixes --- instrumentation/SanitizerCoverageLTO.so.cc | 56 ++++++++++++---------- 1 file changed, 30 insertions(+), 26 deletions(-) diff --git a/instrumentation/SanitizerCoverageLTO.so.cc b/instrumentation/SanitizerCoverageLTO.so.cc index a3074ae2..3a02cf08 100644 --- a/instrumentation/SanitizerCoverageLTO.so.cc +++ b/instrumentation/SanitizerCoverageLTO.so.cc @@ -1299,7 +1299,7 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( bool IsLeafFunc = true; uint32_t skip_next = 0; uint32_t call_counter = 0; - uint32_t inst_save = inst; + uint32_t inst_save = inst, save_global = afl_global_id; uint32_t inst_in_this_func = 0; LLVMContext &Context = F.getContext(); @@ -1332,15 +1332,13 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( SelectInst *selectInst = nullptr; - if (!skip_next && (selectInst = dyn_cast(&IN)) && 1 == 0) { + if ((selectInst = dyn_cast(&IN))) { - uint32_t vector_cnt = 0; - Value *condition = selectInst->getCondition(); - auto t = condition->getType(); + Value *condition = selectInst->getCondition(); + auto t = condition->getType(); if (t->getTypeID() == llvm::Type::IntegerTyID) { - skip_next = 1; inst += 2; } else @@ -1352,8 +1350,7 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( if (tt) { uint32_t elements = tt->getElementCount().getFixedValue(); - vector_cnt = elements; - inst += vector_cnt * 2; + inst += elements * 2; } @@ -1366,12 +1363,6 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( } - skip_next = 1; - - } else { - - skip_next = 0; - } } @@ -1478,19 +1469,22 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( // Now the real instrumentation - IsLeafFunc = true; - skip_next = 0; - if (CTX_add == NULL) { - auto BB = &F.getEntryBlock(); + auto BB = &F.getEntryBlock(); + if (!BB) { + fprintf(stderr, "NULL %s %p\n", F.getName().str().c_str(), BB); - if (!BB) { exit(-1); } - BasicBlock::iterator IP = BB->getFirstInsertionPt(); - IRBuilder<> IRB(&(*IP)); - CTX_add = IRB.CreateAlloca(Type::getInt32Ty(Context), nullptr, "CTX_add"); - auto nosan = IRB.CreateStore(Zero, CTX_add); - nosan->setMetadata("nosanitize", N); + exit(-1); + + } + + BasicBlock::iterator IP = BB->getFirstInsertionPt(); + IRBuilder<> IRB(&(*IP)); + CTX_add = IRB.CreateAlloca(Type::getInt32Ty(Context), nullptr, "CTX_add"); + auto nosan = IRB.CreateStore(Zero, CTX_add); + nosan->setMetadata("nosanitize", N); + } for (auto &BB : F) { @@ -1590,6 +1584,7 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( uint32_t elements = tt->getElementCount().getFixedValue(); vector_cnt = elements; inst += vector_cnt * 2; + if (elements) { FixedVectorType *GuardPtr1 = @@ -1714,6 +1709,15 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( if (inst_in_this_func && call_counter > 1) { + if (inst_in_this_func != afl_global_id - save_global) { + + fprintf( + stderr, + "BUG! inst_in_this_func %u != afl_global_id %u - save_global %u\n", + inst_in_this_func, afl_global_id, save_global); + + } + extra_ctx_inst += inst_in_this_func * (call_counter - 1); afl_global_id += extra_ctx_inst; @@ -1883,7 +1887,7 @@ bool ModuleSanitizerCoverageLTO::Fake_InjectCoverage( } - inst++; // InjectCoverageAtBlock() + ++inst; // InjectCoverageAtBlock() } @@ -2017,7 +2021,7 @@ void ModuleSanitizerCoverageLTO::InjectCoverageAtBlock(Function &F, // done :) - inst++; + ++inst; // AFL++ END /* From e6eee685ceedf92d5ce6dd1c32412e286b4f6104 Mon Sep 17 00:00:00 2001 From: vanhauser-thc Date: Sat, 27 Jan 2024 15:13:27 +0100 Subject: [PATCH 6/9] fix --- instrumentation/SanitizerCoverageLTO.so.cc | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/instrumentation/SanitizerCoverageLTO.so.cc b/instrumentation/SanitizerCoverageLTO.so.cc index 3a02cf08..469df42e 100644 --- a/instrumentation/SanitizerCoverageLTO.so.cc +++ b/instrumentation/SanitizerCoverageLTO.so.cc @@ -261,6 +261,7 @@ class ModuleSanitizerCoverageLTO IntegerType *Int32Tyi = NULL; IntegerType *Int64Tyi = NULL; ConstantInt *Zero = NULL; + ConstantInt *Zero32 = NULL; ConstantInt *One = NULL; AllocaInst *CTX_add = NULL; LLVMContext *Ct = NULL; @@ -513,6 +514,7 @@ bool ModuleSanitizerCoverageLTO::instrumentModule( GlobalVariable::GeneralDynamicTLSModel, 0, false); Zero = ConstantInt::get(Int8Tyi, 0); + Zero32 = ConstantInt::get(Int32Tyi, 0); One = ConstantInt::get(Int8Tyi, 1); initInstrumentList(); @@ -1426,7 +1428,7 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( if (inst == inst_save || call_counter < 2) { fprintf(stderr, "%s: ZERO!\n", F.getName().str().c_str()); - CTX_offset = Zero; + CTX_offset = Zero32; } else { @@ -1453,7 +1455,7 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( IRBuilder<> Builder(IN.getContext()); Builder.SetInsertPoint(IN.getParent(), IN.getIterator()); - StoreInst *StoreCtx = Builder.CreateStore(Zero, AFLContext); + StoreInst *StoreCtx = Builder.CreateStore(Zero32, AFLContext); StoreCtx->setMetadata("nosanitize", N); } @@ -1474,7 +1476,7 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( auto BB = &F.getEntryBlock(); if (!BB) { - fprintf(stderr, "NULL %s %p\n", F.getName().str().c_str(), BB); + fprintf(stderr, "NULL entry %s %p\n", F.getName().str().c_str(), BB); exit(-1); } @@ -1482,7 +1484,7 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( BasicBlock::iterator IP = BB->getFirstInsertionPt(); IRBuilder<> IRB(&(*IP)); CTX_add = IRB.CreateAlloca(Type::getInt32Ty(Context), nullptr, "CTX_add"); - auto nosan = IRB.CreateStore(Zero, CTX_add); + auto nosan = IRB.CreateStore(Zero32, CTX_add); nosan->setMetadata("nosanitize", N); } From 79080355ac1be8f4edc9818daf5530691aad7f3c Mon Sep 17 00:00:00 2001 From: vanhauser-thc Date: Thu, 1 Feb 2024 17:39:23 +0100 Subject: [PATCH 7/9] better CTX instrumentation --- instrumentation/SanitizerCoverageLTO.so.cc | 815 ++++++++++++--------- 1 file changed, 479 insertions(+), 336 deletions(-) diff --git a/instrumentation/SanitizerCoverageLTO.so.cc b/instrumentation/SanitizerCoverageLTO.so.cc index 469df42e..f6d60099 100644 --- a/instrumentation/SanitizerCoverageLTO.so.cc +++ b/instrumentation/SanitizerCoverageLTO.so.cc @@ -194,8 +194,9 @@ class ModuleSanitizerCoverageLTO private: void instrumentFunction(Function &F, DomTreeCallback DTCallback, PostDomTreeCallback PDTCallback); - void InjectCoverageForIndirectCalls(Function &F, - ArrayRef IndirCalls); + /* void InjectCoverageForIndirectCalls(Function &F, + ArrayRef + IndirCalls);*/ bool InjectCoverage(Function &F, ArrayRef AllBlocks, bool IsLeafFunc = true); bool Fake_InjectCoverage(Function &F, ArrayRef AllBlocks, @@ -263,12 +264,12 @@ class ModuleSanitizerCoverageLTO ConstantInt *Zero = NULL; ConstantInt *Zero32 = NULL; ConstantInt *One = NULL; - AllocaInst *CTX_add = NULL; LLVMContext *Ct = NULL; Module *Mo = NULL; - GlobalVariable *AFLMapPtr = NULL; GlobalVariable *AFLContext = NULL; + GlobalVariable *AFLMapPtr = NULL; Value *MapPtrFixed = NULL; + AllocaInst *CTX_add = NULL; std::ofstream dFile; size_t found = 0; // AFL++ END @@ -427,7 +428,11 @@ bool ModuleSanitizerCoverageLTO::instrumentModule( setvbuf(stdout, NULL, _IONBF, 0); if (getenv("AFL_DEBUG")) { debug = 1; } if (getenv("AFL_LLVM_DICT2FILE_NO_MAIN")) { autodictionary_no_main = 1; } - if (getenv("AFL_LLVM_CALLER")) { instrument_ctx = 1; } + if (getenv("AFL_LLVM_CALLER") || getenv("AFL_LLVM_CTX")) { + + instrument_ctx = 1; + + } if ((isatty(2) && !getenv("AFL_QUIET")) || debug) { @@ -435,10 +440,12 @@ bool ModuleSanitizerCoverageLTO::instrumentModule( "%s by Marc \"vanHauser\" Heuse \n", instrument_ctx ? " (CTX mode)" : ""); - } else + } else { be_quiet = 1; + } + skip_nozero = getenv("AFL_LLVM_SKIP_NEVERZERO"); use_threadsafe_counters = getenv("AFL_LLVM_THREADSAFE_INST"); @@ -611,12 +618,12 @@ bool ModuleSanitizerCoverageLTO::instrumentModule( } dictionary.push_back(std::string((char *)&val, len)); - found++; + ++found; if (val2) { dictionary.push_back(std::string((char *)&val2, len)); - found++; + ++found; } @@ -943,7 +950,7 @@ bool ModuleSanitizerCoverageLTO::instrumentModule( '\0') { thestring.append("\0", 1); // add null byte - optLen++; + ++optLen; } @@ -1094,7 +1101,7 @@ bool ModuleSanitizerCoverageLTO::instrumentModule( for (auto token : dictionary) { memlen += token.length(); - count++; + ++count; } @@ -1115,7 +1122,7 @@ bool ModuleSanitizerCoverageLTO::instrumentModule( ptrhld.get()[offset++] = (uint8_t)token.length(); memcpy(ptrhld.get() + offset, token.c_str(), token.length()); offset += token.length(); - count++; + ++count; } @@ -1162,7 +1169,7 @@ bool ModuleSanitizerCoverageLTO::instrumentModule( WARNF("No instrumentation targets found."); else { - char modeline[100]; + char modeline[128]; snprintf(modeline, sizeof(modeline), "%s%s%s%s%s%s", getenv("AFL_HARDEN") ? "hardened" : "non-hardened", getenv("AFL_USE_ASAN") ? ", ASAN" : "", @@ -1170,10 +1177,16 @@ bool ModuleSanitizerCoverageLTO::instrumentModule( getenv("AFL_USE_TSAN") ? ", TSAN" : "", getenv("AFL_USE_CFISAN") ? ", CFISAN" : "", getenv("AFL_USE_UBSAN") ? ", UBSAN" : ""); + char buf[64] = {}; + if (instrument_ctx) { - OKF("Instrumented %u locations (%u selects) with %u extra map entries " - "for CTX (%s mode).", - inst, select_cnt, extra_ctx_inst, modeline); + snprintf(buf, sizeof(buf), " with %u extra map entries for CTX", + extra_ctx_inst); + + } + + OKF("Instrumented %u locations (%u selects)%s (%s mode).", inst, + select_cnt, buf, modeline); } @@ -1254,6 +1267,52 @@ static bool shouldInstrumentBlock(const Function &F, const BasicBlock *BB, } +/// return the number of calls to this function +u32 countCallers(Function *F) { + + u32 callers = 0; + + if (!F) { return 0; } + + for (auto *U : F->users()) { + + if (auto *CI = dyn_cast(U)) { ++callers; } + + } + + return callers; + +} + +/// return the calling function of a function - only if there is a single caller +Function *returnOnlyCaller(Function *F) { + + Function *caller = NULL; + + if (!F) { return NULL; } + + for (auto *U : F->users()) { + + if (auto *CI = dyn_cast(U)) { + + if (caller == NULL) { + + caller = CI->getParent()->getParent(); + + } else { + + return NULL; + + } + + } + + } + + return caller; + +} + void ModuleSanitizerCoverageLTO::instrumentFunction( Function &F, DomTreeCallback DTCallback, PostDomTreeCallback PDTCallback) { @@ -1287,6 +1346,37 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( // AFL++ START if (!F.size()) return; + + LLVMContext &Context = F.getContext(); + MDNode *N = MDNode::get(Context, MDString::get(Context, "nosanitize")); + + if (instrument_ctx) { + + // we have to set __afl_ctx 0 for all indirect calls in all functions, even + // those not to be instrumented. + for (auto &BB : F) { + + for (auto &IN : BB) { + + if (auto *Call = dyn_cast(&IN)) { + + if (Call->isIndirectCall()) { + + IRBuilder<> Builder(IN.getContext()); + Builder.SetInsertPoint(IN.getParent(), IN.getIterator()); + StoreInst *StoreCtx = Builder.CreateStore(Zero32, AFLContext); + StoreCtx->setMetadata("nosanitize", N); + + } + + } + + } + + } + + } + if (!isInInstrumentList(&F, FMNAME)) return; // AFL++ END @@ -1300,207 +1390,238 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( const PostDominatorTree *PDT = PDTCallback(F); bool IsLeafFunc = true; uint32_t skip_next = 0; - uint32_t call_counter = 0; + uint32_t call_counter = 0, call_depth = 0; uint32_t inst_save = inst, save_global = afl_global_id; uint32_t inst_in_this_func = 0; - LLVMContext &Context = F.getContext(); + Function *caller = NULL; - MDNode *N = MDNode::get(Context, MDString::get(Context, "nosanitize")); CTX_add = NULL; - fprintf(stderr, "Function: %s\n", F.getName().str().c_str()); + if (debug) fprintf(stderr, "Function: %s\n", F.getName().str().c_str()); - // Fake instrumentation so we can count how much there will be in this - // function + if (instrument_ctx) { - for (auto &BB : F) { + caller = &F; + call_counter = countCallers(caller); + Function *callee = caller; - for (auto &IN : BB) { + if (call_counter == 1) { - CallInst *callInst = nullptr; + ++call_depth; - if ((callInst = dyn_cast(&IN))) { + while (((caller = returnOnlyCaller(callee)) || 1 == 1) && + (call_counter = countCallers(callee)) == 1) { - Function *Callee = callInst->getCalledFunction(); - if (!Callee) continue; - if (callInst->getCallingConv() != llvm::CallingConv::C) continue; - StringRef FuncName = Callee->getName(); - - if (FuncName.compare(StringRef("__afl_coverage_interesting"))) continue; - - ++inst; + if (debug && caller && callee) + fprintf(stderr, "DEBUG: another depth: %s <- %s [%u]\n", + callee->getName().str().c_str(), + caller->getName().str().c_str(), call_depth); + ++call_depth; + callee = caller; } - SelectInst *selectInst = nullptr; + if (!caller && callee) { - if ((selectInst = dyn_cast(&IN))) { + caller = callee; + if (debug) + fprintf(stderr, "DEBUG: depth found: %s <- %s [count=%u, depth=%u]\n", + caller->getName().str().c_str(), F.getName().str().c_str(), + call_counter, call_depth); - Value *condition = selectInst->getCondition(); - auto t = condition->getType(); + } - if (t->getTypeID() == llvm::Type::IntegerTyID) { + } - inst += 2; + if (debug && call_counter < 2) { - } else + fprintf(stderr, "Function %s only %u (%s)\n", F.getName().str().c_str(), + call_counter, caller->getName().str().c_str()); -#if LLVM_VERSION_MAJOR >= 14 - if (t->getTypeID() == llvm::Type::FixedVectorTyID) { + } - FixedVectorType *tt = dyn_cast(t); - if (tt) { + if (call_counter == 1) { - uint32_t elements = tt->getElementCount().getFixedValue(); - inst += elements * 2; + call_counter = 0; + caller = NULL; + + } + + if (debug) { + + fprintf(stderr, "DEBUG: result: Function=%s callers=%u depth=%u\n", + F.getName().str().c_str(), call_counter, call_depth); + + } + + if (call_counter > 1) { + + // Fake instrumentation so we can count how many instrumentations there + // will be in this function + for (auto &BB : F) { + + for (auto &IN : BB) { + + CallInst *callInst = nullptr; + + if ((callInst = dyn_cast(&IN))) { + + Function *Callee = callInst->getCalledFunction(); + if (!Callee) continue; + if (callInst->getCallingConv() != llvm::CallingConv::C) continue; + StringRef FuncName = Callee->getName(); + + if (FuncName.compare(StringRef("__afl_coverage_interesting"))) + continue; + + ++inst; } - } else + SelectInst *selectInst = nullptr; + + if ((selectInst = dyn_cast(&IN))) { + + Value *condition = selectInst->getCondition(); + auto t = condition->getType(); + + if (t->getTypeID() == llvm::Type::IntegerTyID) { + + inst += 2; + + } else + +#if LLVM_VERSION_MAJOR >= 14 + if (t->getTypeID() == llvm::Type::FixedVectorTyID) { + + FixedVectorType *tt = dyn_cast(t); + if (tt) { + + uint32_t elements = tt->getElementCount().getFixedValue(); + inst += elements * 2; + + } + + } else #endif - { + { - continue; + continue; - } - - } - - } - - if (shouldInstrumentBlock(F, &BB, DT, PDT, Options)) - BlocksToInstrument.push_back(&BB); - /* - if (Options.IndirectCalls) - for (auto &Inst : BB) { - - fprintf(stderr, "FUCK TODO\n"); - CallBase *CB = dyn_cast(&Inst); - if (CB && !CB->getCalledFunction()) IndirCalls.push_back(&Inst); - /// XXX TODO!!!! set 0 ! + } } - */ - - } - - Fake_InjectCoverage(F, BlocksToInstrument, IsLeafFunc); - - // CTX init - - for (auto &BB : F) { - - if (instrument_ctx && &BB == &F.getEntryBlock()) { - - // we insert a CTX value in all our callers: - IRBuilder<> Builder(Context); - for (auto *U : F.users()) { - - if (auto *CI = dyn_cast(U)) { - - fprintf(stderr, "Insert %s [%u] -> %s\n", - CI->getParent()->getParent()->getName().str().c_str(), - call_counter, F.getName().str().c_str()); - Builder.SetInsertPoint(CI); - StoreInst *StoreCtx = Builder.CreateStore( - ConstantInt::get(Type::getInt32Ty(Context), call_counter++), - AFLContext); - StoreCtx->setMetadata("nosanitize", N); - } + if (shouldInstrumentBlock(F, &BB, DT, PDT, Options)) + BlocksToInstrument.push_back(&BB); + } - // We read the CTX for this call - Value *CTX_offset; - BasicBlock::iterator IP = BB.getFirstInsertionPt(); - IRBuilder<> IRB(&(*IP)); - LoadInst *PrevCtxLoad = IRB.CreateLoad( + Fake_InjectCoverage(F, BlocksToInstrument, IsLeafFunc); + + if (debug) + fprintf(stderr, "DEBUG: CTX: %u instrumentations\n", inst - inst_save); + + // we only instrument functions that have more than one instrumented block + if (inst > inst_save + 1) { + + inst_in_this_func = inst - inst_save; + bool done = false; + + // in rare occasions there can be multiple entry points per function + for (auto &BB : F) { + + if (&BB == &F.getEntryBlock() && done == false) { + + // we insert a CTX value in all our callers: + IRBuilder<> Builder(Context); + CallInst *CI = NULL; + Function *F2 = NULL; + uint32_t instrumented_calls = 0; + + for (auto *U : caller->users()) { + + if ((CI = dyn_cast(U))) { + + F2 = CI->getParent()->getParent(); + if (debug) + fprintf(stderr, + "DEBUG: CTX call insert %s [%u/%u] -> %s/%s\n", + F2->getName().str().c_str(), instrumented_calls + 1, + call_counter, caller->getName().str().c_str(), + F.getName().str().c_str()); + + Builder.SetInsertPoint(CI); + StoreInst *StoreCtx = Builder.CreateStore( + ConstantInt::get(Type::getInt32Ty(Context), + instrumented_calls++), + AFLContext); + StoreCtx->setMetadata("nosanitize", N); + + } + + } + + if (instrumented_calls != call_counter) { + + fprintf(stderr, "BUG! %s/%s <=> %u vs %u\n", + caller->getName().str().c_str(), + F.getName().str().c_str(), instrumented_calls, + call_counter); + exit(-1); + + } + + done = true; + + } + + // in all entrypoints we have to load the CTX value + if (&BB == &F.getEntryBlock()) { + + Value *CTX_offset; + BasicBlock::iterator IP = BB.getFirstInsertionPt(); + IRBuilder<> IRB(&(*IP)); + LoadInst *PrevCtxLoad = IRB.CreateLoad( #if LLVM_VERSION_MAJOR >= 14 - IRB.getInt32Ty(), + IRB.getInt32Ty(), #endif - AFLContext); + AFLContext); + PrevCtxLoad->setMetadata("nosanitize", N); - PrevCtxLoad->setMetadata("nosanitize", N); + CTX_offset = IRB.CreateMul( + ConstantInt::get(Type::getInt32Ty(Context), inst_in_this_func), + PrevCtxLoad, "CTXmul", false, true); - if (inst == inst_save || call_counter < 2) { + CTX_add = + IRB.CreateAlloca(Type::getInt32Ty(Context), nullptr, "CTX_add"); + auto nosan = IRB.CreateStore(CTX_offset, CTX_add); + nosan->setMetadata("nosanitize", N); - fprintf(stderr, "%s: ZERO!\n", F.getName().str().c_str()); - CTX_offset = Zero32; + if (debug) + fprintf( + stderr, "DEBUG: extra CTX instrumentations for %s: %u * %u\n", + F.getName().str().c_str(), inst - inst_save, call_counter); - } else { - - fprintf(stderr, "%s: %u * %u\n", F.getName().str().c_str(), - inst - inst_save, call_counter); - CTX_offset = IRB.CreateMul( - ConstantInt::get(Type::getInt32Ty(Context), inst - inst_save), - PrevCtxLoad, "CTXmul", false, true); - - } - - CTX_add = IRB.CreateAlloca(Type::getInt32Ty(Context), nullptr, "CTX_add"); - auto nosan = IRB.CreateStore(CTX_offset, CTX_add); - nosan->setMetadata("nosanitize", N); - - } - - // we have to set __afl_ctx 0 for all indirect calls - for (auto &IN : BB) { - - if (auto *Call = dyn_cast(&IN)) { - - if (Call->isIndirectCall()) { - - IRBuilder<> Builder(IN.getContext()); - Builder.SetInsertPoint(IN.getParent(), IN.getIterator()); - StoreInst *StoreCtx = Builder.CreateStore(Zero32, AFLContext); - StoreCtx->setMetadata("nosanitize", N); + } } + inst = inst_save; + } } } - inst_in_this_func = inst - inst_save; - inst = inst_save; - - // Now the real instrumentation - - if (CTX_add == NULL) { - - auto BB = &F.getEntryBlock(); - if (!BB) { - - fprintf(stderr, "NULL entry %s %p\n", F.getName().str().c_str(), BB); - exit(-1); - - } - - BasicBlock::iterator IP = BB->getFirstInsertionPt(); - IRBuilder<> IRB(&(*IP)); - CTX_add = IRB.CreateAlloca(Type::getInt32Ty(Context), nullptr, "CTX_add"); - auto nosan = IRB.CreateStore(Zero32, CTX_add); - nosan->setMetadata("nosanitize", N); - - } - for (auto &BB : F) { - // fprintf(stderr, "BB: %s\n", BB.getName().str().c_str()); - for (auto &IN : BB) { - /* - std::string Str; - llvm::raw_string_ostream RSO(Str); - IN.print(RSO); - errs() << RSO.str() << "\n"; - */ CallInst *callInst = nullptr; if ((callInst = dyn_cast(&IN))) { @@ -1523,16 +1644,21 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( if (FuncName.compare(StringRef("__afl_coverage_interesting"))) continue; - IRBuilder<> Builder(Context); - Value *val = ConstantInt::get(Int32Ty, ++afl_global_id); - LoadInst *CTX_load = Builder.CreateLoad( + Value *val = ConstantInt::get(Int32Ty, ++afl_global_id); + if (CTX_add) { + + IRBuilder<> Builder(Context); + LoadInst *CTX_load = Builder.CreateLoad( #if LLVM_VERSION_MAJOR >= 14 - Builder.getInt32Ty(), + Builder.getInt32Ty(), #endif - CTX_add); - ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(CTX_load); - Value *val2 = Builder.CreateAdd(val, CTX_load); - callInst->setOperand(1, val2); + CTX_add); + ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(CTX_load); + val = Builder.CreateAdd(val, CTX_load); + + } + + callInst->setOperand(1, val); ++inst; } @@ -1545,164 +1671,191 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( IN.print(os); fprintf(stderr, "X(%u): %s\n", skip_next, os.str().c_str()); */ - if ((selectInst = dyn_cast(&IN))) { + if (!skip_next && (selectInst = dyn_cast(&IN))) { - if (!skip_next) { + uint32_t vector_cnt = 0; + Value *condition = selectInst->getCondition(); + Value *result; + auto t = condition->getType(); + IRBuilder<> IRB(selectInst->getNextNode()); - uint32_t vector_cnt = 0; - Value *condition = selectInst->getCondition(); - Value *result; - auto t = condition->getType(); - IRBuilder<> IRB(selectInst->getNextNode()); + ++select_cnt; - // fprintf(stderr, "ADDING!!!\n"); - LoadInst *CTX_load = IRB.CreateLoad( + if (t->getTypeID() == llvm::Type::IntegerTyID) { + + Value *val1 = ConstantInt::get(Int32Ty, ++afl_global_id); + Value *val2 = ConstantInt::get(Int32Ty, ++afl_global_id); + if (CTX_add) { + + LoadInst *CTX_load = IRB.CreateLoad( #if LLVM_VERSION_MAJOR >= 14 - IRB.getInt32Ty(), + IRB.getInt32Ty(), #endif - CTX_add); - ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(CTX_load); - - ++select_cnt; - - if (t->getTypeID() == llvm::Type::IntegerTyID) { - - Value *val1 = ConstantInt::get(Int32Ty, ++afl_global_id); - Value *val11 = IRB.CreateAdd(val1, CTX_load); - Value *val2 = ConstantInt::get(Int32Ty, ++afl_global_id); - Value *val22 = IRB.CreateAdd(val2, CTX_load); - result = IRB.CreateSelect(condition, val11, val22); - skip_next = 1; - inst += 2; - - } else - -#if LLVM_VERSION_MAJOR >= 14 - if (t->getTypeID() == llvm::Type::FixedVectorTyID) { - - FixedVectorType *tt = dyn_cast(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 *val11 = IRB.CreateAdd(val1, CTX_load); - Value *val2 = ConstantInt::get(Int32Ty, ++afl_global_id); - Value *val22 = IRB.CreateAdd(val2, CTX_load); - x = IRB.CreateInsertElement(GuardPtr1, val11, (uint64_t)0); - y = IRB.CreateInsertElement(GuardPtr2, val22, (uint64_t)0); - - for (uint64_t i = 1; i < elements; i++) { - - val1 = ConstantInt::get(Int32Ty, ++afl_global_id); - val11 = IRB.CreateAdd(val1, CTX_load); - val2 = ConstantInt::get(Int32Ty, ++afl_global_id); - val11 = IRB.CreateAdd(val1, CTX_load); - x = IRB.CreateInsertElement(GuardPtr1, val11, i); - y = IRB.CreateInsertElement(GuardPtr2, val22, i); - - } - - result = IRB.CreateSelect(condition, x, y); - skip_next = 1; - - } - - } - - } else - -#endif - { - - unhandled++; - continue; - - } - - uint32_t vector_cur = 0; - /* Load SHM pointer */ - LoadInst *MapPtr = - IRB.CreateLoad(PointerType::get(Int8Ty, 0), AFLMapPtr); - ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(MapPtr); - - while (1) { - - /* Get CurLoc */ - Value *MapPtrIdx = nullptr; - - /* Load counter for CurLoc */ - if (!vector_cnt) { - - MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, result); - - } else { - - auto element = IRB.CreateExtractElement(result, vector_cur++); - MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, element); - - } - - if (use_threadsafe_counters) { - - auto nosan = IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, - MapPtrIdx, One, -#if LLVM_VERSION_MAJOR >= 13 - llvm::MaybeAlign(1), -#endif - llvm::AtomicOrdering::Monotonic); - ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(nosan); - - } else { - - LoadInst *Counter = IRB.CreateLoad(IRB.getInt8Ty(), MapPtrIdx); - ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(Counter); - - /* 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); - - } - - auto nosan = IRB.CreateStore(Incr, MapPtrIdx); - ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(nosan); - - } - - if (!vector_cnt || vector_cnt == vector_cur) { break; } + CTX_add); + val1 = IRB.CreateAdd(val1, CTX_load); + val2 = IRB.CreateAdd(val2, CTX_load); } + result = IRB.CreateSelect(condition, val1, val2); skip_next = 1; + inst += 2; - } else { + } else - skip_next = 0; +#if LLVM_VERSION_MAJOR >= 14 + if (t->getTypeID() == llvm::Type::FixedVectorTyID) { + + FixedVectorType *tt = dyn_cast(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); + if (CTX_add) { + + LoadInst *CTX_load = IRB.CreateLoad( + #if LLVM_VERSION_MAJOR >= 14 + IRB.getInt32Ty(), + #endif + CTX_add); + val1 = IRB.CreateAdd(val1, CTX_load); + val2 = IRB.CreateAdd(val2, CTX_load); + + } + + 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); + /*if (CTX_add) { // already loaded I guess + + LoadInst *CTX_load = IRB.CreateLoad( + #if LLVM_VERSION_MAJOR >= 14 + IRB.getInt32Ty(), + #endif + CTX_add); + val1 = IRB.CreateAdd(val1, CTX_load); + val2 = IRB.CreateAdd(val2, CTX_load); + + }*/ + + x = IRB.CreateInsertElement(GuardPtr1, val1, i); + y = IRB.CreateInsertElement(GuardPtr2, val2, i); + + } + + result = IRB.CreateSelect(condition, x, y); + skip_next = 1; + + } + + } + + } else + +#endif + { + + ++unhandled; + continue; } + uint32_t vector_cur = 0; + /* Load SHM pointer */ + LoadInst *MapPtr = + IRB.CreateLoad(PointerType::get(Int8Ty, 0), AFLMapPtr); + ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(MapPtr); + + while (1) { + + /* Get CurLoc */ + Value *MapPtrIdx = nullptr; + + /* Load counter for CurLoc */ + if (!vector_cnt) { + + MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, result); + + } else { + + auto element = IRB.CreateExtractElement(result, vector_cur++); + MapPtrIdx = IRB.CreateGEP(Int8Ty, 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(IRB.getInt8Ty(), MapPtrIdx); + ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(Counter); + + /* 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); + + } + + auto nosan = IRB.CreateStore(Incr, MapPtrIdx); + ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(nosan); + + } + + if (!vector_cnt || vector_cnt == vector_cur) { break; } + + } + + skip_next = 1; + + } else { + + skip_next = 0; + } } - // if (shouldInstrumentBlock(F, &BB, DT, PDT, Options)) - // BlocksToInstrument.push_back(&BB); + // if (shouldInstrumentBlock(F, &BB, DT, PDT, Options)) + // BlocksToInstrument.push_back(&BB); + + /* + for (auto &Inst : BB) { + + if (Options.IndirectCalls) { + + CallBase *CB = dyn_cast(&Inst); + if (CB && !CB->getCalledFunction()) IndirCalls.push_back(&Inst); + + } + + }*/ } @@ -1717,6 +1870,7 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( stderr, "BUG! inst_in_this_func %u != afl_global_id %u - save_global %u\n", inst_in_this_func, afl_global_id, save_global); + exit(-1); } @@ -1725,24 +1879,6 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( } - /* - fprintf(stderr, "FUNCTION: %s [%u]\n", F.getName().str().c_str(), - extra_ctx_inst); int n = 0; for (auto &BB : F) { - - fprintf(stderr, "BB %d\n", n++); - for (auto &IN : BB) { - - std::string Str; - llvm::raw_string_ostream RSO(Str); - IN.print(RSO); - errs() << RSO.str() << "\n"; - - } - - } - - */ - } GlobalVariable *ModuleSanitizerCoverageLTO::CreateFunctionLocalArrayInSection( @@ -1867,7 +2003,6 @@ bool ModuleSanitizerCoverageLTO::InjectCoverage( } -// Fake InjectCoverage to count the instrumentations bool ModuleSanitizerCoverageLTO::Fake_InjectCoverage( Function &F, ArrayRef AllBlocks, bool IsLeafFunc) { @@ -1875,7 +2010,6 @@ bool ModuleSanitizerCoverageLTO::Fake_InjectCoverage( for (size_t i = 0, N = AllBlocks.size(); i < N; i++) { - // AFL++ START if (BlockList.size()) { int skip = 0; @@ -1904,6 +2038,7 @@ bool ModuleSanitizerCoverageLTO::Fake_InjectCoverage( // The cache is used to speed up recording the caller-callee pairs. // The address of the caller is passed implicitly via caller PC. // CacheSize is encoded in the name of the run-time function. +/* void ModuleSanitizerCoverageLTO::InjectCoverageForIndirectCalls( Function &F, ArrayRef IndirCalls) { @@ -1922,6 +2057,8 @@ void ModuleSanitizerCoverageLTO::InjectCoverageForIndirectCalls( } +*/ + void ModuleSanitizerCoverageLTO::InjectCoverageAtBlock(Function &F, BasicBlock &BB, size_t Idx, @@ -1968,13 +2105,19 @@ void ModuleSanitizerCoverageLTO::InjectCoverageAtBlock(Function &F, /* Set the ID of the inserted basic block */ ConstantInt *CurLoc = ConstantInt::get(Int32Tyi, afl_global_id); - LoadInst *CTX_load = IRB.CreateLoad( + Value *val = CurLoc; + + if (CTX_add) { + + LoadInst *CTX_load = IRB.CreateLoad( #if LLVM_VERSION_MAJOR >= 14 - IRB.getInt32Ty(), + IRB.getInt32Ty(), #endif - CTX_add); - ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(CTX_load); - Value *offset = IRB.CreateAdd(CurLoc, CTX_load); + CTX_add); + ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(CTX_load); + val = IRB.CreateAdd(CurLoc, CTX_load); + + } /* Load SHM pointer */ @@ -1982,13 +2125,13 @@ void ModuleSanitizerCoverageLTO::InjectCoverageAtBlock(Function &F, if (map_addr) { - MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtrFixed, offset); + MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtrFixed, val); } else { LoadInst *MapPtr = IRB.CreateLoad(PointerType::get(Int8Ty, 0), AFLMapPtr); ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(MapPtr); - MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, offset); + MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, val); } From bd13d32437ebf0c1f7304dc4c8f9797dc4cce7fb Mon Sep 17 00:00:00 2001 From: vanhauser-thc Date: Fri, 2 Feb 2024 09:54:24 +0100 Subject: [PATCH 8/9] final touches --- instrumentation/SanitizerCoverageLTO.so.cc | 325 +++++++++++++-------- src/afl-cc.c | 2 - 2 files changed, 198 insertions(+), 129 deletions(-) diff --git a/instrumentation/SanitizerCoverageLTO.so.cc b/instrumentation/SanitizerCoverageLTO.so.cc index f6d60099..31d26ca3 100644 --- a/instrumentation/SanitizerCoverageLTO.so.cc +++ b/instrumentation/SanitizerCoverageLTO.so.cc @@ -1394,10 +1394,14 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( uint32_t inst_save = inst, save_global = afl_global_id; uint32_t inst_in_this_func = 0; Function *caller = NULL; + LoadInst *PrevCtxLoad = NULL; CTX_add = NULL; - if (debug) fprintf(stderr, "Function: %s\n", F.getName().str().c_str()); + if (debug) + fprintf(stderr, + "Function: %s (%u %u) XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX\n", + F.getName().str().c_str(), inst, afl_global_id); if (instrument_ctx) { @@ -1585,7 +1589,8 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( Value *CTX_offset; BasicBlock::iterator IP = BB.getFirstInsertionPt(); IRBuilder<> IRB(&(*IP)); - LoadInst *PrevCtxLoad = IRB.CreateLoad( + + PrevCtxLoad = IRB.CreateLoad( #if LLVM_VERSION_MAJOR >= 14 IRB.getInt32Ty(), #endif @@ -1608,20 +1613,78 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( } - } + // bool loaded = false, multicall = false; + for (auto &IN : BB) { - inst = inst_save; + // check all calls and where callee count == 1 instrument + // our current caller_id to __afl_ctx + if (auto callInst = dyn_cast(&IN)) { + + Function *Callee = callInst->getCalledFunction(); + if (countCallers(Callee) == 1) { + + if (debug) + fprintf(stderr, "DEBUG: %s call to %s with only one caller\n", + F.getName().str().c_str(), + Callee->getName().str().c_str()); + /* if (loaded == false || multicall == true) { // } */ + IRBuilder<> Builder(IN.getContext()); + Builder.SetInsertPoint(callInst); + StoreInst *StoreCtx = + Builder.CreateStore(PrevCtxLoad, AFLContext); + StoreCtx->setMetadata("nosanitize", N); + // multicall = false; loaded = true; + + } // else { multicall = true; } + + } + + } + + } } } + inst = inst_save; + } + /* if (debug) + fprintf(stderr, "Next instrumentation (%u-%u=%u %u-%u=%u)\n", inst, + inst_save, inst - inst_save, afl_global_id, save_global, + afl_global_id - save_global);*/ + for (auto &BB : F) { + skip_next = 0; + + /* + uint32_t j = 0; + fprintf(stderr, "BB %p ============================================\n", + CTX_add);*/ + for (auto &IN : BB) { + /* j++; + uint32_t i = 1; + std::string errMsg; + raw_string_ostream os(errMsg); + IN.print(os); + fprintf(stderr, "Next instruction, BB size now %zu: %02u %s\n", + BB.size(), j, os.str().c_str()); for (auto &IN2 : BB) { + + std::string errMsg2; + raw_string_ostream os2(errMsg2); + IN2.print(os2); + fprintf( + stderr, "%s %02u: %s\n", + strcmp(os.str().c_str(), os2.str().c_str()) == 0 ? ">>>" : " + ", i++, os2.str().c_str()); + + }*/ + CallInst *callInst = nullptr; if ((callInst = dyn_cast(&IN))) { @@ -1665,83 +1728,62 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( SelectInst *selectInst = nullptr; - /* - std::string errMsg; - raw_string_ostream os(errMsg); - IN.print(os); - fprintf(stderr, "X(%u): %s\n", skip_next, os.str().c_str()); - */ - if (!skip_next && (selectInst = dyn_cast(&IN))) { + if ((selectInst = dyn_cast(&IN))) { - uint32_t vector_cnt = 0; - Value *condition = selectInst->getCondition(); - Value *result; - auto t = condition->getType(); - IRBuilder<> IRB(selectInst->getNextNode()); + if (!skip_next) { - ++select_cnt; + // fprintf(stderr, "Select in\n"); - if (t->getTypeID() == llvm::Type::IntegerTyID) { + uint32_t vector_cnt = 0; + Value *condition = selectInst->getCondition(); + Value *result; + auto t = condition->getType(); + IRBuilder<> IRB(selectInst->getNextNode()); - Value *val1 = ConstantInt::get(Int32Ty, ++afl_global_id); - Value *val2 = ConstantInt::get(Int32Ty, ++afl_global_id); - if (CTX_add) { + ++select_cnt; - LoadInst *CTX_load = IRB.CreateLoad( + if (t->getTypeID() == llvm::Type::IntegerTyID) { + + Value *val1 = ConstantInt::get(Int32Ty, ++afl_global_id); + Value *val2 = ConstantInt::get(Int32Ty, ++afl_global_id); + if (CTX_add) { + + LoadInst *CTX_load = IRB.CreateLoad( #if LLVM_VERSION_MAJOR >= 14 - IRB.getInt32Ty(), + IRB.getInt32Ty(), #endif - CTX_add); - val1 = IRB.CreateAdd(val1, CTX_load); - val2 = IRB.CreateAdd(val2, CTX_load); + CTX_add); + val1 = IRB.CreateAdd(val1, CTX_load); + val2 = IRB.CreateAdd(val2, CTX_load); - } + } - result = IRB.CreateSelect(condition, val1, val2); - skip_next = 1; - inst += 2; + result = IRB.CreateSelect(condition, val1, val2); + skip_next = 1; + inst += 2; - } else + } else #if LLVM_VERSION_MAJOR >= 14 - if (t->getTypeID() == llvm::Type::FixedVectorTyID) { + if (t->getTypeID() == llvm::Type::FixedVectorTyID) { - FixedVectorType *tt = dyn_cast(t); - if (tt) { + FixedVectorType *tt = dyn_cast(t); + if (tt) { - uint32_t elements = tt->getElementCount().getFixedValue(); - vector_cnt = elements; - inst += vector_cnt * 2; - if (elements) { + 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; + 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); - if (CTX_add) { - - LoadInst *CTX_load = IRB.CreateLoad( - #if LLVM_VERSION_MAJOR >= 14 - IRB.getInt32Ty(), - #endif - CTX_add); - val1 = IRB.CreateAdd(val1, CTX_load); - val2 = IRB.CreateAdd(val2, CTX_load); - - } - - 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); - /*if (CTX_add) { // already loaded I guess + Value *val1 = ConstantInt::get(Int32Ty, ++afl_global_id); + Value *val2 = ConstantInt::get(Int32Ty, ++afl_global_id); + if (CTX_add) { LoadInst *CTX_load = IRB.CreateLoad( #if LLVM_VERSION_MAJOR >= 14 @@ -1751,93 +1793,117 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( val1 = IRB.CreateAdd(val1, CTX_load); val2 = IRB.CreateAdd(val2, CTX_load); - }*/ + } - x = IRB.CreateInsertElement(GuardPtr1, val1, i); - y = IRB.CreateInsertElement(GuardPtr2, val2, i); + 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); + /*if (CTX_add) { // already loaded I guess + + LoadInst *CTX_load = IRB.CreateLoad( + #if LLVM_VERSION_MAJOR >= 14 + IRB.getInt32Ty(), + #endif + CTX_add); + val1 = IRB.CreateAdd(val1, CTX_load); + val2 = IRB.CreateAdd(val2, CTX_load); + + }*/ + + x = IRB.CreateInsertElement(GuardPtr1, val1, i); + y = IRB.CreateInsertElement(GuardPtr2, val2, i); + + } + + result = IRB.CreateSelect(condition, x, y); + skip_next = 1; } - result = IRB.CreateSelect(condition, x, y); - skip_next = 1; + } + + } else + +#endif + { + + ++unhandled; + continue; + + } + + uint32_t vector_cur = 0; + /* Load SHM pointer */ + LoadInst *MapPtr = + IRB.CreateLoad(PointerType::get(Int8Ty, 0), AFLMapPtr); + ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(MapPtr); + + while (1) { + + /* Get CurLoc */ + Value *MapPtrIdx = nullptr; + + /* Load counter for CurLoc */ + if (!vector_cnt) { + + MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, result); + + } else { + + auto element = IRB.CreateExtractElement(result, vector_cur++); + MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, element); } - } + if (use_threadsafe_counters) { - } else - -#endif - { - - ++unhandled; - continue; - - } - - uint32_t vector_cur = 0; - /* Load SHM pointer */ - LoadInst *MapPtr = - IRB.CreateLoad(PointerType::get(Int8Ty, 0), AFLMapPtr); - ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(MapPtr); - - while (1) { - - /* Get CurLoc */ - Value *MapPtrIdx = nullptr; - - /* Load counter for CurLoc */ - if (!vector_cnt) { - - MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, result); - - } else { - - auto element = IRB.CreateExtractElement(result, vector_cur++); - MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, element); - - } - - if (use_threadsafe_counters) { - - IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One, + IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, + One, #if LLVM_VERSION_MAJOR >= 13 - llvm::MaybeAlign(1), + llvm::MaybeAlign(1), #endif - llvm::AtomicOrdering::Monotonic); + llvm::AtomicOrdering::Monotonic); - } else { + } else { - LoadInst *Counter = IRB.CreateLoad(IRB.getInt8Ty(), MapPtrIdx); - ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(Counter); + LoadInst *Counter = IRB.CreateLoad(IRB.getInt8Ty(), MapPtrIdx); + ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(Counter); - /* Update bitmap */ + /* Update bitmap */ - Value *Incr = IRB.CreateAdd(Counter, One); + Value *Incr = IRB.CreateAdd(Counter, One); - if (skip_nozero == NULL) { + if (skip_nozero == NULL) { - auto cf = IRB.CreateICmpEQ(Incr, Zero); - auto carry = IRB.CreateZExt(cf, Int8Ty); - Incr = IRB.CreateAdd(Incr, carry); + auto cf = IRB.CreateICmpEQ(Incr, Zero); + auto carry = IRB.CreateZExt(cf, Int8Ty); + Incr = IRB.CreateAdd(Incr, carry); + + } + + auto nosan = IRB.CreateStore(Incr, MapPtrIdx); + ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(nosan); } - auto nosan = IRB.CreateStore(Incr, MapPtrIdx); - ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(nosan); + if (!vector_cnt || vector_cnt == vector_cur) { break; } } - if (!vector_cnt || vector_cnt == vector_cur) { break; } + skip_next = 1; + // fprintf(stderr, "Select out\n"); + + } else { + + // fprintf(stderr, "Select skip\n"); + skip_next = 0; } - skip_next = 1; - - } else { - - skip_next = 0; - } } @@ -1862,6 +1928,11 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( InjectCoverage(F, BlocksToInstrument, IsLeafFunc); // InjectCoverageForIndirectCalls(F, IndirCalls); + /*if (debug) + fprintf(stderr, "Done instrumentation (%u-%u=%u %u-%u=%u)\n", inst, + inst_save, inst - inst_save, afl_global_id, save_global, + afl_global_id - save_global);*/ + if (inst_in_this_func && call_counter > 1) { if (inst_in_this_func != afl_global_id - save_global) { diff --git a/src/afl-cc.c b/src/afl-cc.c index 4f6745ed..fd466541 100644 --- a/src/afl-cc.c +++ b/src/afl-cc.c @@ -1103,8 +1103,6 @@ static void instrument_opt_mode_exclude(aflcc_state_t *aflcc) { } - fprintf(stderr, "X %u %u\n", aflcc->compiler_mode, LTO); - if (aflcc->instrument_opt_mode && aflcc->compiler_mode != LLVM && !((aflcc->instrument_opt_mode & INSTRUMENT_OPT_CALLER) && aflcc->compiler_mode == LTO)) From ba3a039e457025399f7b58905d3923d068ef0eef Mon Sep 17 00:00:00 2001 From: vanhauser-thc Date: Fri, 2 Feb 2024 13:17:20 +0100 Subject: [PATCH 9/9] finish lto-ctx --- instrumentation/SanitizerCoverageLTO.so.cc | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/instrumentation/SanitizerCoverageLTO.so.cc b/instrumentation/SanitizerCoverageLTO.so.cc index 31d26ca3..65602109 100644 --- a/instrumentation/SanitizerCoverageLTO.so.cc +++ b/instrumentation/SanitizerCoverageLTO.so.cc @@ -250,7 +250,7 @@ class ModuleSanitizerCoverageLTO uint32_t afl_global_id = 0; uint32_t unhandled = 0; uint32_t select_cnt = 0; - uint32_t instrument_ctx = 1; + uint32_t instrument_ctx = 0; uint32_t extra_ctx_inst = 0; uint64_t map_addr = 0; const char *skip_nozero = NULL; @@ -771,12 +771,12 @@ bool ModuleSanitizerCoverageLTO::instrumentModule( else Str2 = TmpStr.str(); - if (debug) + /*if (debug) fprintf(stderr, "F:%s %p(%s)->\"%s\"(%s) %p(%s)->\"%s\"(%s)\n", FuncName.c_str(), Str1P, Str1P->getName().str().c_str(), Str1.c_str(), HasStr1 == true ? "true" : "false", Str2P, Str2P->getName().str().c_str(), Str2.c_str(), - HasStr2 == true ? "true" : "false"); + HasStr2 == true ? "true" : "false");*/ // we handle the 2nd parameter first because of llvm memcpy if (!HasStr2) { @@ -1398,10 +1398,7 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( CTX_add = NULL; - if (debug) - fprintf(stderr, - "Function: %s (%u %u) XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX\n", - F.getName().str().c_str(), inst, afl_global_id); + if (debug) fprintf(stderr, "Function: %s\n", F.getName().str().c_str()); if (instrument_ctx) { @@ -1613,7 +1610,6 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( } - // bool loaded = false, multicall = false; for (auto &IN : BB) { // check all calls and where callee count == 1 instrument @@ -1627,15 +1623,14 @@ void ModuleSanitizerCoverageLTO::instrumentFunction( fprintf(stderr, "DEBUG: %s call to %s with only one caller\n", F.getName().str().c_str(), Callee->getName().str().c_str()); - /* if (loaded == false || multicall == true) { // } */ + IRBuilder<> Builder(IN.getContext()); Builder.SetInsertPoint(callInst); StoreInst *StoreCtx = Builder.CreateStore(PrevCtxLoad, AFLContext); StoreCtx->setMetadata("nosanitize", N); - // multicall = false; loaded = true; - } // else { multicall = true; } + } }