afl-tmin is fixed via default initialization in forkserver

3 files changed, 43 insertions, 71 deletions

diff --git a/llvm_mode/LLVMInsTrim.so.cc b/llvm_mode/LLVMInsTrim.so.cc
index 8127c893..0fddad24 100644
--- a/llvm_mode/LLVMInsTrim.so.cc
+++ b/llvm_mode/LLVMInsTrim.so.cc
@@ -54,6 +54,7 @@ struct InsTrim : public ModulePass {
 
  protected:
   std::list<std::string> myWhitelist;
+  uint32_t function_minimum_size = 1;
 
  private:
   std::mt19937 generator;
@@ -152,6 +153,9 @@ struct InsTrim : public ModulePass {
 
     }
 
+    if (getenv("AFL_LLVM_INSTRIM_SKIPSINGLEBLOCK") != NULL)
+      function_minimum_size = 2;
+
     // this is our default
     MarkSetOpt = true;
 
@@ -176,8 +180,8 @@ struct InsTrim : public ModulePass {
 
     for (Function &F : M) {
 
-      // if it is external or only contains one basic block: skip it
-      if (F.size() < 2) { continue; }
+      // if the function below our minimum size skip it (1 or 2)
+      if (F.size() < function_minimum_size) { continue; }
 
       if (!myWhitelist.empty()) {
 
@@ -383,67 +387,15 @@ struct InsTrim : public ModulePass {
 
         }
 
-        // Bugfix #1: remove single block function instrumentation
-
-        for (BasicBlock &BB : F) {
-
-          if (MarkSetOpt && MS.find(&BB) == MS.end()) {
-
-            // Bugfix #2: instrument blocks that should be but InsTrim
-            //            doesn't due to an algorithmic bug
-            int more_than_one = -1;
-
-            for (pred_iterator PI = pred_begin(&BB), E = pred_end(&BB); PI != E;
-                 ++PI) {
-
-              BasicBlock *Pred = *PI;
-              int         count = 0;
-
-              if (more_than_one == -1) more_than_one = 0;
-              for (succ_iterator SI = succ_begin(Pred), E = succ_end(Pred);
-                   SI != E; ++SI) {
-
-                BasicBlock *Succ = *SI;
-                if (Succ != NULL) count++;
-
-              }
-
-              if (count > 1) more_than_one = 1;
-
+	if (function_minimum_size < 2) {
+          for (BasicBlock &BB : F) {
+            if (MS.find(&BB) == MS.end()) {
+              continue;
             }
-
-            if (more_than_one != 1) continue;
-            for (succ_iterator SI = succ_begin(&BB), E = succ_end(&BB); SI != E;
-                 ++SI) {
-
-              BasicBlock *Succ = *SI;
-              if (Succ != NULL && MS.find(Succ) == MS.end()) {
-
-                int cnt = 0;
-                for (succ_iterator SI2 = succ_begin(Succ), E2 = succ_end(Succ);
-                     SI2 != E2; ++SI2) {
-
-                  BasicBlock *Succ2 = *SI2;
-                  if (Succ2 != NULL) cnt++;
-
-                }
-
-                if (cnt == 0) {
-
-                  // fprintf(stderr, "INSERT!\n");
-                  MS.insert(Succ);
-                  total_rs += 1;
-
-                }
-
-              }
-
-            }
-
+            IRBuilder<> IRB(&*BB.getFirstInsertionPt());
+            IRB.CreateStore(ConstantInt::get(Int32Ty, genLabel()), OldPrev);
           }
-
         }
-
       }
 
       for (BasicBlock &BB : F) {
diff --git a/llvm_mode/MarkNodes.cc b/llvm_mode/MarkNodes.cc
index 7b22bac0..b8587826 100644
--- a/llvm_mode/MarkNodes.cc
+++ b/llvm_mode/MarkNodes.cc
@@ -376,10 +376,10 @@ void MakeUniq(uint32_t now) {
 
 }
 
-void MarkSubGraph(uint32_t ss, uint32_t tt) {
+bool MarkSubGraph(uint32_t ss, uint32_t tt) {
 
   TopologicalSort(ss, tt);
-  if (TopoOrder.empty()) return;
+  if (TopoOrder.empty()) return false;
 
   for (uint32_t i : TopoOrder) {
 
@@ -393,6 +393,11 @@ void MarkSubGraph(uint32_t ss, uint32_t tt) {
     MakeUniq(TopoOrder[i]);
 
   }
+  
+   // Check if there is an empty path.
+  if (NextMarked[tt].count(TopoOrder[0]) > 0)
+    return true;
+  return false;
 
 }
 
@@ -417,13 +422,19 @@ void MarkVertice() {
 
   timeStamp = 0;
   uint32_t t = 0;
+  bool emptyPathExists = true;
 
   while (s != t) {
 
-    MarkSubGraph(DominatorTree::idom[t], t);
+    emptyPathExists &= MarkSubGraph(DominatorTree::idom[t], t);
     t = DominatorTree::idom[t];
 
   }
+  
+  if (emptyPathExists) {
+    // Mark all exit blocks to catch the empty path.
+    Marked.insert(t_Pred[0].begin(), t_Pred[0].end());
+  }
 
 }
 
diff --git a/src/afl-forkserver.c b/src/afl-forkserver.c
index 60c84cde..087f89ac 100644
--- a/src/afl-forkserver.c
+++ b/src/afl-forkserver.c
@@ -135,15 +135,24 @@ void handle_timeout(int sig) {
 
 void afl_fsrv_init(afl_forkserver_t *fsrv) {
 
-  uint32_t i, j = 0;
-
-  // this is the default and is != 0 so we need to set it if fsrv is still
-  // uninitialized
-  for (i = 0; i < sizeof(afl_forkserver_t) && j == 0; i++)
-    if (((char*)fsrv)[i] != 0)
-      j = 1;
-  if (j == 0)
+  // this structure needs default so we initialize it if this was not done already
+
+  if (!fsrv->use_stdin) {
+
     fsrv->use_stdin = 1;
+    fsrv->out_fd = -1;
+    fsrv->out_dir_fd = -1;
+    fsrv->dev_null_fd = -1;
+#ifndef HAVE_ARC4RANDOM
+    fsrv->dev_urandom_fd = -1;
+#endif
+    fsrv->exec_tmout = EXEC_TIMEOUT;
+    fsrv->mem_limit = MEM_LIMIT;
+    fsrv->child_pid = -1;
+    fsrv->out_dir_fd = -1;
+    
+  }
+  
   list_append(&fsrv_list, fsrv);
 
 }