added AFL_LLVM_SKIPSINGLEBLOCK and changed default behaviour to instrument single block functions

4 files changed, 23 insertions, 9 deletions

diff --git a/llvm_mode/LLVMInsTrim.so.cc b/llvm_mode/LLVMInsTrim.so.cc
index ad046a8b..ed6c79e8 100644
--- a/llvm_mode/LLVMInsTrim.so.cc
+++ b/llvm_mode/LLVMInsTrim.so.cc
@@ -134,7 +134,8 @@ struct InsTrim : public ModulePass {
 
     }
 
-    if (getenv("AFL_LLVM_INSTRIM_SKIPSINGLEBLOCK") != NULL)
+    if (getenv("AFL_LLVM_INSTRIM_SKIPSINGLEBLOCK") ||
+        getenv("AFL_LLVM_SKIPSINGLEBLOCK"))
       function_minimum_size = 2;
 
     unsigned PrevLocSize = 0;
@@ -394,7 +395,7 @@ struct InsTrim : public ModulePass {
               if ((callInst = dyn_cast<CallInst>(&IN))) {
 
                 Function *Callee = callInst->getCalledFunction();
-                if (!Callee || Callee->size() < 2)
+                if (!Callee || Callee->size() < function_minimum_size)
                   continue;
                 else {
 
diff --git a/llvm_mode/afl-llvm-lto-instrim.so.cc b/llvm_mode/afl-llvm-lto-instrim.so.cc
index 6371a3cc..a686bb81 100644
--- a/llvm_mode/afl-llvm-lto-instrim.so.cc
+++ b/llvm_mode/afl-llvm-lto-instrim.so.cc
@@ -172,7 +172,8 @@ struct InsTrimLTO : public ModulePass {
 
     }
 
-    if (getenv("AFL_LLVM_INSTRIM_SKIPSINGLEBLOCK") != NULL)
+    if (getenv("AFL_LLVM_INSTRIM_SKIPSINGLEBLOCK") ||
+        getenv("AFL_LLVM_SKIPSINGLEBLOCK"))
       function_minimum_size = 2;
 
     // this is our default
diff --git a/llvm_mode/afl-llvm-lto-instrumentation.so.cc b/llvm_mode/afl-llvm-lto-instrumentation.so.cc
index 0e353fdf..f44b336e 100644
--- a/llvm_mode/afl-llvm-lto-instrumentation.so.cc
+++ b/llvm_mode/afl-llvm-lto-instrumentation.so.cc
@@ -87,6 +87,7 @@ class AFLLTOPass : public ModulePass {
 
  protected:
   int      afl_global_id = 1, debug = 0, autodictionary = 0;
+  uint32_t function_minimum_size = 1;
   uint32_t be_quiet = 0, inst_blocks = 0, inst_funcs = 0, total_instr = 0;
   uint64_t map_addr = 0x10000;
   char *   skip_nozero = NULL;
@@ -124,6 +125,10 @@ bool AFLLTOPass::runOnModule(Module &M) {
 
   if (getenv("AFL_LLVM_MAP_DYNAMIC")) map_addr = 0;
 
+  if (getenv("AFL_LLVM_INSTRIM_SKIPSINGLEBLOCK") ||
+      getenv("AFL_LLVM_SKIPSINGLEBLOCK"))
+    function_minimum_size = 2;
+
   if ((ptr = getenv("AFL_LLVM_MAP_ADDR"))) {
 
     uint64_t val;
@@ -189,7 +194,7 @@ bool AFLLTOPass::runOnModule(Module &M) {
 
     // fprintf(stderr, "DEBUG: Function %s\n", F.getName().str().c_str());
 
-    if (F.size() < 2) continue;
+    if (F.size() < function_minimum_size) continue;
     if (isBlacklisted(&F)) continue;
 
     std::vector<BasicBlock *> InsBlocks;
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)) {