diff options
Diffstat (limited to 'instrumentation/LLVMInsTrim.so.cc')
| -rw-r--r-- | instrumentation/LLVMInsTrim.so.cc | 598 |
1 files changed, 598 insertions, 0 deletions
diff --git a/instrumentation/LLVMInsTrim.so.cc b/instrumentation/LLVMInsTrim.so.cc new file mode 100644 index 00000000..61a420ba --- /dev/null +++ b/instrumentation/LLVMInsTrim.so.cc @@ -0,0 +1,598 @@ +#include <stdio.h> +#include <stdlib.h> +#include <stdarg.h> +#include <unistd.h> + +#include "llvm/Config/llvm-config.h" +#if LLVM_VERSION_MAJOR == 3 && LLVM_VERSION_MINOR < 5 +typedef long double max_align_t; +#endif + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.h" +#if LLVM_VERSION_MAJOR > 3 || \ + (LLVM_VERSION_MAJOR == 3 && LLVM_VERSION_MINOR > 4) + #include "llvm/IR/CFG.h" + #include "llvm/IR/Dominators.h" + #include "llvm/IR/DebugInfo.h" +#else + #include "llvm/Support/CFG.h" + #include "llvm/Analysis/Dominators.h" + #include "llvm/DebugInfo.h" +#endif +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Instructions.h" +#include "llvm/IR/LegacyPassManager.h" +#include "llvm/IR/Module.h" +#include "llvm/Pass.h" +#include "llvm/Support/raw_ostream.h" +#include "llvm/Support/CommandLine.h" +#include "llvm/Transforms/IPO/PassManagerBuilder.h" +#include "llvm/Transforms/Utils/BasicBlockUtils.h" +#include "llvm/IR/BasicBlock.h" +#include <unordered_set> +#include <random> +#include <list> +#include <string> +#include <fstream> + +#include "MarkNodes.h" +#include "afl-llvm-common.h" +#include "llvm-ngram-coverage.h" + +#include "config.h" +#include "debug.h" + +using namespace llvm; + +static cl::opt<bool> MarkSetOpt("markset", cl::desc("MarkSet"), + cl::init(false)); +static cl::opt<bool> LoopHeadOpt("loophead", cl::desc("LoopHead"), + cl::init(false)); + +namespace { + +struct InsTrim : public ModulePass { + + protected: + uint32_t function_minimum_size = 1; + char * skip_nozero = NULL; + + private: + std::mt19937 generator; + int total_instr = 0; + + unsigned int genLabel() { + + return generator() & (MAP_SIZE - 1); + + } + + public: + static char ID; + + InsTrim() : ModulePass(ID), generator(0) { + + initInstrumentList(); + + } + + void getAnalysisUsage(AnalysisUsage &AU) const override { + + AU.addRequired<DominatorTreeWrapperPass>(); + + } + +#if LLVM_VERSION_MAJOR < 4 + const char * +#else + StringRef +#endif + getPassName() const override { + + return "InstTrim Instrumentation"; + + } + +#if LLVM_VERSION_MAJOR > 4 || \ + (LLVM_VERSION_MAJOR == 4 && LLVM_VERSION_PATCH >= 1) + #define AFL_HAVE_VECTOR_INTRINSICS 1 +#endif + + bool runOnModule(Module &M) override { + + setvbuf(stdout, NULL, _IONBF, 0); + + if ((isatty(2) && !getenv("AFL_QUIET")) || getenv("AFL_DEBUG") != NULL) { + + SAYF(cCYA "LLVMInsTrim" VERSION cRST " by csienslab\n"); + + } else + + be_quiet = 1; + + if (getenv("AFL_DEBUG") != NULL) debug = 1; + + LLVMContext &C = M.getContext(); + + IntegerType *Int8Ty = IntegerType::getInt8Ty(C); + IntegerType *Int32Ty = IntegerType::getInt32Ty(C); + +#if LLVM_VERSION_MAJOR < 9 + char *neverZero_counters_str; + if ((neverZero_counters_str = getenv("AFL_LLVM_NOT_ZERO")) != NULL) + if (!be_quiet) OKF("LLVM neverZero activated (by hexcoder)\n"); +#endif + skip_nozero = getenv("AFL_LLVM_SKIP_NEVERZERO"); + + if (getenv("AFL_LLVM_INSTRIM_LOOPHEAD") != NULL || + getenv("LOOPHEAD") != NULL) { + + LoopHeadOpt = true; + + } + + unsigned int PrevLocSize = 0; + char * ngram_size_str = getenv("AFL_LLVM_NGRAM_SIZE"); + if (!ngram_size_str) ngram_size_str = getenv("AFL_NGRAM_SIZE"); + char *ctx_str = getenv("AFL_LLVM_CTX"); + +#ifdef AFL_HAVE_VECTOR_INTRINSICS + unsigned int ngram_size = 0; + /* Decide previous location vector size (must be a power of two) */ + VectorType *PrevLocTy = NULL; + + if (ngram_size_str) + if (sscanf(ngram_size_str, "%u", &ngram_size) != 1 || ngram_size < 2 || + ngram_size > NGRAM_SIZE_MAX) + FATAL( + "Bad value of AFL_NGRAM_SIZE (must be between 2 and NGRAM_SIZE_MAX " + "(%u))", + NGRAM_SIZE_MAX); + + if (ngram_size) + PrevLocSize = ngram_size - 1; + else +#else + if (ngram_size_str) + #ifdef LLVM_VERSION_STRING + FATAL( + "Sorry, NGRAM branch coverage is not supported with llvm version %s!", + LLVM_VERSION_STRING); + #else + #ifndef LLVM_VERSION_PATCH + FATAL( + "Sorry, NGRAM branch coverage is not supported with llvm version " + "%d.%d.%d!", + LLVM_VERSION_MAJOR, LLVM_VERSION_MINOR, 0); + #else + FATAL( + "Sorry, NGRAM branch coverage is not supported with llvm version " + "%d.%d.%d!", + LLVM_VERSION_MAJOR, LLVM_VERSION_MINOR, LLVM_VERISON_PATCH); + #endif + #endif +#endif + PrevLocSize = 1; + +#ifdef AFL_HAVE_VECTOR_INTRINSICS + // IntegerType *Int64Ty = IntegerType::getInt64Ty(C); + int PrevLocVecSize = PowerOf2Ceil(PrevLocSize); + IntegerType *IntLocTy = + IntegerType::getIntNTy(C, sizeof(PREV_LOC_T) * CHAR_BIT); + if (ngram_size) + PrevLocTy = VectorType::get(IntLocTy, PrevLocVecSize + #if LLVM_VERSION_MAJOR >= 12 + , + false + #endif + ); +#endif + + /* Get globals for the SHM region and the previous location. Note that + __afl_prev_loc is thread-local. */ + + GlobalVariable *AFLMapPtr = + new GlobalVariable(M, PointerType::get(Int8Ty, 0), false, + GlobalValue::ExternalLinkage, 0, "__afl_area_ptr"); + GlobalVariable *AFLPrevLoc; + GlobalVariable *AFLContext = NULL; + LoadInst * PrevCtx = NULL; // for CTX sensitive coverage + + if (ctx_str) +#ifdef __ANDROID__ + AFLContext = new GlobalVariable( + M, Int32Ty, false, GlobalValue::ExternalLinkage, 0, "__afl_prev_ctx"); +#else + AFLContext = new GlobalVariable( + M, Int32Ty, false, GlobalValue::ExternalLinkage, 0, "__afl_prev_ctx", + 0, GlobalVariable::GeneralDynamicTLSModel, 0, false); +#endif + +#ifdef AFL_HAVE_VECTOR_INTRINSICS + if (ngram_size) + #ifdef __ANDROID__ + AFLPrevLoc = new GlobalVariable( + M, PrevLocTy, /* isConstant */ false, GlobalValue::ExternalLinkage, + /* Initializer */ nullptr, "__afl_prev_loc"); + #else + AFLPrevLoc = new GlobalVariable( + M, PrevLocTy, /* isConstant */ false, GlobalValue::ExternalLinkage, + /* Initializer */ nullptr, "__afl_prev_loc", + /* InsertBefore */ nullptr, GlobalVariable::GeneralDynamicTLSModel, + /* AddressSpace */ 0, /* IsExternallyInitialized */ false); + #endif + else +#endif +#ifdef __ANDROID__ + AFLPrevLoc = new GlobalVariable( + M, Int32Ty, false, GlobalValue::ExternalLinkage, 0, "__afl_prev_loc"); +#else + AFLPrevLoc = new GlobalVariable( + M, Int32Ty, false, GlobalValue::ExternalLinkage, 0, "__afl_prev_loc", 0, + GlobalVariable::GeneralDynamicTLSModel, 0, false); +#endif + +#ifdef AFL_HAVE_VECTOR_INTRINSICS + /* Create the vector shuffle mask for updating the previous block history. + Note that the first element of the vector will store cur_loc, so just set + it to undef to allow the optimizer to do its thing. */ + + SmallVector<Constant *, 32> PrevLocShuffle = {UndefValue::get(Int32Ty)}; + + for (unsigned I = 0; I < PrevLocSize - 1; ++I) + PrevLocShuffle.push_back(ConstantInt::get(Int32Ty, I)); + + for (int I = PrevLocSize; I < PrevLocVecSize; ++I) + PrevLocShuffle.push_back(ConstantInt::get(Int32Ty, PrevLocSize)); + + Constant *PrevLocShuffleMask = ConstantVector::get(PrevLocShuffle); +#endif + + // this is our default + MarkSetOpt = true; + + ConstantInt *Zero = ConstantInt::get(Int8Ty, 0); + ConstantInt *One = ConstantInt::get(Int8Ty, 1); + + u64 total_rs = 0; + u64 total_hs = 0; + + scanForDangerousFunctions(&M); + + for (Function &F : M) { + + if (debug) { + + uint32_t bb_cnt = 0; + + for (auto &BB : F) + if (BB.size() > 0) ++bb_cnt; + SAYF(cMGN "[D] " cRST "Function %s size %zu %u\n", + F.getName().str().c_str(), F.size(), bb_cnt); + + } + + if (!isInInstrumentList(&F)) continue; + + // if the function below our minimum size skip it (1 or 2) + if (F.size() < function_minimum_size) { continue; } + + std::unordered_set<BasicBlock *> MS; + if (!MarkSetOpt) { + + for (auto &BB : F) { + + MS.insert(&BB); + + } + + total_rs += F.size(); + + } else { + + auto Result = markNodes(&F); + auto RS = Result.first; + auto HS = Result.second; + + MS.insert(RS.begin(), RS.end()); + if (!LoopHeadOpt) { + + MS.insert(HS.begin(), HS.end()); + total_rs += MS.size(); + + } else { + + DenseSet<std::pair<BasicBlock *, BasicBlock *>> EdgeSet; + DominatorTreeWrapperPass * DTWP = + &getAnalysis<DominatorTreeWrapperPass>(F); + auto DT = &DTWP->getDomTree(); + + total_rs += RS.size(); + total_hs += HS.size(); + + for (BasicBlock *BB : HS) { + + bool Inserted = false; + for (auto BI = pred_begin(BB), BE = pred_end(BB); BI != BE; ++BI) { + + auto Edge = BasicBlockEdge(*BI, BB); + if (Edge.isSingleEdge() && DT->dominates(Edge, BB)) { + + EdgeSet.insert({*BI, BB}); + Inserted = true; + break; + + } + + } + + if (!Inserted) { + + MS.insert(BB); + total_rs += 1; + total_hs -= 1; + + } + + } + + for (auto I = EdgeSet.begin(), E = EdgeSet.end(); I != E; ++I) { + + auto PredBB = I->first; + auto SuccBB = I->second; + auto NewBB = + SplitBlockPredecessors(SuccBB, {PredBB}, ".split", DT, nullptr, +#if LLVM_VERSION_MAJOR >= 8 + nullptr, +#endif + false); + MS.insert(NewBB); + + } + + } + + for (BasicBlock &BB : F) { + + if (MS.find(&BB) == MS.end()) { continue; } + IRBuilder<> IRB(&*BB.getFirstInsertionPt()); + +#ifdef AFL_HAVE_VECTOR_INTRINSICS + if (ngram_size) { + + LoadInst *PrevLoc = IRB.CreateLoad(AFLPrevLoc); + PrevLoc->setMetadata(M.getMDKindID("nosanitize"), + MDNode::get(C, None)); + + Value *ShuffledPrevLoc = IRB.CreateShuffleVector( + PrevLoc, UndefValue::get(PrevLocTy), PrevLocShuffleMask); + Value *UpdatedPrevLoc = IRB.CreateInsertElement( + ShuffledPrevLoc, ConstantInt::get(Int32Ty, genLabel()), + (uint64_t)0); + + IRB.CreateStore(UpdatedPrevLoc, AFLPrevLoc) + ->setMetadata(M.getMDKindID("nosanitize"), + MDNode::get(C, None)); + + } else + +#endif + { + + IRB.CreateStore(ConstantInt::get(Int32Ty, genLabel()), AFLPrevLoc); + + } + + } + + } + + int has_calls = 0; + for (BasicBlock &BB : F) { + + auto PI = pred_begin(&BB); + auto PE = pred_end(&BB); + IRBuilder<> IRB(&*BB.getFirstInsertionPt()); + Value * L = NULL; + unsigned int cur_loc; + + // Context sensitive coverage + if (ctx_str && &BB == &F.getEntryBlock()) { + + PrevCtx = IRB.CreateLoad(AFLContext); + PrevCtx->setMetadata(M.getMDKindID("nosanitize"), + MDNode::get(C, None)); + + // does the function have calls? and is any of the calls larger than + // one basic block? + has_calls = 0; + for (auto &BB : F) { + + if (has_calls) break; + for (auto &IN : BB) { + + CallInst *callInst = nullptr; + if ((callInst = dyn_cast<CallInst>(&IN))) { + + Function *Callee = callInst->getCalledFunction(); + if (!Callee || Callee->size() < function_minimum_size) + continue; + else { + + has_calls = 1; + break; + + } + + } + + } + + } + + // if yes we store a context ID for this function in the global var + if (has_calls) { + + ConstantInt *NewCtx = ConstantInt::get(Int32Ty, genLabel()); + StoreInst * StoreCtx = IRB.CreateStore(NewCtx, AFLContext); + StoreCtx->setMetadata(M.getMDKindID("nosanitize"), + MDNode::get(C, None)); + + } + + } // END of ctx_str + + if (MarkSetOpt && MS.find(&BB) == MS.end()) { continue; } + + if (PI == PE) { + + cur_loc = genLabel(); + L = ConstantInt::get(Int32Ty, cur_loc); + + } else { + + auto *PN = PHINode::Create(Int32Ty, 0, "", &*BB.begin()); + DenseMap<BasicBlock *, unsigned> PredMap; + for (auto PI = pred_begin(&BB), PE = pred_end(&BB); PI != PE; ++PI) { + + BasicBlock *PBB = *PI; + auto It = PredMap.insert({PBB, genLabel()}); + unsigned Label = It.first->second; + cur_loc = Label; + PN->addIncoming(ConstantInt::get(Int32Ty, Label), PBB); + + } + + L = PN; + + } + + /* Load prev_loc */ + LoadInst *PrevLoc = IRB.CreateLoad(AFLPrevLoc); + PrevLoc->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + Value *PrevLocTrans; + +#ifdef AFL_HAVE_VECTOR_INTRINSICS + /* "For efficiency, we propose to hash the tuple as a key into the + hit_count map as (prev_block_trans << 1) ^ curr_block_trans, where + prev_block_trans = (block_trans_1 ^ ... ^ block_trans_(n-1)" */ + + if (ngram_size) + PrevLocTrans = + IRB.CreateZExt(IRB.CreateXorReduce(PrevLoc), IRB.getInt32Ty()); + else +#endif + PrevLocTrans = IRB.CreateZExt(PrevLoc, IRB.getInt32Ty()); + + if (ctx_str) + PrevLocTrans = + IRB.CreateZExt(IRB.CreateXor(PrevLocTrans, PrevCtx), Int32Ty); + + /* Load SHM pointer */ + LoadInst *MapPtr = IRB.CreateLoad(AFLMapPtr); + MapPtr->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + Value *MapPtrIdx; +#ifdef AFL_HAVE_VECTOR_INTRINSICS + if (ngram_size) + MapPtrIdx = IRB.CreateGEP( + MapPtr, IRB.CreateZExt(IRB.CreateXor(PrevLocTrans, L), Int32Ty)); + else +#endif + MapPtrIdx = IRB.CreateGEP(MapPtr, IRB.CreateXor(PrevLocTrans, L)); + + /* Update bitmap */ + LoadInst *Counter = IRB.CreateLoad(MapPtrIdx); + Counter->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + + 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) +#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 + */ + 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)); + + if (ctx_str && has_calls) { + + // in CTX mode we have to restore the original context for the + // caller - she might be calling other functions which need the + // correct CTX + Instruction *Inst = BB.getTerminator(); + if (isa<ReturnInst>(Inst) || isa<ResumeInst>(Inst)) { + + IRBuilder<> Post_IRB(Inst); + StoreInst * RestoreCtx = Post_IRB.CreateStore(PrevCtx, AFLContext); + RestoreCtx->setMetadata(M.getMDKindID("nosanitize"), + MDNode::get(C, None)); + + } + + } + + total_instr++; + + } + + } + + if (!be_quiet) { + + char modeline[100]; + snprintf(modeline, sizeof(modeline), "%s%s%s%s%s", + getenv("AFL_HARDEN") ? "hardened" : "non-hardened", + getenv("AFL_USE_ASAN") ? ", ASAN" : "", + getenv("AFL_USE_MSAN") ? ", MSAN" : "", + getenv("AFL_USE_CFISAN") ? ", CFISAN" : "", + getenv("AFL_USE_UBSAN") ? ", UBSAN" : ""); + + OKF("Instrumented %u locations (%llu, %llu) (%s mode)\n", total_instr, + total_rs, total_hs, modeline); + + } + + return false; + + } + +}; // end of struct InsTrim + +} // end of anonymous namespace + +char InsTrim::ID = 0; + +static void registerAFLPass(const PassManagerBuilder &, + legacy::PassManagerBase &PM) { + + PM.add(new InsTrim()); + +} + +static RegisterStandardPasses RegisterAFLPass( + PassManagerBuilder::EP_OptimizerLast, registerAFLPass); + +static RegisterStandardPasses RegisterAFLPass0( + PassManagerBuilder::EP_EnabledOnOptLevel0, registerAFLPass); + |