-l fix (first iteration)

4 files changed, 45 insertions, 6 deletions

diff --git a/CHANGES b/CHANGES
index 149b3d7..5f6d4b1 100644
--- a/CHANGES
+++ b/CHANGES
@@ -2,6 +2,8 @@ Changelog
 =========
 
 https://github.com/vanhauser-thc/afl-dyninst
+ - First fix for -l option, did copy and instrument all libs
+ - Only compile dyninst9 bug workaround when necessary
  - added -I option (only instrument specific functions)
  - updated the README for guidance to build against dyninst version 10
  - added support for dyninst version 10
diff --git a/Makefile b/Makefile
index ca093f7..d029936 100644
--- a/Makefile
+++ b/Makefile
@@ -38,7 +38,7 @@ afl-dyninst:	afl-dyninst.o
 		-ldyninstAPI
 
 libAflDyninst.so: libAflDyninst.cpp
-	$(CXX) -O3 -std=c++11 $(LIBFLAGS) -I$(AFL_ROOT) -I$(DEPS_ROOT)/include libAflDyninst.cpp -o libAflDyninst.so
+	$(CXX) -O3 -std=c++11 $(LIBFLAGS) -I$(AFL_ROOT) -I$(DYNINST_ROOT)/include -I$(DEPS_ROOT)/include libAflDyninst.cpp -o libAflDyninst.so
 
 afl-dyninst.o: afl-dyninst.cpp
 	$(CXX) $(CXXFLAGS) $(DYNINST_OPT) -I$(DEPS_ROOT)/include -I$(DYNINST_ROOT)/include  -c afl-dyninst.cpp
