added AFL_LLVM_SKIPSINGLEBLOCK and changed default behaviour to instrument single block functions

1 files changed, 12 insertions, 5 deletions

diff --git a/llvm_mode/afl-llvm-pass.so.cc b/llvm_mode/afl-llvm-pass.so.cc
index 0d9e0aba..2d23ad21 100644
--- a/llvm_mode/afl-llvm-pass.so.cc
+++ b/llvm_mode/afl-llvm-pass.so.cc
@@ -84,6 +84,7 @@ class AFLCoverage : public ModulePass {
   uint32_t ngram_size = 0;
   uint32_t debug = 0;
   uint32_t map_size = MAP_SIZE;
+  uint32_t function_minimum_size = 1;
   char *   ctx_str = NULL, *skip_nozero = NULL;
 
 };
@@ -182,6 +183,10 @@ bool AFLCoverage::runOnModule(Module &M) {
 #endif
   skip_nozero = getenv("AFL_LLVM_SKIP_NEVERZERO");
 
+  if (getenv("AFL_LLVM_INSTRIM_SKIPSINGLEBLOCK") ||
+      getenv("AFL_LLVM_SKIPSINGLEBLOCK"))
+    function_minimum_size = 2;
+
   unsigned PrevLocSize = 0;
 
   char *ngram_size_str = getenv("AFL_LLVM_NGRAM_SIZE");
@@ -294,13 +299,15 @@ bool AFLCoverage::runOnModule(Module &M) {
 
     if (!isInWhitelist(&F)) continue;
 
+    if (F.size() < function_minimum_size) continue;
+
     for (auto &BB : F) {
 
       BasicBlock::iterator IP = BB.getFirstInsertionPt();
       IRBuilder<>          IRB(&(*IP));
 
       // Context sensitive coverage
-      if (ctx_str && &BB == &F.getEntryBlock() && F.size() > 1) {
+      if (ctx_str && &BB == &F.getEntryBlock()) {
 
         // load the context ID of the previous function and write to to a local
         // variable on the stack
@@ -318,7 +325,7 @@ bool AFLCoverage::runOnModule(Module &M) {
             if ((callInst = dyn_cast<CallInst>(&IN))) {
 
               Function *Callee = callInst->getCalledFunction();
-              if (!Callee || Callee->size() < 2)
+              if (!Callee || Callee->size() < function_minimum_size)
                 continue;
               else {
 
@@ -389,11 +396,11 @@ bool AFLCoverage::runOnModule(Module &M) {
       }
 
       // fprintf(stderr, " == %d\n", more_than_one);
-      if (more_than_one != 1) {
+      if (F.size() > 1 && more_than_one != 1) {
 
         // 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 && has_calls && F.size() > 1) {
+        if (ctx_str && has_calls) {
 
           Instruction *Inst = BB.getTerminator();
           if (isa<ReturnInst>(Inst) || isa<ResumeInst>(Inst)) {
@@ -526,7 +533,7 @@ 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 && has_calls && F.size() > 1) {
+      if (ctx_str && has_calls) {
 
         Instruction *Inst = BB.getTerminator();
         if (isa<ReturnInst>(Inst) || isa<ResumeInst>(Inst)) {