added whitelist+blacklist to all llvm_mode passes

7 files changed, 343 insertions, 48 deletions

diff --git a/docs/ChangeLog b/docs/ChangeLog
index 33c6f618..c1d53379 100644
--- a/docs/ChangeLog
+++ b/docs/ChangeLog
@@ -27,7 +27,7 @@ Version ++2.60d (develop):
        instrumentation. compile normally and set AFL_LLVM_USE_TRACE_PC :)
   - afl-cmin is now a sh script (invoking awk) instead of bash for portability
     the original script is still present as afl-cmin.bash
-  - added blacklisted function check in all modules of llvm_mode
+  - added blacklist and whitelisting function check in all modules of llvm_mode
   - added fix from Debian project to compile libdislocator and libtokencap
 
 
diff --git a/llvm_mode/LLVMInsTrim.so.cc b/llvm_mode/LLVMInsTrim.so.cc
index 11451b43..24df6d42 100644
--- a/llvm_mode/LLVMInsTrim.so.cc
+++ b/llvm_mode/LLVMInsTrim.so.cc
@@ -144,19 +144,6 @@ struct InsTrim : public ModulePass {
     // this is our default
     MarkSetOpt = true;
 
-    /*    // I dont think this makes sense to port into LLVMInsTrim
-          char* inst_ratio_str = getenv("AFL_INST_RATIO");
-          unsigned int inst_ratio = 100;
-          if (inst_ratio_str) {
-
-           if (sscanf(inst_ratio_str, "%u", &inst_ratio) != 1 || !inst_ratio ||
-       inst_ratio > 100) FATAL("Bad value of AFL_INST_RATIO (must be between 1
-       and 100)");
-
-          }
-
-    */
-
     LLVMContext &C = M.getContext();
     IntegerType *Int8Ty = IntegerType::getInt8Ty(C);
     IntegerType *Int32Ty = IntegerType::getInt32Ty(C);
@@ -203,8 +190,7 @@ struct InsTrim : public ModulePass {
 
           if (instFilename.str().empty()) {
 
-            /* If the original location is empty, try using the inlined location
-             */
+            /* If the original location is empty, try using the inlined location */
             DILocation *oDILoc = cDILoc->getInlinedAt();
             if (oDILoc) {
 
@@ -432,28 +418,19 @@ struct InsTrim : public ModulePass {
         IRB.CreateStore(Incr, MapPtrIdx)
             ->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
 
-        /* Set prev_loc to cur_loc >> 1 */
-        /*
-        StoreInst *Store = IRB.CreateStore(ConstantInt::get(Int32Ty, L >> 1),
-        OldPrev); Store->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C,
-        None));
-        */
-
         total_instr++;
 
       }
 
     }
 
-    OKF("Instrumented %u locations (%llu, %llu) (%s mode)\n" /*", ratio
-                                                                %u%%)."*/
-        ,
+    OKF("Instrumented %u locations (%llu, %llu) (%s mode)\n",
         total_instr, total_rs, total_hs,
         getenv("AFL_HARDEN")
             ? "hardened"
             : ((getenv("AFL_USE_ASAN") || getenv("AFL_USE_MSAN"))
                    ? "ASAN/MSAN"
-                   : "non-hardened") /*, inst_ratio*/);
+                   : "non-hardened"));
     return false;
 
   }
diff --git a/llvm_mode/MarkNodes.cc b/llvm_mode/MarkNodes.cc
index 2aeeda8d..caa8cede 100644
--- a/llvm_mode/MarkNodes.cc
+++ b/llvm_mode/MarkNodes.cc
@@ -65,16 +65,11 @@ void buildCFG(Function *F) {
 
   }
 
-  // uint32_t FakeID = 0;
   for (auto S = F->begin(), E = F->end(); S != E; ++S) {
 
     BasicBlock *BB = &*S;
     uint32_t    MyID = LMap[BB];
-    // if (succ_begin(BB) == succ_end(BB)) {
 
-    // Succs[MyID].push_back(FakeID);
-    // Marked.insert(MyID);
-    //}
     for (auto I = succ_begin(BB), E = succ_end(BB); I != E; ++I) {
 
       Succs[MyID].push_back(LMap[*I]);
@@ -113,7 +108,7 @@ void DFStree(size_t now_id) {
 
 }
 
-void turnCFGintoDAG(Function *F) {
+void turnCFGintoDAG() {
 
   tSuccs = Succs;
   tag.resize(Blocks.size());
@@ -176,7 +171,7 @@ void DFS(uint32_t now) {
 
 }
 
-void DominatorTree(Function *F) {
+void DominatorTree() {
 
   if (Blocks.empty()) return;
   uint32_t s = start_point;
@@ -390,7 +385,7 @@ void MarkSubGraph(uint32_t ss, uint32_t tt) {
 
 }
 
-void MarkVertice(Function *F) {
+void MarkVertice() {
 
   uint32_t s = start_point;
 
@@ -411,8 +406,6 @@ void MarkVertice(Function *F) {
 
   timeStamp = 0;
   uint32_t t = 0;
-  // MarkSubGraph(s, t);
-  // return;
 
   while (s != t) {
 
@@ -432,9 +425,9 @@ std::pair<std::vector<BasicBlock *>, std::vector<BasicBlock *> > markNodes(
   reset();
   labelEachBlock(F);
   buildCFG(F);
-  turnCFGintoDAG(F);
-  DominatorTree::DominatorTree(F);
-  MarkVertice(F);
+  turnCFGintoDAG();
+  DominatorTree::DominatorTree();
+  MarkVertice();
 
   std::vector<BasicBlock *> Result, ResultAbove;
   for (uint32_t x : Markabove) {
diff --git a/llvm_mode/compare-transform-pass.so.cc b/llvm_mode/compare-transform-pass.so.cc
index 0ccce875..5d924b63 100644
--- a/llvm_mode/compare-transform-pass.so.cc
+++ b/llvm_mode/compare-transform-pass.so.cc
@@ -18,7 +18,13 @@
 #include <stdlib.h>
 #include <unistd.h>
 
+#include <list>
+#include <string>
+#include <fstream>
+#include <sys/time.h>
+
 #include "llvm/ADT/Statistic.h"
+#include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Module.h"
@@ -42,6 +48,23 @@ class CompareTransform : public ModulePass {
   static char ID;
   CompareTransform() : ModulePass(ID) {
 
+    char *instWhiteListFilename = getenv("AFL_LLVM_WHITELIST");
+    if (instWhiteListFilename) {
+
+      std::string   line;
+      std::ifstream fileStream;
+      fileStream.open(instWhiteListFilename);
+      if (!fileStream) report_fatal_error("Unable to open AFL_LLVM_WHITELIST");
+      getline(fileStream, line);
+      while (fileStream) {
+
+        myWhitelist.push_back(line);
+        getline(fileStream, line);
+
+      }
+
+    }
+
   }
 
   bool runOnModule(Module &M) override;
@@ -57,6 +80,9 @@ class CompareTransform : public ModulePass {
 
   }
 
+ protected:
+  std::list<std::string> myWhitelist;
+
  private:
   bool transformCmps(Module &M, const bool processStrcmp,
                      const bool processMemcmp, const bool processStrncmp,
@@ -104,6 +130,74 @@ bool CompareTransform::transformCmps(Module &M, const bool processStrcmp,
 
     for (auto &BB : F) {
 
+      if (!myWhitelist.empty()) {
+
+        BasicBlock::iterator IP = BB.getFirstInsertionPt();
+
+        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) {
+
+          DILocation *cDILoc = dyn_cast<DILocation>(Loc.getAsMDNode());
+
+          unsigned int instLine = cDILoc->getLine();
+          StringRef    instFilename = cDILoc->getFilename();
+
+          if (instFilename.str().empty()) {
+
+            /* If the original location is empty, try using the inlined location
+             */
+            DILocation *oDILoc = cDILoc->getInlinedAt();
+            if (oDILoc) {
+
+              instFilename = oDILoc->getFilename();
+              instLine = oDILoc->getLine();
+
+            }
+
+          }
+
+          (void)instLine;
+
+          /* Continue only if we know where we actually are */
+          if (!instFilename.str().empty()) {
+
+            for (std::list<std::string>::iterator it = myWhitelist.begin();
+                 it != myWhitelist.end(); ++it) {
+
+              /* 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()) {
+
+                if (instFilename.str().compare(
+                        instFilename.str().length() - it->length(),
+                        it->length(), *it) == 0) {
+
+                  instrumentBlock = true;
+                  break;
+
+                }
+
+              }
+
+            }
+
+          }
+
+        }
+
+        /* Either we couldn't figure out our location or the location is
+         * not whitelisted, so we skip instrumentation. */
+        if (!instrumentBlock) continue;
+
+      }
+
       for (auto &IN : BB) {
 
         CallInst *callInst = nullptr;
diff --git a/llvm_mode/split-compares-pass.so.cc b/llvm_mode/split-compares-pass.so.cc
index eeac4a55..bc25b322 100644
--- a/llvm_mode/split-compares-pass.so.cc
+++ b/llvm_mode/split-compares-pass.so.cc
@@ -15,7 +15,17 @@
  * limitations under the License.
  */
 
+#include <stdio.h>
+#include <stdlib.h>
+#include <unistd.h>
+
+#include <list>
+#include <string>
+#include <fstream>
+#include <sys/time.h>
+
 #include "llvm/Pass.h"
+#include "llvm/IR/DebugInfo.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/Transforms/IPO/PassManagerBuilder.h"
@@ -35,6 +45,41 @@ class SplitComparesTransform : public ModulePass {
   static char ID;
   SplitComparesTransform() : ModulePass(ID) {
 
+    char *instWhiteListFilename = getenv("AFL_LLVM_WHITELIST");
+    if (instWhiteListFilename) {
+
+      std::string   line;
+      std::ifstream fileStream;
+      fileStream.open(instWhiteListFilename);
+      if (!fileStream) report_fatal_error("Unable to open AFL_LLVM_WHITELIST");
+      getline(fileStream, line);
+      while (fileStream) {
+
+        myWhitelist.push_back(line);
+        getline(fileStream, line);
+
+      }
+
+    }
+
+  }
+
+  static bool isBlacklisted(const Function *F) {
+
+    static const SmallVector<std::string, 5> Blacklist = {
+
+        "asan.", "llvm.", "sancov.", "__ubsan_handle_", "ign."
+
+    };
+
+    for (auto const &BlacklistFunc : Blacklist) {
+
+      if (F->getName().startswith(BlacklistFunc)) { return true; }
+
+    }
+
+    return false;
+
   }
 
   bool runOnModule(Module &M) override;
@@ -49,6 +94,9 @@ class SplitComparesTransform : public ModulePass {
 
   }
 
+ protected:
+  std::list<std::string> myWhitelist;
+
  private:
   int enableFPSplit;
 
@@ -77,8 +125,78 @@ bool SplitComparesTransform::simplifyCompares(Module &M) {
    * all integer comparisons with >= and <= predicates to the icomps vector */
   for (auto &F : M) {
 
+    if (isBlacklisted(&F)) continue;
+
     for (auto &BB : F) {
 
+      if (!myWhitelist.empty()) {
+
+        bool instrumentBlock = false;
+
+        BasicBlock::iterator IP = BB.getFirstInsertionPt();
+
+        /* 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());
+
+          unsigned int instLine = cDILoc->getLine();
+          StringRef    instFilename = cDILoc->getFilename();
+
+          if (instFilename.str().empty()) {
+
+            /* If the original location is empty, try using the inlined location
+             */
+            DILocation *oDILoc = cDILoc->getInlinedAt();
+            if (oDILoc) {
+
+              instFilename = oDILoc->getFilename();
+              instLine = oDILoc->getLine();
+
+            }
+
+          }
+
+          (void)instLine;
+
+          /* Continue only if we know where we actually are */
+          if (!instFilename.str().empty()) {
+
+            for (std::list<std::string>::iterator it = myWhitelist.begin();
+                 it != myWhitelist.end(); ++it) {
+
+              /* 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()) {
+
+                if (instFilename.str().compare(
+                        instFilename.str().length() - it->length(),
+                        it->length(), *it) == 0) {
+
+                  instrumentBlock = true;
+                  break;
+
+                }
+
+              }
+
+            }
+
+          }
+
+        }
+
+        /* Either we couldn't figure out our location or the location is
+         * not whitelisted, so we skip instrumentation. */
+        if (!instrumentBlock) continue;
+
+      }
+
       for (auto &IN : BB) {
 
         CmpInst *selectcmpInst = nullptr;
diff --git a/llvm_mode/split-switches-pass.so.cc b/llvm_mode/split-switches-pass.so.cc
index 2743a71a..3a2838c0 100644
--- a/llvm_mode/split-switches-pass.so.cc
+++ b/llvm_mode/split-switches-pass.so.cc
@@ -18,7 +18,13 @@
 #include <stdlib.h>
 #include <unistd.h>
 
+#include <list>
+#include <string>
+#include <fstream>
+#include <sys/time.h>
+
 #include "llvm/ADT/Statistic.h"
+#include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Module.h"
@@ -42,6 +48,41 @@ class SplitSwitchesTransform : public ModulePass {
   static char ID;
   SplitSwitchesTransform() : ModulePass(ID) {
 
+    char *instWhiteListFilename = getenv("AFL_LLVM_WHITELIST");
+    if (instWhiteListFilename) {
+
+      std::string   line;
+      std::ifstream fileStream;
+      fileStream.open(instWhiteListFilename);
+      if (!fileStream) report_fatal_error("Unable to open AFL_LLVM_WHITELIST");
+      getline(fileStream, line);
+      while (fileStream) {
+
+        myWhitelist.push_back(line);
+        getline(fileStream, line);
+
+      }
+
+    }
+
+  }
+
+  static bool isBlacklisted(const Function *F) {
+
+    static const SmallVector<std::string, 5> Blacklist = {
+
+        "asan.", "llvm.", "sancov.", "__ubsan_handle_", "ign."
+
+    };
+
+    for (auto const &BlacklistFunc : Blacklist) {
+
+      if (F->getName().startswith(BlacklistFunc)) { return true; }
+
+    }
+
+    return false;
+
   }
 
   bool runOnModule(Module &M) override;
@@ -71,6 +112,9 @@ class SplitSwitchesTransform : public ModulePass {
 
   typedef std::vector<CaseExpr> CaseVector;
 
+ protected:
+  std::list<std::string> myWhitelist;
+
  private:
   bool        splitSwitches(Module &M);
   bool        transformCmps(Module &M, const bool processStrcmp,
@@ -268,10 +312,79 @@ bool SplitSwitchesTransform::splitSwitches(Module &M) {
    * all switches to switches vector for later processing */
   for (auto &F : M) {
 
+    if (isBlacklisted(&F)) continue;
+
     for (auto &BB : F) {
 
       SwitchInst *switchInst = nullptr;
 
+      if (!myWhitelist.empty()) {
+
+        bool                 instrumentBlock = false;
+        BasicBlock::iterator IP = BB.getFirstInsertionPt();
+
+        /* 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());
+
+          unsigned int instLine = cDILoc->getLine();
+          StringRef    instFilename = cDILoc->getFilename();
+
+          if (instFilename.str().empty()) {
+
+            /* If the original location is empty, try using the inlined location
+             */
+            DILocation *oDILoc = cDILoc->getInlinedAt();
+            if (oDILoc) {
+
+              instFilename = oDILoc->getFilename();
+              instLine = oDILoc->getLine();
+
+            }
+
+          }
+
+          (void)instLine;
+
+          /* Continue only if we know where we actually are */
+          if (!instFilename.str().empty()) {
+
+            for (std::list<std::string>::iterator it = myWhitelist.begin();
+                 it != myWhitelist.end(); ++it) {
+
+              /* 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()) {
+
+                if (instFilename.str().compare(
+                        instFilename.str().length() - it->length(),
+                        it->length(), *it) == 0) {
+
+                  instrumentBlock = true;
+                  break;
+
+                }
+
+              }
+
+            }
+
+          }
+
+        }
+
+        /* Either we couldn't figure out our location or the location is
+         * not whitelisted, so we skip instrumentation. */
+        if (!instrumentBlock) continue;
+
+      }
+
       if ((switchInst = dyn_cast<SwitchInst>(BB.getTerminator()))) {
 
         if (switchInst->getNumCases() < 1) continue;
diff --git a/test/test.sh b/test/test.sh
index 93a4e008..9676d22d 100755
--- a/test/test.sh
+++ b/test/test.sh
@@ -153,10 +153,10 @@ test "$SYS" = "i686" -o "$SYS" = "x86_64" -o "$SYS" = "amd64" && {
     ../afl-cmin -i in -o in2 -- ./test-instr.plain > /dev/null
     CNT=`ls in2/ | wc -l`
     case "$CNT" in
-1| *1) $ECHO "$GREEN[+] afl-cmin correctly minimized the number of testcases" ;;
-*) $ECHO "$RED[!] afl-cmin did not correctly minimizethe  number of testcases"
-       CODE=1
-       ;;
+      *1) $ECHO "$GREEN[+] afl-cmin correctly minimized the number of testcases" ;;
+      *)  $ECHO "$RED[!] afl-cmin did not correctly minimize the number of testcases"
+          CODE=1
+          ;;
     esac
     ../afl-tmin -i in/in2 -o in2/in2 -- ./test-instr.plain > /dev/null 2>&1
     SIZE=`ls -l in2/in2 2> /dev/null | awk '{print$5}'`
@@ -259,10 +259,10 @@ test -e ../afl-clang-fast -a -e ../split-switches-pass.so && {
       ../afl-cmin -i in -o in2 -- ./test-instr.plain > /dev/null
       CNT=`ls in2/ | wc -l`
       case "$CNT" in
-1| *1) $ECHO "$GREEN[+] afl-cmin correctly minimized the number of testcases" ;;
-*) $ECHO "$RED[!] afl-cmin did not correctly minimize the number of testcases"
-         CODE=1
-         ;;
+        *1) $ECHO "$GREEN[+] afl-cmin correctly minimized the number of testcases" ;;
+        *)  $ECHO "$RED[!] afl-cmin did not correctly minimize the number of testcases"
+            CODE=1
+            ;;
       esac
       ../afl-tmin -i in/in2 -o in2/in2 -- ./test-instr.plain > /dev/null 2>&1
       SIZE=`ls -l in2/in2 2> /dev/null | awk '{print$5}'`