From 0aa93afeb82c610fd39b5dd4a3dd7482b9f86b1e Mon Sep 17 00:00:00 2001 From: Andrea Fioraldi Date: Thu, 4 Mar 2021 14:50:26 +0100 Subject: vectorial top-k CTX first implementation --- instrumentation/LLVMInsTrim.so.cc | 2 +- instrumentation/afl-compiler-rt.o.c | 4 +- instrumentation/afl-llvm-pass.so.cc | 154 ++++++++++++++++++++++++---- instrumentation/llvm-alternative-coverage.h | 21 ++++ instrumentation/llvm-ngram-coverage.h | 18 ---- src/afl-cc.c | 2 +- 6 files changed, 162 insertions(+), 39 deletions(-) create mode 100644 instrumentation/llvm-alternative-coverage.h delete mode 100644 instrumentation/llvm-ngram-coverage.h diff --git a/instrumentation/LLVMInsTrim.so.cc b/instrumentation/LLVMInsTrim.so.cc index 948f8f3a..83f11411 100644 --- a/instrumentation/LLVMInsTrim.so.cc +++ b/instrumentation/LLVMInsTrim.so.cc @@ -38,7 +38,7 @@ typedef long double max_align_t; #include "MarkNodes.h" #include "afl-llvm-common.h" -#include "llvm-ngram-coverage.h" +#include "llvm-alternative-coverage.h" #include "config.h" #include "debug.h" diff --git a/instrumentation/afl-compiler-rt.o.c b/instrumentation/afl-compiler-rt.o.c index c9577a55..d583e246 100644 --- a/instrumentation/afl-compiler-rt.o.c +++ b/instrumentation/afl-compiler-rt.o.c @@ -20,7 +20,7 @@ #include "config.h" #include "types.h" #include "cmplog.h" -#include "llvm-ngram-coverage.h" +#include "llvm-alternative-coverage.h" #include #include @@ -97,10 +97,12 @@ int __afl_selective_coverage_temp = 1; #if defined(__ANDROID__) || defined(__HAIKU__) PREV_LOC_T __afl_prev_loc[NGRAM_SIZE_MAX]; +PREV_LOC_T __afl_prev_caller[CTX_MAX_K]; u32 __afl_prev_ctx; u32 __afl_cmp_counter; #else __thread PREV_LOC_T __afl_prev_loc[NGRAM_SIZE_MAX]; +__thread PREV_LOC_T __afl_prev_caller[CTX_MAX_K]; __thread u32 __afl_prev_ctx; __thread u32 __afl_cmp_counter; #endif diff --git a/instrumentation/afl-llvm-pass.so.cc b/instrumentation/afl-llvm-pass.so.cc index 33898aec..c58e9d95 100644 --- a/instrumentation/afl-llvm-pass.so.cc +++ b/instrumentation/afl-llvm-pass.so.cc @@ -62,7 +62,7 @@ typedef long double max_align_t; #endif #include "afl-llvm-common.h" -#include "llvm-ngram-coverage.h" +#include "llvm-alternative-coverage.h" using namespace llvm; @@ -82,6 +82,7 @@ class AFLCoverage : public ModulePass { protected: uint32_t ngram_size = 0; + uint32_t ctx_k = 0; uint32_t map_size = MAP_SIZE; uint32_t function_minimum_size = 1; char * ctx_str = NULL, *caller_str = NULL, *skip_nozero = NULL; @@ -183,12 +184,17 @@ bool AFLCoverage::runOnModule(Module &M) { skip_nozero = getenv("AFL_LLVM_SKIP_NEVERZERO"); unsigned PrevLocSize = 0; + unsigned PrevCallerSize = 0; char *ngram_size_str = getenv("AFL_LLVM_NGRAM_SIZE"); if (!ngram_size_str) ngram_size_str = getenv("AFL_NGRAM_SIZE"); + char *ctx_k_str = getenv("AFL_LLVM_CTX_K"); + if (!ctx_k_str) ctx_k_str = getenv("AFL_CTX_K"); ctx_str = getenv("AFL_LLVM_CTX"); caller_str = getenv("AFL_LLVM_CALLER"); + bool instrument_ctx = ctx_str || caller_str; + #ifdef AFL_HAVE_VECTOR_INTRINSICS /* Decide previous location vector size (must be a power of two) */ VectorType *PrevLocTy = NULL; @@ -205,6 +211,25 @@ bool AFLCoverage::runOnModule(Module &M) { if (ngram_size) PrevLocSize = ngram_size - 1; else + PrevLocSize = 1; + + /* Decide K-ctx vector size (must be a power of two) */ + VectorType *PrevCallerTy = NULL; + + if (ctx_k_str) + if (sscanf(ctx_k_str, "%u", &ctx_k) != 1 || ctx_k < 2 || + ctx_k > CTX_MAX_K) + FATAL("Bad value of AFL_CTX_K (must be between 2 and CTX_MAX_K (%u))", CTX_MAX_K); + + if (ctx_k == 1) { + ctx_k = 0; + instrument_ctx = true; + caller_str = ctx_k_str; // Enable CALLER instead + } + if (ctx_k) { + PrevCallerSize = ctx_k; + instrument_ctx = true; + } #else if (ngram_size_str) #ifndef LLVM_VERSION_PATCH @@ -218,8 +243,20 @@ bool AFLCoverage::runOnModule(Module &M) { "%d.%d.%d!", LLVM_VERSION_MAJOR, LLVM_VERSION_MINOR, LLVM_VERSION_PATCH); #endif + if (ctx_k_str) + #ifndef LLVM_VERSION_PATCH + FATAL( + "Sorry, K-CTX branch coverage is not supported with llvm version " + "%d.%d.%d!", + LLVM_VERSION_MAJOR, LLVM_VERSION_MINOR, 0); + #else + FATAL( + "Sorry, K-CTX branch coverage is not supported with llvm version " + "%d.%d.%d!", + LLVM_VERSION_MAJOR, LLVM_VERSION_MINOR, LLVM_VERSION_PATCH); + #endif + PrevLocSize = 1; #endif - PrevLocSize = 1; #ifdef AFL_HAVE_VECTOR_INTRINSICS int PrevLocVecSize = PowerOf2Ceil(PrevLocSize); @@ -232,6 +269,17 @@ bool AFLCoverage::runOnModule(Module &M) { ); #endif +#ifdef AFL_HAVE_VECTOR_INTRINSICS + int PrevCallerVecSize = PowerOf2Ceil(PrevCallerSize); + if (ctx_k) + PrevCallerTy = VectorType::get(IntLocTy, PrevCallerVecSize + #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. */ @@ -239,6 +287,7 @@ bool AFLCoverage::runOnModule(Module &M) { new GlobalVariable(M, PointerType::get(Int8Ty, 0), false, GlobalValue::ExternalLinkage, 0, "__afl_area_ptr"); GlobalVariable *AFLPrevLoc; + GlobalVariable *AFLPrevCaller; GlobalVariable *AFLContext = NULL; if (ctx_str || caller_str) @@ -275,6 +324,30 @@ bool AFLCoverage::runOnModule(Module &M) { GlobalVariable::GeneralDynamicTLSModel, 0, false); #endif +#ifdef AFL_HAVE_VECTOR_INTRINSICS + if (ctx_k) + #if defined(__ANDROID__) || defined(__HAIKU__) + AFLPrevCaller = new GlobalVariable( + M, PrevCallerTy, /* isConstant */ false, GlobalValue::ExternalLinkage, + /* Initializer */ nullptr, "__afl_prev_caller"); + #else + AFLPrevCaller = new GlobalVariable( + M, PrevCallerTy, /* isConstant */ false, GlobalValue::ExternalLinkage, + /* Initializer */ nullptr, "__afl_prev_caller", + /* InsertBefore */ nullptr, GlobalVariable::GeneralDynamicTLSModel, + /* AddressSpace */ 0, /* IsExternallyInitialized */ false); + #endif + else +#endif +#if defined(__ANDROID__) || defined(__HAIKU__) + AFLPrevCaller = new GlobalVariable( + M, Int32Ty, false, GlobalValue::ExternalLinkage, 0, "__afl_prev_caller"); +#else + AFLPrevCaller = new GlobalVariable( + M, Int32Ty, false, GlobalValue::ExternalLinkage, 0, "__afl_prev_caller", 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 @@ -289,13 +362,24 @@ bool AFLCoverage::runOnModule(Module &M) { PrevLocShuffle.push_back(ConstantInt::get(Int32Ty, PrevLocSize)); Constant *PrevLocShuffleMask = ConstantVector::get(PrevLocShuffle); + + SmallVector PrevCallerShuffle = {UndefValue::get(Int32Ty)}; + + for (unsigned I = 0; I < PrevCallerSize - 1; ++I) + PrevCallerShuffle.push_back(ConstantInt::get(Int32Ty, I)); + + for (int I = PrevCallerSize; I < PrevCallerVecSize; ++I) + PrevCallerShuffle.push_back(ConstantInt::get(Int32Ty, PrevCallerSize)); + + Constant *PrevCallerShuffleMask = ConstantVector::get(PrevCallerShuffle); #endif // other constants we need ConstantInt *Zero = ConstantInt::get(Int8Ty, 0); ConstantInt *One = ConstantInt::get(Int8Ty, 1); - LoadInst *PrevCtx = NULL; // CTX sensitive coverage + Value *PrevCtx = NULL; // CTX sensitive coverage + LoadInst *PrevCaller = NULL; // K-CTX coverage /* Instrument all the things! */ @@ -319,12 +403,21 @@ bool AFLCoverage::runOnModule(Module &M) { IRBuilder<> IRB(&(*IP)); // Context sensitive coverage - if ((ctx_str || caller_str) && &BB == &F.getEntryBlock()) { + if (instrument_ctx && &BB == &F.getEntryBlock()) { - // load the context ID of the previous function and write to to a local - // variable on the stack - PrevCtx = IRB.CreateLoad(AFLContext); - PrevCtx->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); +#ifdef AFL_HAVE_VECTOR_INTRINSICS + if (ctx_k) { + PrevCaller = IRB.CreateLoad(AFLPrevCaller); + PrevCaller->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + PrevCtx = IRB.CreateZExt(IRB.CreateXorReduce(PrevCaller), IRB.getInt32Ty()); + } else +#endif + { + // load the context ID of the previous function and write to to a local variable on the stack + LoadInst* PrevCtxLoad = IRB.CreateLoad(AFLContext); + PrevCtxLoad->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + PrevCtx = PrevCtxLoad; + } // does the function have calls? and is any of the calls larger than one // basic block? @@ -356,10 +449,22 @@ bool AFLCoverage::runOnModule(Module &M) { if (has_calls) { Value *NewCtx = ConstantInt::get(Int32Ty, AFL_R(map_size)); - if (ctx_str) NewCtx = IRB.CreateXor(PrevCtx, NewCtx); - StoreInst *StoreCtx = IRB.CreateStore(NewCtx, AFLContext); - StoreCtx->setMetadata(M.getMDKindID("nosanitize"), - MDNode::get(C, None)); +#ifdef AFL_HAVE_VECTOR_INTRINSICS + if (ctx_k) { + Value *ShuffledPrevCaller = IRB.CreateShuffleVector( + PrevCaller, UndefValue::get(PrevCallerTy), PrevCallerShuffleMask); + Value *UpdatedPrevCaller = IRB.CreateInsertElement(ShuffledPrevCaller, NewCtx, (uint64_t)0); + + StoreInst * Store = IRB.CreateStore(UpdatedPrevCaller, AFLPrevCaller); + Store->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + } else +#endif + { + if (ctx_str) NewCtx = IRB.CreateXor(PrevCtx, NewCtx); + StoreInst *StoreCtx = IRB.CreateStore(NewCtx, AFLContext); + StoreCtx->setMetadata(M.getMDKindID("nosanitize"), + MDNode::get(C, None)); + } } @@ -413,16 +518,22 @@ bool AFLCoverage::runOnModule(Module &M) { // in CTX mode we have to restore the original context for the caller - // she might be calling other functions which need the correct CTX - if ((ctx_str || caller_str) && has_calls) { + if (instrument_ctx && has_calls) { Instruction *Inst = BB.getTerminator(); if (isa(Inst) || isa(Inst)) { IRBuilder<> Post_IRB(Inst); - StoreInst * RestoreCtx = Post_IRB.CreateStore(PrevCtx, AFLContext); + + StoreInst * RestoreCtx; +#ifdef AFL_HAVE_VECTOR_INTRINSICS + if (ctx_k) + RestoreCtx = IRB.CreateStore(PrevCaller, AFLPrevCaller); + else +#endif + RestoreCtx = Post_IRB.CreateStore(PrevCtx, AFLContext); RestoreCtx->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); - } } @@ -460,7 +571,7 @@ bool AFLCoverage::runOnModule(Module &M) { #endif PrevLocTrans = PrevLoc; - if (ctx_str || caller_str) + if (instrument_ctx) PrevLocTrans = IRB.CreateZExt(IRB.CreateXor(PrevLocTrans, PrevCtx), Int32Ty); else @@ -547,13 +658,20 @@ bool AFLCoverage::runOnModule(Module &M) { // in CTX mode we have to restore the original context for the caller - // she might be calling other functions which need the correct CTX. // Currently this is only needed for the Ubuntu clang-6.0 bug - if ((ctx_str || caller_str) && has_calls) { + if (instrument_ctx && has_calls) { Instruction *Inst = BB.getTerminator(); if (isa(Inst) || isa(Inst)) { IRBuilder<> Post_IRB(Inst); - StoreInst * RestoreCtx = Post_IRB.CreateStore(PrevCtx, AFLContext); + + StoreInst * RestoreCtx; +#ifdef AFL_HAVE_VECTOR_INTRINSICS + if (ctx_k) + RestoreCtx = IRB.CreateStore(PrevCaller, AFLPrevCaller); + else +#endif + RestoreCtx = Post_IRB.CreateStore(PrevCtx, AFLContext); RestoreCtx->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); diff --git a/instrumentation/llvm-alternative-coverage.h b/instrumentation/llvm-alternative-coverage.h new file mode 100644 index 00000000..0d7b3957 --- /dev/null +++ b/instrumentation/llvm-alternative-coverage.h @@ -0,0 +1,21 @@ +#ifndef AFL_NGRAM_CONFIG_H +#define AFL_NGRAM_CONFIG_H + +#include "types.h" + +#if (MAP_SIZE_POW2 <= 16) +typedef u16 PREV_LOC_T; +#elif (MAP_SIZE_POW2 <= 32) +typedef u32 PREV_LOC_T; +#else +typedef u64 PREV_LOC_T; +#endif + +/* Maximum ngram size */ +#define NGRAM_SIZE_MAX 16U + +/* Maximum K for top-K context sensitivity */ +#define CTX_MAX_K 32U + +#endif + diff --git a/instrumentation/llvm-ngram-coverage.h b/instrumentation/llvm-ngram-coverage.h deleted file mode 100644 index 666839c8..00000000 --- a/instrumentation/llvm-ngram-coverage.h +++ /dev/null @@ -1,18 +0,0 @@ -#ifndef AFL_NGRAM_CONFIG_H -#define AFL_NGRAM_CONFIG_H - -#include "types.h" - -#if (MAP_SIZE_POW2 <= 16) -typedef u16 PREV_LOC_T; -#elif (MAP_SIZE_POW2 <= 32) -typedef u32 PREV_LOC_T; -#else -typedef u64 PREV_LOC_T; -#endif - -/* Maximum ngram size */ -#define NGRAM_SIZE_MAX 16U - -#endif - diff --git a/src/afl-cc.c b/src/afl-cc.c index 0c689286..3c96beac 100644 --- a/src/afl-cc.c +++ b/src/afl-cc.c @@ -22,7 +22,7 @@ #include "types.h" #include "debug.h" #include "alloc-inl.h" -#include "llvm-ngram-coverage.h" +#include "llvm-alternative-coverage.h" #include #include -- cgit 1.4.1 From be5274d4a9c7fb835530be5054132253d2559ade Mon Sep 17 00:00:00 2001 From: Andrea Fioraldi Date: Thu, 4 Mar 2021 15:12:08 +0100 Subject: fix kctx compilation hang --- include/envs.h | 1 + instrumentation/afl-llvm-pass.so.cc | 13 ++++++++----- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/include/envs.h b/include/envs.h index 26f4de90..e8595ef7 100644 --- a/include/envs.h +++ b/include/envs.h @@ -81,6 +81,7 @@ static char *afl_environment_variables[] = { "AFL_LLVM_CMPLOG", "AFL_LLVM_INSTRIM", "AFL_LLVM_CTX", + "AFL_LLVM_CTX_K", "AFL_LLVM_DICT2FILE", "AFL_LLVM_DOCUMENT_IDS", "AFL_LLVM_INSTRIM_LOOPHEAD", diff --git a/instrumentation/afl-llvm-pass.so.cc b/instrumentation/afl-llvm-pass.so.cc index c58e9d95..fbf55f81 100644 --- a/instrumentation/afl-llvm-pass.so.cc +++ b/instrumentation/afl-llvm-pass.so.cc @@ -363,15 +363,18 @@ bool AFLCoverage::runOnModule(Module &M) { Constant *PrevLocShuffleMask = ConstantVector::get(PrevLocShuffle); + Constant *PrevCallerShuffleMask = NULL; SmallVector PrevCallerShuffle = {UndefValue::get(Int32Ty)}; - for (unsigned I = 0; I < PrevCallerSize - 1; ++I) - PrevCallerShuffle.push_back(ConstantInt::get(Int32Ty, I)); + if (ctx_k) { + for (unsigned I = 0; I < PrevCallerSize - 1; ++I) + PrevCallerShuffle.push_back(ConstantInt::get(Int32Ty, I)); - for (int I = PrevCallerSize; I < PrevCallerVecSize; ++I) - PrevCallerShuffle.push_back(ConstantInt::get(Int32Ty, PrevCallerSize)); + for (int I = PrevCallerSize; I < PrevCallerVecSize; ++I) + PrevCallerShuffle.push_back(ConstantInt::get(Int32Ty, PrevCallerSize)); - Constant *PrevCallerShuffleMask = ConstantVector::get(PrevCallerShuffle); + PrevCallerShuffleMask = ConstantVector::get(PrevCallerShuffle); + } #endif // other constants we need -- cgit 1.4.1 From af9aeb89d43919fb5e538778e5df9bf3ffaba0aa Mon Sep 17 00:00:00 2001 From: Andrea Fioraldi Date: Thu, 4 Mar 2021 15:26:15 +0100 Subject: afl-cc code for k-ctx --- src/afl-cc.c | 64 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 60 insertions(+), 4 deletions(-) diff --git a/src/afl-cc.c b/src/afl-cc.c index 3c96beac..97f32b2b 100644 --- a/src/afl-cc.c +++ b/src/afl-cc.c @@ -50,7 +50,7 @@ static u8 **cc_params; /* Parameters passed to the real CC */ static u32 cc_par_cnt = 1; /* Param count, including argv0 */ static u8 clang_mode; /* Invoked as afl-clang*? */ static u8 llvm_fullpath[PATH_MAX]; -static u8 instrument_mode, instrument_opt_mode, ngram_size, lto_mode; +static u8 instrument_mode, instrument_opt_mode, ngram_size, ctx_k, lto_mode; static u8 compiler_mode, plusplus_mode, have_instr_env = 0; static u8 have_gcc, have_llvm, have_gcc_plugin, have_lto, have_instr_list = 0; static u8 * lto_flag = AFL_CLANG_FLTO, *argvnull; @@ -75,6 +75,7 @@ enum { INSTRUMENT_OPT_CTX = 8, INSTRUMENT_OPT_NGRAM = 16, INSTRUMENT_OPT_CALLER = 32, + INSTRUMENT_OPT_CTX_K = 64, }; @@ -1282,11 +1283,21 @@ int main(int argc, char **argv, char **envp) { ngram_size = atoi(getenv("AFL_LLVM_NGRAM_SIZE")); if (ngram_size < 2 || ngram_size > NGRAM_SIZE_MAX) FATAL( - "NGRAM instrumentation mode must be between 2 and NGRAM_SIZE_MAX " + "K-CTX instrumentation mode must be between 2 and NGRAM_SIZE_MAX " "(%u)", NGRAM_SIZE_MAX); } + + if (getenv("AFL_LLVM_CTX_K")) { + + instrument_opt_mode |= INSTRUMENT_OPT_CTX_K; + ctx_k = atoi(getenv("AFL_LLVM_CTX_K")); + if (ctx_k < 1 || ctx_k > CTX_MAX_K) + FATAL( + "NGRAM instrumentation mode must be between 1 and CTX_MAX_K (%u)", CTX_MAX_K); + + } if (getenv("AFL_LLVM_INSTRUMENT")) { @@ -1382,6 +1393,32 @@ int main(int argc, char **argv, char **envp) { compiler_mode = CLANG; } + + if (strncasecmp(ptr2, "ctx-", strlen("ctx-")) == 0) { + + u8 *ptr3 = ptr2 + strlen("ctx-"); + while (*ptr3 && (*ptr3 < '0' || *ptr3 > '9')) + ptr3++; + + if (!*ptr3) { + + if ((ptr3 = getenv("AFL_LLVM_CTX_K")) == NULL) + FATAL( + "you must set the K-CTX K with (e.g. for value 2) " + "AFL_LLVM_INSTRUMENT=ctx-2"); + + } + + ctx_k = atoi(ptr3); + if (ctx_k < 1 || ctx_k > CTX_MAX_K) + FATAL( + "K-CTX instrumentation option must be between 1 and CTX_MAX_K (%u)", + CTX_MAX_K); + instrument_opt_mode |= (INSTRUMENT_OPT_CTX_K); + u8 *ptr4 = alloc_printf("%u", ctx_k); + setenv("AFL_LLVM_CTX_K", ptr4, 1); + + } if (strncasecmp(ptr2, "ctx", strlen("ctx")) == 0) { @@ -1437,6 +1474,20 @@ int main(int argc, char **argv, char **envp) { } + if ((instrument_opt_mode & INSTRUMENT_OPT_CTX) && + (instrument_opt_mode & INSTRUMENT_OPT_CTX_K)) { + + FATAL("you cannot set CTX and K-CTX together"); + + } + + if ((instrument_opt_mode & INSTRUMENT_OPT_CALLER) && + (instrument_opt_mode & INSTRUMENT_OPT_CTX_K)) { + + FATAL("you cannot set CALLER and K-CTX together"); + + } + if (instrument_opt_mode && instrument_mode == INSTRUMENT_DEFAULT && (compiler_mode == LLVM || compiler_mode == UNSET)) { @@ -1797,13 +1848,18 @@ int main(int argc, char **argv, char **envp) { } else { char *ptr2 = alloc_printf(" + NGRAM-%u", ngram_size); + char *ptr3 = alloc_printf(" + K-CTX-%u", ctx_k); + ptr = alloc_printf( - "%s%s%s%s", instrument_mode_string[instrument_mode], + "%s%s%s%s%s", instrument_mode_string[instrument_mode], (instrument_opt_mode & INSTRUMENT_OPT_CTX) ? " + CTX" : "", (instrument_opt_mode & INSTRUMENT_OPT_CALLER) ? " + CALLER" : "", - (instrument_opt_mode & INSTRUMENT_OPT_NGRAM) ? ptr2 : ""); + (instrument_opt_mode & INSTRUMENT_OPT_NGRAM) ? ptr2 : "", + (instrument_opt_mode & INSTRUMENT_OPT_CTX_K) ? ptr3 : "" + ); ck_free(ptr2); + ck_free(ptr3); } -- cgit 1.4.1 From c429021de10cf01878bd2150cc68c6b403db9335 Mon Sep 17 00:00:00 2001 From: Andrea Fioraldi Date: Fri, 5 Mar 2021 15:27:10 +0100 Subject: fix typos and format --- instrumentation/afl-llvm-pass.so.cc | 99 ++++++++++++++++++++++++------------- src/afl-cc.c | 16 +++--- 2 files changed, 72 insertions(+), 43 deletions(-) diff --git a/instrumentation/afl-llvm-pass.so.cc b/instrumentation/afl-llvm-pass.so.cc index fbf55f81..f4717345 100644 --- a/instrumentation/afl-llvm-pass.so.cc +++ b/instrumentation/afl-llvm-pass.so.cc @@ -212,24 +212,30 @@ bool AFLCoverage::runOnModule(Module &M) { PrevLocSize = ngram_size - 1; else PrevLocSize = 1; - + /* Decide K-ctx vector size (must be a power of two) */ VectorType *PrevCallerTy = NULL; if (ctx_k_str) - if (sscanf(ctx_k_str, "%u", &ctx_k) != 1 || ctx_k < 2 || - ctx_k > CTX_MAX_K) - FATAL("Bad value of AFL_CTX_K (must be between 2 and CTX_MAX_K (%u))", CTX_MAX_K); + if (sscanf(ctx_k_str, "%u", &ctx_k) != 1 || ctx_k < 2 || ctx_k > CTX_MAX_K) + FATAL("Bad value of AFL_CTX_K (must be between 2 and CTX_MAX_K (%u))", + CTX_MAX_K); if (ctx_k == 1) { + ctx_k = 0; instrument_ctx = true; - caller_str = ctx_k_str; // Enable CALLER instead + caller_str = ctx_k_str; // Enable CALLER instead + } + if (ctx_k) { + PrevCallerSize = ctx_k; instrument_ctx = true; + } + #else if (ngram_size_str) #ifndef LLVM_VERSION_PATCH @@ -274,8 +280,8 @@ bool AFLCoverage::runOnModule(Module &M) { if (ctx_k) PrevCallerTy = VectorType::get(IntLocTy, PrevCallerVecSize #if LLVM_VERSION_MAJOR >= 12 - , - false + , + false #endif ); #endif @@ -340,12 +346,13 @@ bool AFLCoverage::runOnModule(Module &M) { else #endif #if defined(__ANDROID__) || defined(__HAIKU__) - AFLPrevCaller = new GlobalVariable( - M, Int32Ty, false, GlobalValue::ExternalLinkage, 0, "__afl_prev_caller"); + AFLPrevCaller = + new GlobalVariable(M, Int32Ty, false, GlobalValue::ExternalLinkage, 0, + "__afl_prev_caller"); #else AFLPrevCaller = new GlobalVariable( - M, Int32Ty, false, GlobalValue::ExternalLinkage, 0, "__afl_prev_caller", 0, - GlobalVariable::GeneralDynamicTLSModel, 0, false); + M, Int32Ty, false, GlobalValue::ExternalLinkage, 0, "__afl_prev_caller", + 0, GlobalVariable::GeneralDynamicTLSModel, 0, false); #endif #ifdef AFL_HAVE_VECTOR_INTRINSICS @@ -362,11 +369,12 @@ bool AFLCoverage::runOnModule(Module &M) { PrevLocShuffle.push_back(ConstantInt::get(Int32Ty, PrevLocSize)); Constant *PrevLocShuffleMask = ConstantVector::get(PrevLocShuffle); - - Constant *PrevCallerShuffleMask = NULL; + + Constant * PrevCallerShuffleMask = NULL; SmallVector PrevCallerShuffle = {UndefValue::get(Int32Ty)}; if (ctx_k) { + for (unsigned I = 0; I < PrevCallerSize - 1; ++I) PrevCallerShuffle.push_back(ConstantInt::get(Int32Ty, I)); @@ -374,15 +382,17 @@ bool AFLCoverage::runOnModule(Module &M) { PrevCallerShuffle.push_back(ConstantInt::get(Int32Ty, PrevCallerSize)); PrevCallerShuffleMask = ConstantVector::get(PrevCallerShuffle); + } + #endif // other constants we need ConstantInt *Zero = ConstantInt::get(Int8Ty, 0); ConstantInt *One = ConstantInt::get(Int8Ty, 1); - Value *PrevCtx = NULL; // CTX sensitive coverage - LoadInst *PrevCaller = NULL; // K-CTX coverage + Value * PrevCtx = NULL; // CTX sensitive coverage + LoadInst *PrevCaller = NULL; // K-CTX coverage /* Instrument all the things! */ @@ -410,16 +420,25 @@ bool AFLCoverage::runOnModule(Module &M) { #ifdef AFL_HAVE_VECTOR_INTRINSICS if (ctx_k) { + PrevCaller = IRB.CreateLoad(AFLPrevCaller); - PrevCaller->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); - PrevCtx = IRB.CreateZExt(IRB.CreateXorReduce(PrevCaller), IRB.getInt32Ty()); - } else + PrevCaller->setMetadata(M.getMDKindID("nosanitize"), + MDNode::get(C, None)); + PrevCtx = + IRB.CreateZExt(IRB.CreateXorReduce(PrevCaller), IRB.getInt32Ty()); + + } else + #endif { - // load the context ID of the previous function and write to to a local variable on the stack - LoadInst* PrevCtxLoad = IRB.CreateLoad(AFLContext); - PrevCtxLoad->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + + // load the context ID of the previous function and write to to a + // local variable on the stack + LoadInst *PrevCtxLoad = IRB.CreateLoad(AFLContext); + PrevCtxLoad->setMetadata(M.getMDKindID("nosanitize"), + MDNode::get(C, None)); PrevCtx = PrevCtxLoad; + } // does the function have calls? and is any of the calls larger than one @@ -454,19 +473,28 @@ bool AFLCoverage::runOnModule(Module &M) { Value *NewCtx = ConstantInt::get(Int32Ty, AFL_R(map_size)); #ifdef AFL_HAVE_VECTOR_INTRINSICS if (ctx_k) { + Value *ShuffledPrevCaller = IRB.CreateShuffleVector( - PrevCaller, UndefValue::get(PrevCallerTy), PrevCallerShuffleMask); - Value *UpdatedPrevCaller = IRB.CreateInsertElement(ShuffledPrevCaller, NewCtx, (uint64_t)0); + PrevCaller, UndefValue::get(PrevCallerTy), + PrevCallerShuffleMask); + Value *UpdatedPrevCaller = IRB.CreateInsertElement( + ShuffledPrevCaller, NewCtx, (uint64_t)0); + + StoreInst *Store = + IRB.CreateStore(UpdatedPrevCaller, AFLPrevCaller); + Store->setMetadata(M.getMDKindID("nosanitize"), + MDNode::get(C, None)); + + } else - StoreInst * Store = IRB.CreateStore(UpdatedPrevCaller, AFLPrevCaller); - Store->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); - } else #endif { + if (ctx_str) NewCtx = IRB.CreateXor(PrevCtx, NewCtx); StoreInst *StoreCtx = IRB.CreateStore(NewCtx, AFLContext); StoreCtx->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + } } @@ -528,15 +556,16 @@ bool AFLCoverage::runOnModule(Module &M) { IRBuilder<> Post_IRB(Inst); - StoreInst * RestoreCtx; -#ifdef AFL_HAVE_VECTOR_INTRINSICS + StoreInst *RestoreCtx; + #ifdef AFL_HAVE_VECTOR_INTRINSICS if (ctx_k) RestoreCtx = IRB.CreateStore(PrevCaller, AFLPrevCaller); - else -#endif - RestoreCtx = Post_IRB.CreateStore(PrevCtx, AFLContext); + else + #endif + RestoreCtx = Post_IRB.CreateStore(PrevCtx, AFLContext); RestoreCtx->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); + } } @@ -667,14 +696,14 @@ bool AFLCoverage::runOnModule(Module &M) { if (isa(Inst) || isa(Inst)) { IRBuilder<> Post_IRB(Inst); - - StoreInst * RestoreCtx; + + StoreInst *RestoreCtx; #ifdef AFL_HAVE_VECTOR_INTRINSICS if (ctx_k) RestoreCtx = IRB.CreateStore(PrevCaller, AFLPrevCaller); - else + else #endif - RestoreCtx = Post_IRB.CreateStore(PrevCtx, AFLContext); + RestoreCtx = Post_IRB.CreateStore(PrevCtx, AFLContext); RestoreCtx->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None)); diff --git a/src/afl-cc.c b/src/afl-cc.c index 97f32b2b..e4ea66e4 100644 --- a/src/afl-cc.c +++ b/src/afl-cc.c @@ -1283,19 +1283,19 @@ int main(int argc, char **argv, char **envp) { ngram_size = atoi(getenv("AFL_LLVM_NGRAM_SIZE")); if (ngram_size < 2 || ngram_size > NGRAM_SIZE_MAX) FATAL( - "K-CTX instrumentation mode must be between 2 and NGRAM_SIZE_MAX " + "NGRAM instrumentation mode must be between 2 and NGRAM_SIZE_MAX " "(%u)", NGRAM_SIZE_MAX); } - + if (getenv("AFL_LLVM_CTX_K")) { instrument_opt_mode |= INSTRUMENT_OPT_CTX_K; ctx_k = atoi(getenv("AFL_LLVM_CTX_K")); if (ctx_k < 1 || ctx_k > CTX_MAX_K) - FATAL( - "NGRAM instrumentation mode must be between 1 and CTX_MAX_K (%u)", CTX_MAX_K); + FATAL("K-CTX instrumentation mode must be between 1 and CTX_MAX_K (%u)", + CTX_MAX_K); } @@ -1393,7 +1393,7 @@ int main(int argc, char **argv, char **envp) { compiler_mode = CLANG; } - + if (strncasecmp(ptr2, "ctx-", strlen("ctx-")) == 0) { u8 *ptr3 = ptr2 + strlen("ctx-"); @@ -1412,7 +1412,8 @@ int main(int argc, char **argv, char **envp) { ctx_k = atoi(ptr3); if (ctx_k < 1 || ctx_k > CTX_MAX_K) FATAL( - "K-CTX instrumentation option must be between 1 and CTX_MAX_K (%u)", + "K-CTX instrumentation option must be between 1 and CTX_MAX_K " + "(%u)", CTX_MAX_K); instrument_opt_mode |= (INSTRUMENT_OPT_CTX_K); u8 *ptr4 = alloc_printf("%u", ctx_k); @@ -1855,8 +1856,7 @@ int main(int argc, char **argv, char **envp) { (instrument_opt_mode & INSTRUMENT_OPT_CTX) ? " + CTX" : "", (instrument_opt_mode & INSTRUMENT_OPT_CALLER) ? " + CALLER" : "", (instrument_opt_mode & INSTRUMENT_OPT_NGRAM) ? ptr2 : "", - (instrument_opt_mode & INSTRUMENT_OPT_CTX_K) ? ptr3 : "" - ); + (instrument_opt_mode & INSTRUMENT_OPT_CTX_K) ? ptr3 : ""); ck_free(ptr2); ck_free(ptr3); -- cgit 1.4.1