llvm_mode whitelist (partial instrumentation) support added

1 files changed, 69 insertions, 1 deletions

diff --git a/llvm_mode/afl-llvm-pass.so.cc b/llvm_mode/afl-llvm-pass.so.cc
index 15b3764a..d46db7c0 100644
--- a/llvm_mode/afl-llvm-pass.so.cc
+++ b/llvm_mode/afl-llvm-pass.so.cc
@@ -31,6 +31,11 @@
 #include <stdlib.h>
 #include <unistd.h>
 
+#include <list>
+#include <string>
+#include <fstream>
+
+#include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LegacyPassManager.h"
@@ -48,7 +53,21 @@ namespace {
     public:
 
       static char ID;
-      AFLCoverage() : ModulePass(ID) { }
+      AFLCoverage() : 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;
 
@@ -56,6 +75,10 @@ namespace {
       //  return "American Fuzzy Lop Instrumentation";
       // }
 
+    protected:
+
+      std::list<std::string> myWhitelist;
+
   };
 
 }
@@ -115,6 +138,51 @@ bool AFLCoverage::runOnModule(Module &M) {
 
       BasicBlock::iterator IP = BB.getFirstInsertionPt();
       IRBuilder<> IRB(&(*IP));
+      
+      if (!myWhitelist.empty()) {
+          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();
+                  }
+              }
+
+              /* 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 (AFL_R(100) >= inst_ratio) continue;