LTO whitelist functionality rewritten, now anything can be skipped

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

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

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) {
 

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) {

llvm_mode/afl-llvm-lto-whitelist.so.cc

@@ -122,16 +122,16 @@ 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 = BB.getFirstInsertionPt();
+    BasicBlock::iterator IP = F.getEntryBlock().getFirstInsertionPt();
     IRBuilder<>          IRB(&(*IP));
 
     if (!myWhitelist.empty()) {
 
-        bool instrumentBlock = false;
+      bool instrumentFunction = false;
 
       /* Get the current location using debug information.
        * For now, just instrument the block if we are not able
@@ -160,6 +160,9 @@ bool AFLwhitelist::runOnModule(Module &M) {
 
         (void)instLine;
 
+        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()) {
 
@@ -173,10 +176,10 @@ bool AFLwhitelist::runOnModule(Module &M) {
             if (instFilename.str().length() >= it->length()) {
 
               if (instFilename.str().compare(
-                        instFilename.str().length() - it->length(),
-                        it->length(), *it) == 0) {
+                      instFilename.str().length() - it->length(), it->length(),
+                      *it) == 0) {
 
-                  instrumentBlock = true;
+                instrumentFunction = true;
                 break;
 
               }
@@ -192,40 +195,30 @@ 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 (instrumentFunction == true) {
 
-          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)
-
           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");
 
     }