added AFL_LLVM_SKIPSINGLEBLOCK and changed default behaviour to instrument single block functions

7 files changed, 38 insertions, 22 deletions

diff --git a/docs/Changelog.md b/docs/Changelog.md
index 35240021..8bcc8949 100644
--- a/docs/Changelog.md
+++ b/docs/Changelog.md
@@ -16,6 +16,8 @@ sending a mail to <afl-users+subscribe@googlegroups.com>.
      - an old, old bug in afl that would show negative stability in rare
        circumstances is now hopefully fixed
   - llvm_mode:
+     - afl-clang-fast/lto now do not skip single block functions. This
+       behaviour can be reactivated with AFL_LLVM_SKIPSINGLEBLOCK
      - if LLVM 11 is installed the posix shm_open+mmap is used and a fixed
        address for the shared memory map is used as this increases the
        fuzzing speed
diff --git a/docs/env_variables.md b/docs/env_variables.md
index f5d28981..36e5a432 100644
--- a/docs/env_variables.md
+++ b/docs/env_variables.md
@@ -83,6 +83,10 @@ tools make fairly broad use of environmental variables:
 The native instrumentation helpers (llvm_mode and gcc_plugin) accept a subset
 of the settings discussed in section #1, with the exception of:
 
+    - Setting AFL_LLVM_SKIPSINGLEBLOCK=1 will skip instrumenting
+      functions with a single basic block. This is useful for most C and
+      some C++ targets. This works for all instrumentation modes.
+
   - AFL_AS, since this toolchain does not directly invoke GNU as.
 
   - TMPDIR and AFL_KEEP_ASSEMBLY, since no temporary assembly files are
@@ -156,10 +160,6 @@ Then there are a few specific features that are only available in llvm_mode:
       afl-fuzz will only be able to see the path the loop took, but not how
       many times it was called (unless it is a complex loop).
 
-    - Setting AFL_LLVM_INSTRIM_SKIPSINGLEBLOCK=1 will skip instrumenting
-      functions with a single basic block. This is useful for most C and
-      some C++ targets.
-
     See llvm_mode/README.instrim.md
 
 ### NGRAM
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)) {
diff --git a/src/afl-common.c b/src/afl-common.c
index 54b2e790..d9d57863 100644
--- a/src/afl-common.c
+++ b/src/afl-common.c
@@ -64,15 +64,15 @@ char *afl_environment_variables[] = {
     "AFL_LD_PRELOAD", "AFL_LD_VERBOSE", "AFL_LLVM_CMPLOG", "AFL_LLVM_INSTRIM",
     "AFL_LLVM_CTX", "AFL_LLVM_INSTRUMENT", "AFL_LLVM_INSTRIM_LOOPHEAD",
     "AFL_LLVM_LTO_AUTODICTIONARY", "AFL_LLVM_AUTODICTIONARY",
-    "AFL_LLVM_INSTRIM_SKIPSINGLEBLOCK", "AFL_LLVM_LAF_SPLIT_COMPARES",
-    "AFL_LLVM_LAF_SPLIT_COMPARES_BITW", "AFL_LLVM_LAF_SPLIT_FLOATS",
-    "AFL_LLVM_LAF_SPLIT_SWITCHES", "AFL_LLVM_LAF_TRANSFORM_COMPARES",
-    "AFL_LLVM_MAP_ADDR", "AFL_LLVM_MAP_DYNAMIC", "AFL_LLVM_NGRAM_SIZE",
-    "AFL_NGRAM_SIZE", "AFL_LLVM_NOT_ZERO", "AFL_LLVM_WHITELIST",
-    "AFL_LLVM_SKIP_NEVERZERO", "AFL_NO_AFFINITY", "AFL_LLVM_LTO_STARTID",
-    "AFL_LLVM_LTO_DONTWRITEID", "AFL_NO_ARITH", "AFL_NO_BUILTIN",
-    "AFL_NO_CPU_RED", "AFL_NO_FORKSRV", "AFL_NO_UI", "AFL_NO_PYTHON",
-    "AFL_UNTRACER_FILE",
+    "AFL_LLVM_SKIPSINGLEBLOCK", "AFL_LLVM_INSTRIM_SKIPSINGLEBLOCK",
+    "AFL_LLVM_LAF_SPLIT_COMPARES", "AFL_LLVM_LAF_SPLIT_COMPARES_BITW",
+    "AFL_LLVM_LAF_SPLIT_FLOATS", "AFL_LLVM_LAF_SPLIT_SWITCHES",
+    "AFL_LLVM_LAF_TRANSFORM_COMPARES", "AFL_LLVM_MAP_ADDR",
+    "AFL_LLVM_MAP_DYNAMIC", "AFL_LLVM_NGRAM_SIZE", "AFL_NGRAM_SIZE",
+    "AFL_LLVM_NOT_ZERO", "AFL_LLVM_WHITELIST", "AFL_LLVM_SKIP_NEVERZERO",
+    "AFL_NO_AFFINITY", "AFL_LLVM_LTO_STARTID", "AFL_LLVM_LTO_DONTWRITEID",
+    "AFL_NO_ARITH", "AFL_NO_BUILTIN", "AFL_NO_CPU_RED", "AFL_NO_FORKSRV",
+    "AFL_NO_UI", "AFL_NO_PYTHON", "AFL_UNTRACER_FILE",
     "AFL_NO_X86",  // not really an env but we dont want to warn on it
     "AFL_MAP_SIZE", "AFL_MAPSIZE", "AFL_PATH", "AFL_PERFORMANCE_FILE",
     //"AFL_PERSISTENT", // not implemented anymore, so warn additionally