finished dynamorio changes

2 files changed, 54 insertions, 9 deletions

diff --git a/afl-dyninst.cpp b/afl-dyninst.cpp
index 7822fd9..d9db33a 100644
--- a/afl-dyninst.cpp
+++ b/afl-dyninst.cpp
@@ -52,7 +52,7 @@ static const char *USAGE = "-dfvD -i <binary> -o <binary> -l <library> -e <addre
   -m: minimum size of a basic bock to instrument (default: 1)\n \
   -f: try to fix a dyninst bug that leads to crashes\n \
   -S: do not instrument this function (repeat for more than one)\n \
-  -D: instrument fork server and forced exit functions but no basic blocks\n \
+  -D: instrument only a simple fork server and also forced exit functions\n \
   -v: verbose output\n";
 
 bool parseOptions(int argc, char **argv) {
@@ -246,12 +246,14 @@ int main(int argc, char **argv) {
     return EXIT_FAILURE;
   }
   
-  if (DYNINST_MAJOR_VERSION < 9 || (DYNINST_MAJOR_VERSION == 9 && DYNINST_MINOR_VERSION < 3) || (DYNINST_MAJOR_VERSION == 9 && DYNINST_MINOR_VERSION == 3 && DYNINST_PATCH_VERSION <= 2)) {
-    if (dynfix == false)
-      fprintf(stderr, "Warning: your dyninst version does not include a critical fix, you should use the -f option!\n");
-  } else {
-    if (dynfix == true)
-      fprintf(stderr, "Notice: your dyninst version is fixed, the -f option should not be necessary.\n");
+  if (do_bb == true) {
+    if (DYNINST_MAJOR_VERSION < 9 || (DYNINST_MAJOR_VERSION == 9 && DYNINST_MINOR_VERSION < 3) || (DYNINST_MAJOR_VERSION == 9 && DYNINST_MINOR_VERSION == 3 && DYNINST_PATCH_VERSION <= 2)) {
+      if (dynfix == false)
+        fprintf(stderr, "Warning: your dyninst version does not include a critical fix, you should use the -f option!\n");
+    } else {
+      if (dynfix == true)
+        fprintf(stderr, "Notice: your dyninst version is fixed, the -f option should not be necessary.\n");
+    }
   }
 
   BPatch bpatch;
@@ -308,12 +310,17 @@ int main(int argc, char **argv) {
   appImage = appBin->getImage();
 
   /* Find code coverage functions in the instrumentation library */
-  BPatch_function *initAflForkServer = findFuncByName(appImage, (char *) "initAflForkServer");
+  BPatch_function *initAflForkServer;
   save_rdi = findFuncByName(appImage, (char *) "save_rdi");
   restore_rdi = findFuncByName(appImage, (char *) "restore_rdi");
   BPatch_function *bbCallback = findFuncByName(appImage, (char *) "bbCallback");
   BPatch_function *forceCleanExit = findFuncByName(appImage, (char *) "forceCleanExit");
 
+  if (do_bb == true)
+    initAflForkServer = findFuncByName(appImage, (char *) "initAflForkServer");
+  else
+    initAflForkServer = findFuncByName(appImage, (char *) "initOnlyAflForkServer");
+
   if (!initAflForkServer || !bbCallback || !save_rdi || !restore_rdi || !forceCleanExit) {
     cerr << "Instrumentation library lacks callbacks!" << endl;
     return EXIT_FAILURE;
@@ -340,7 +347,7 @@ int main(int argc, char **argv) {
         continue;
     }
     
-    if (do_bb) {
+    if (do_bb == true) {
       cout << "Instrumenting module: " << moduleName << endl;
       vector < BPatch_function * >*allFunctions = (*moduleIter)->getProcedures();
       vector < BPatch_function * >::iterator funcIter;
diff --git a/libAflDyninst.cpp b/libAflDyninst.cpp
index 43d5d78..9b48102 100644
--- a/libAflDyninst.cpp
+++ b/libAflDyninst.cpp
@@ -93,3 +93,41 @@ void save_rdi() {
 void restore_rdi() {
   rdi = saved_di;
 }
+
+
+void initOnlyAflForkServer() {
+  // enter fork() server thyme!
+  int n = write(FORKSRV_FD + 1, &__afl_temp_data, 4);
+
+  if (n != 4) {
+    PRINT_ERROR("Error writting fork server\n");
+    return;
+  }
+  while (1) {
+    n = read(FORKSRV_FD, &__afl_temp_data, 4);
+    if (n != 4) {
+      PRINT_ERROR("Error reading fork server\n");
+      return;
+    }
+
+    __afl_fork_pid = fork();
+    if (__afl_fork_pid < 0) {
+      PRINT_ERROR("Error on fork()\n");
+      return;
+    }
+    if (__afl_fork_pid == 0) {
+      close(FORKSRV_FD);
+      close(FORKSRV_FD + 1);
+      break;
+    } else {
+      // parrent stuff
+      n = write(FORKSRV_FD + 1, &__afl_fork_pid, 4);
+      pid_t temp_pid = waitpid(__afl_fork_pid, &__afl_temp_data, 2);
+
+      if (temp_pid == 0) {
+        return;
+      }
+      n = write(FORKSRV_FD + 1, &__afl_temp_data, 4);
+    }
+  }
+}