diff --git a/afl-dyninst.cpp b/afl-dyninst.cpp
index d736ad7..01c5c5e 100644
--- a/afl-dyninst.cpp
+++ b/afl-dyninst.cpp
@@ -320,15 +320,21 @@ bool insertBBCallback(BPatch_addressSpace *appBin, BPatch_function *curFunc, cha
       BPatch_constExpr bbId(randID);
 
       instArgs.push_back(&bbId);
+#if (DYNINST_MAJOR_VERSION < 10)
       BPatch_funcCallExpr instIncExpr1(*save_rdi, instArgs1);
       BPatch_funcCallExpr instIncExpr3(*restore_rdi, instArgs1);
+#endif
       BPatch_funcCallExpr instIncExpr(*instBBIncFunc, instArgs);
 
+#if (DYNINST_MAJOR_VERSION < 10)
       if (dynfix == true)
         handle = appBin->insertSnippet(instIncExpr1, *bbEntry, BPatch_callBefore, BPatch_firstSnippet);
+#endif
       handle = appBin->insertSnippet(instIncExpr, *bbEntry, BPatch_callBefore);
+#if (DYNINST_MAJOR_VERSION < 10)
       if (dynfix == true)
         handle = appBin->insertSnippet(instIncExpr3, *bbEntry, BPatch_callBefore, BPatch_lastSnippet);
+#endif
     }
 
     if (!handle) {
@@ -445,8 +451,10 @@ int main(int argc, char **argv) {
   /* Find code coverage functions in the instrumentation library */
   BPatch_function *initAflForkServer;
 
+#if (DYNINST_MAJOR_VERSION < 10)
   save_rdi = findFuncByName(appImage, (char *)"save_rdi");
   restore_rdi = findFuncByName(appImage, (char *)"restore_rdi");
+#endif
   BPatch_function *bbCallback = findFuncByName(appImage, (char *)"bbCallback");
   BPatch_function *forceCleanExit = findFuncByName(appImage, (char *)"forceCleanExit");
 
@@ -458,7 +466,11 @@ int main(int argc, char **argv) {
   } else
     initAflForkServer = findFuncByName(appImage, (char *)"initOnlyAflForkServer");
 
-  if (!initAflForkServer || !bbCallback || !save_rdi || !restore_rdi || !forceCleanExit) {
+  if (!initAflForkServer || !bbCallback || !forceCleanExit
+#if (DYNINST_MAJOR_VERSION < 10)
+      || !save_rdi || !restore_rdi
+#endif
+  ) {
     cerr << "Instrumentation library lacks callbacks!" << endl;
     return EXIT_FAILURE;
   }
@@ -467,7 +479,6 @@ int main(int argc, char **argv) {
 
   // if an entrypoint was set then find function, else find _init
   BPatch_function *funcToPatch = NULL;
-
   if (entryPoint == 0 && entryPointName == NULL) {
     if (func2patch == NULL) {
       cerr << "Couldn't locate _init, specify entry point manually with -e 0xaddr" << endl;
@@ -536,20 +547,39 @@ int main(int argc, char **argv) {
     return EXIT_FAILURE;
   }
 
+  bool skip_until_next_library = false;
+
   for (moduleIter = modules->begin(); moduleIter != modules->end(); ++moduleIter) {
     char moduleName[1024];
 
     (*moduleIter)->getName(moduleName, 1024);
-    if ((*moduleIter)->isSharedLib()) {
-      if (instrumentLibraries.find(moduleName) == instrumentLibraries.end() && string(moduleName).find(".so") != string::npos) {
+    if ((*moduleIter)->isSharedLib() && (strstr(moduleName, ".so.") != NULL || (strlen(moduleName) > 3 && strncmp(moduleName + strlen(moduleName) - 3, ".so", 3) == 0))) {
+      bool skip_this_lib = true;
+      for (std::set<std::string>::iterator libIter = instrumentLibraries.begin(); libIter != instrumentLibraries.end(); ++libIter)
+        if (strncmp(libIter->c_str(), moduleName, strlen(libIter->c_str())) == 0)
+          skip_this_lib = false;
+      if (skip_this_lib == true) {
+        skip_until_next_library = true;
         cout << "Skipping library: " << moduleName << endl;
         continue;
+      } else {
+        skip_until_next_library = false;
       }
     }
 
     if (string(moduleName).find(defaultModuleName) != string::npos) {
-      if (skipMainModule)
+      if (skipMainModule) {
+        skip_until_next_library = true;
         continue;
+      } else {
+        skip_until_next_library = false;
+      }
+    }
+
+    if (skip_until_next_library == true) {
+      if (verbose)
+        cout << "Skipping " << moduleName << " because skip_until_next_library is active" << endl;
+      continue;
     }
 
     if (do_bb == true) {
diff --git a/libAflDyninst.cpp b/libAflDyninst.cpp
index e7c2a0f..af2d9f4 100644
--- a/libAflDyninst.cpp
+++ b/libAflDyninst.cpp
@@ -1,4 +1,5 @@
 #include "config.h"
+#include "dyninstversion.h" // if this include errors, compile and install https://github.com/dyninst/dyninst
 #include <algorithm>
 #include <cstdio>
 #include <cstdlib>
@@ -23,8 +24,10 @@ static unsigned short int prev_id = 0;
 static bool forkserver_installed = false;
 #if (__amd64__ || __x86_64__)
 static long saved_di;
+#if (DYNINST_MAJOR_VERSION < 10)
 register long rdi asm("di"); // the warning is fine - we need the warning because of a bug in dyninst9
 #endif
+#endif
 
 #define PRINT_ERROR(string) (void)(write(2, string, strlen(string)) + 1) // the (...+1) weirdness is so we do not get an ignoring return value warning
 
@@ -91,17 +94,21 @@ void bbCallback(unsigned short id) {
 
 void forceCleanExit() { exit(0); }
 
+#if (DYNINST_MAJOR_VERSION < 10)
 void save_rdi() {
 #if __amd64__ || __x86_64__
   saved_di = rdi;
 #endif
 }
+#endif
 
+#if (DYNINST_MAJOR_VERSION < 10)
 void restore_rdi() {
 #if __amd64__ || __x86_64__
   rdi = saved_di;
 #endif
 }
+#endif
 
 void initOnlyAflForkServer() {
   if (forkserver_installed == true)