1 files changed, 148 insertions, 2 deletions

diff --git a/llvm_mode/compare-transform-pass.so.cc b/llvm_mode/compare-transform-pass.so.cc
index 0ccce875..e1332a9d 100644
--- a/llvm_mode/compare-transform-pass.so.cc
+++ b/llvm_mode/compare-transform-pass.so.cc
@@ -18,6 +18,12 @@
 #include <stdlib.h>
 #include <unistd.h>
 
+#include <list>
+#include <string>
+#include <fstream>
+#include <sys/time.h>
+#include "llvm/Config/llvm-config.h"
+
 #include "llvm/ADT/Statistic.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LegacyPassManager.h"
@@ -26,10 +32,19 @@
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/IPO/PassManagerBuilder.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
-#include "llvm/IR/Verifier.h"
 #include "llvm/Pass.h"
 #include "llvm/Analysis/ValueTracking.h"
 
+#if LLVM_VERSION_MAJOR > 3 || \
+    (LLVM_VERSION_MAJOR == 3 && LLVM_VERSION_MINOR > 4)
+#include "llvm/IR/Verifier.h"
+#include "llvm/IR/DebugInfo.h"
+#else
+#include "llvm/Analysis/Verifier.h"
+#include "llvm/DebugInfo.h"
+#define nullptr 0
+#endif
+
 #include <set>
 
 using namespace llvm;
@@ -42,6 +57,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 +89,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,
@@ -89,7 +124,7 @@ bool CompareTransform::transformCmps(Module &M, const bool processStrcmp,
       c = M.getOrInsertFunction("tolower", Int32Ty, Int32Ty
 #if LLVM_VERSION_MAJOR < 5
                                 ,
-                                nullptr
+                                NULL
 #endif
       );
 #if LLVM_VERSION_MAJOR < 9
@@ -104,6 +139,117 @@ 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 LLVM_VERSION_MAJOR >= 4 || \
+    (LLVM_VERSION_MAJOR == 3 && LLVM_VERSION_MINOR >= 7)
+        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;
+
+                }
+
+              }
+
+            }
+
+          }
+
+        }
+
+#else
+        if (!Loc.isUnknown()) {
+
+          DILocation cDILoc(Loc.getAsMDNode(C));
+
+          unsigned int instLine = cDILoc.getLineNumber();
+          StringRef    instFilename = cDILoc.getFilename();
+
+          (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;
+
+                }
+
+              }
+
+            }
+
+          }
+
+        }
+
+#endif
+
+        /* 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;