#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;
}
if (getenv("AFL_LLVM_INSTRIM_SKIPSINGLEBLOCK") ||
getenv("AFL_LLVM_SKIPSINGLEBLOCK"))
function_minimum_size = 2;
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;
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);
uint64_t PrevLocVecSize = PowerOf2Ceil(PrevLocSize);
IntegerType *IntLocTy =
IntegerType::getIntNTy(C, sizeof(PREV_LOC_T) * CHAR_BIT);
if (ngram_size) PrevLocTy = VectorType::get(IntLocTy, PrevLocVecSize);
#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;
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 (unsigned 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;
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);
