LTO whitelist functionality rewritten, now anything can be skipped

5 files changed, 83 insertions, 71 deletions

diff --git a/docs/Changelog.md b/docs/Changelog.md
index 08952717..ae398b66 100644
--- a/docs/Changelog.md
+++ b/docs/Changelog.md
@@ -25,6 +25,8 @@ sending a mail to <afl-users+subscribe@googlegroups.com>.
       which needs 3.8.0)
     - small change to cmplog to make it work with current llvm 11-dev
     - added AFL_LLVM_LAF_ALL, sets all laf-intel settings
+    - LTO whitelist functionality rewritten, now main, _init etc functions
+      need not to be whitelisted anymore
   - fixed afl-gcc/afl-as that could break on fast systems reusing pids in
     the same second
   - added lots of dictionaries from oss-fuzz, go-fuzz and Jakub Wilk
diff --git a/llvm_mode/README.lto.md b/llvm_mode/README.lto.md
index 0415076a..4790c167 100644
--- a/llvm_mode/README.lto.md
+++ b/llvm_mode/README.lto.md
@@ -190,11 +190,6 @@ target will likely crash when started. This can be avoided by compiling with
 
 This can e.g. happen with OpenSSL.
 
-## Upcoming Work
-
-1. Currently the LTO whitelist feature does not allow to instrument main,
-   start and init functions
-
 ## History
 
 This was originally envisioned by hexcoder- in Summer 2019, however we saw no
