From cc1fe2f2d29b4b306c1558189251adae781dbb70 Mon Sep 17 00:00:00 2001
From: van Hauser <vh@thc.org>
Date: Sat, 15 Aug 2020 22:59:49 +0200
Subject: skip instrumenting blocks following __afl_loop to improve stability
 in LTO. not in afl-llvm-pass and instrim because they are outdated, sancov
 cant be fixed

---
 docs/Changelog.md                            |  2 ++
 llvm_mode/afl-llvm-lto-instrumentation.so.cc | 41 ++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/docs/Changelog.md b/docs/Changelog.md
index ea7c7caf..ead4ff26 100644
--- a/docs/Changelog.md
+++ b/docs/Changelog.md
@@ -35,6 +35,8 @@ sending a mail to <afl-users+subscribe@googlegroups.com>.
      - skipping ctors and ifuncs for instrumentation
      - LTO: switch default to the dynamic memory map, set AFL_LLVM_MAP_ADDR
             for a fixed map address (eg. 0x10000)
+     - LTO: improved stability for persistent mode, no other instrumentation
+            has that advantage
      - LTO: laf-intel and redqueen/cmplog are now applied at link time
             to prevent llvm optimizing away the splits
      - LTO: autodictionary mode is a default
diff --git a/llvm_mode/afl-llvm-lto-instrumentation.so.cc b/llvm_mode/afl-llvm-lto-instrumentation.so.cc
index a4caf77b..2b99d4c6 100644
--- a/llvm_mode/afl-llvm-lto-instrumentation.so.cc
+++ b/llvm_mode/afl-llvm-lto-instrumentation.so.cc
@@ -103,6 +103,7 @@ bool AFLLTOPass::runOnModule(Module &M) {
   std::vector<std::string>         dictionary;
   std::vector<CallInst *>          calls;
   DenseMap<Value *, std::string *> valueMap;
+  std::vector<BasicBlock *>        BlockList;
   char *                           ptr;
   FILE *                           documentFile = NULL;
 
@@ -310,6 +311,24 @@ bool AFLLTOPass::runOnModule(Module &M) {
             isStrncasecmp &= !FuncName.compare("strncasecmp");
             isIntMemcpy &= !FuncName.compare("llvm.memcpy.p0i8.p0i8.i64");
 
+            /* we do something different here, putting this BB and the
+               successors in a block map */
+            if (!FuncName.compare("__afl_persistent_loop")) {
+
+              BlockList.push_back(&BB);
+              /*
+                            for (succ_iterator SI = succ_begin(&BB), SE =
+                 succ_end(&BB); SI != SE; ++SI) {
+
+                              BasicBlock *succ = *SI;
+                              BlockList.push_back(succ);
+
+                            }
+
+              */
+
+            }
+
             if (!isStrcmp && !isMemcmp && !isStrncmp && !isStrcasecmp &&
                 !isStrncasecmp && !isIntMemcpy)
               continue;
@@ -603,6 +622,28 @@ bool AFLLTOPass::runOnModule(Module &M) {
         uint32_t                  fs = origBB->getParent()->size();
         uint32_t                  countto;
 
+        if (BlockList.size()) {
+
+          int skip = 0;
+          for (uint32_t k = 0; k < BlockList.size(); k++) {
+
+            if (origBB == BlockList[k]) {
+
+              if (debug)
+                fprintf(
+                    stderr,
+                    "DEBUG: Function %s skipping BB with/after __afl_loop\n",
+                    F.getName().str().c_str());
+              skip = 1;
+
+            }
+
+          }
+
+          if (skip) continue;
+
+        }
+
         for (succ_iterator SI = succ_begin(origBB), SE = succ_end(origBB);
              SI != SE; ++SI) {
 
-- 
cgit 1.4.1