diff --git a/llvm_mode/afl-llvm-lto-instrim.so.cc b/llvm_mode/afl-llvm-lto-instrim.so.cc
index a7d9b756..27504e8d 100644
--- a/llvm_mode/afl-llvm-lto-instrim.so.cc
+++ b/llvm_mode/afl-llvm-lto-instrim.so.cc
@@ -561,6 +561,17 @@ struct InsTrimLTO : public ModulePass {
       if (F.size() < function_minimum_size) continue;
       if (isBlacklisted(&F)) continue;
 
+      // whitelist check
+      AttributeList Attrs = F.getAttributes();
+      if (Attrs.hasAttribute(-1, StringRef("skipinstrument"))) {
+
+        if (debug)
+          fprintf(stderr, "DEBUG: Function %s is not whitelisted\n",
+                  F.getName().str().c_str());
+        continue;
+
+      }
+
       std::unordered_set<BasicBlock *> MS;
       if (!MarkSetOpt) {
 
diff --git a/llvm_mode/afl-llvm-lto-instrumentation.so.cc b/llvm_mode/afl-llvm-lto-instrumentation.so.cc
index f44b336e..cbe68171 100644
--- a/llvm_mode/afl-llvm-lto-instrumentation.so.cc
+++ b/llvm_mode/afl-llvm-lto-instrumentation.so.cc
@@ -197,6 +197,17 @@ bool AFLLTOPass::runOnModule(Module &M) {
     if (F.size() < function_minimum_size) continue;
     if (isBlacklisted(&F)) continue;
 
+    // whitelist check
+    AttributeList Attrs = F.getAttributes();
+    if (Attrs.hasAttribute(-1, StringRef("skipinstrument"))) {
+
+      if (debug)
+        fprintf(stderr, "DEBUG: Function %s is not whitelisted\n",
+                F.getName().str().c_str());
+      continue;
+
+    }
+
     std::vector<BasicBlock *> InsBlocks;
 
     if (autodictionary) {
diff --git a/llvm_mode/afl-llvm-lto-whitelist.so.cc b/llvm_mode/afl-llvm-lto-whitelist.so.cc
index a116c4ea..8856ce21 100644
--- a/llvm_mode/afl-llvm-lto-whitelist.so.cc
+++ b/llvm_mode/afl-llvm-lto-whitelist.so.cc
@@ -122,64 +122,65 @@ bool AFLwhitelist::runOnModule(Module &M) {
 
   for (auto &F : M) {
 
+    if (F.size() < 1) continue;
+    // fprintf(stderr, "F:%s\n", F.getName().str().c_str());
     if (isBlacklisted(&F)) continue;
 
-    for (auto &BB : F) {
+    BasicBlock::iterator IP = F.getEntryBlock().getFirstInsertionPt();
+    IRBuilder<>          IRB(&(*IP));
 
-      BasicBlock::iterator IP = BB.getFirstInsertionPt();
-      IRBuilder<>          IRB(&(*IP));
+    if (!myWhitelist.empty()) {
 
-      if (!myWhitelist.empty()) {
+      bool instrumentFunction = false;
 
-        bool instrumentBlock = false;
+      /* Get the current location using debug information.
+       * For now, just instrument the block if we are not able
+       * to determine our location. */
+      DebugLoc Loc = IP->getDebugLoc();
+      if (Loc) {
 
-        /* Get the current location using debug information.
-         * For now, just instrument the block if we are not able
-         * to determine our location. */
-        DebugLoc Loc = IP->getDebugLoc();
-        if (Loc) {
+        DILocation *cDILoc = dyn_cast<DILocation>(Loc.getAsMDNode());
 
-          DILocation *cDILoc = dyn_cast<DILocation>(Loc.getAsMDNode());
+        unsigned int instLine = cDILoc->getLine();
+        StringRef    instFilename = cDILoc->getFilename();
 
-          unsigned int instLine = cDILoc->getLine();
-          StringRef    instFilename = cDILoc->getFilename();
+        if (instFilename.str().empty()) {
 
-          if (instFilename.str().empty()) {
+          /* If the original location is empty, try using the inlined location
+           */
+          DILocation *oDILoc = cDILoc->getInlinedAt();
+          if (oDILoc) {
 
-            /* If the original location is empty, try using the inlined location
-             */
-            DILocation *oDILoc = cDILoc->getInlinedAt();
-            if (oDILoc) {
-
-              instFilename = oDILoc->getFilename();
-              instLine = oDILoc->getLine();
-
-            }
+            instFilename = oDILoc->getFilename();
+            instLine = oDILoc->getLine();
 
           }
 
-          (void)instLine;
+        }
 
-          /* Continue only if we know where we actually are */
-          if (!instFilename.str().empty()) {
+        (void)instLine;
 
-            for (std::list<std::string>::iterator it = myWhitelist.begin();
-                 it != myWhitelist.end(); ++it) {
+        if (debug)
+          SAYF(cMGN "[D] " cRST "function %s is in file %s\n",
+               F.getName().str().c_str(), instFilename.str().c_str());
+        /* Continue only if we know where we actually are */
+        if (!instFilename.str().empty()) {
 
-              /* We don't check for filename equality here because
-               * filenames might actually be full paths. Instead we
-               * check that the actual filename ends in the filename
-               * specified in the list. */
-              if (instFilename.str().length() >= it->length()) {
+          for (std::list<std::string>::iterator it = myWhitelist.begin();
+               it != myWhitelist.end(); ++it) {
 
-                if (instFilename.str().compare(
-                        instFilename.str().length() - it->length(),
-                        it->length(), *it) == 0) {
+            /* We don't check for filename equality here because
+             * filenames might actually be full paths. Instead we
+             * check that the actual filename ends in the filename
+             * specified in the list. */
+            if (instFilename.str().length() >= it->length()) {
 
-                  instrumentBlock = true;
-                  break;
+              if (instFilename.str().compare(
+                      instFilename.str().length() - it->length(), it->length(),
+                      *it) == 0) {
 
-                }
+                instrumentFunction = true;
+                break;
 
               }
 
@@ -189,43 +190,35 @@ bool AFLwhitelist::runOnModule(Module &M) {
 
         }
 
-        /* Either we couldn't figure out our location or the location is
-         * not whitelisted, so we skip instrumentation.
-         * We do this by renaming the function. */
-        if (!instrumentBlock) {
-
-          if (F.getName().compare("main") == 0 ||
-              F.getName().compare("start") == 0 ||
-              F.getName().compare("_start") == 0 ||
-              F.getName().compare("init") == 0 ||
-              F.getName().compare("_init") == 0) {
-
-            // We do not honor be_quiet for this one
-            WARNF("Cannot ignore functions main/init/start");
-
-          } else {
-
-            // StringRef newName = StringRef("ign.") + F.getName();
-            if (debug)
-              SAYF(cMGN "[D] " cRST "renamed %s to ign.%s\n",
-                   F.getName().str().c_str(), F.getName().str().c_str());
-            Function *_F(&F);
-            _F->setName("ign." + F.getName());
-
-          }
+      }
 
-        } else if (debug)
+      /* Either we couldn't figure out our location or the location is
+       * not whitelisted, so we skip instrumentation.
+       * We do this by renaming the function. */
+      if (instrumentFunction == true) {
 
+        if (debug)
           SAYF(cMGN "[D] " cRST "function %s is in whitelist\n",
                F.getName().str().c_str());
 
       } else {
 
-        PFATAL("Whitelist is empty");
+        if (debug)
+          SAYF(cMGN "[D] " cRST "function %s is NOT in whitelist\n",
+               F.getName().str().c_str());
+
+        auto &        Ctx = F.getContext();
+        AttributeList Attrs = F.getAttributes();
+        AttrBuilder   NewAttrs;
+        NewAttrs.addAttribute("skipinstrument");
+        F.setAttributes(
+            Attrs.addAttributes(Ctx, AttributeList::FunctionIndex, NewAttrs));
 
       }
 
-      break;
+    } else {
+
+      PFATAL("Whitelist is empty");
 
     }