3 files changed, 73 insertions, 31 deletions
diff --git a/llvm_mode/LLVMInsTrim.so.cc b/llvm_mode/LLVMInsTrim.so.cc
index 08d3f68f..390e0697 100644
--- a/llvm_mode/LLVMInsTrim.so.cc
+++ b/llvm_mode/LLVMInsTrim.so.cc
@@ -169,6 +169,7 @@ struct InsTrim : public ModulePass {
 
     ConstantInt *Zero = ConstantInt::get(Int8Ty, 0);
     ConstantInt *One = ConstantInt::get(Int8Ty, 1);
+    ConstantInt *One32 = ConstantInt::get(Int32Ty, 1);
 
     u64 total_rs = 0;
     u64 total_hs = 0;
@@ -382,19 +383,64 @@ struct InsTrim : public ModulePass {
 
         }
 
-        auto *EBB = &F.getEntryBlock();
-        if (succ_begin(EBB) == succ_end(EBB)) {
+        // Bugfix #1: remove single block function instrumentation
 
-          MS.insert(EBB);
-          total_rs += 1;
+        for (BasicBlock &BB : F) {
 
-        }
+          if (MarkSetOpt && MS.find(&BB) == MS.end()) {
 
-        for (BasicBlock &BB : F) {
+            // 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 (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 (MS.find(&BB) == MS.end()) { continue; }
-          IRBuilder<> IRB(&*BB.getFirstInsertionPt());
-          IRB.CreateStore(ConstantInt::get(Int32Ty, genLabel()), OldPrev);
+                if (cnt == 0) {
+
+                  // fprintf(stderr, "INSERT!\n");
+                  MS.insert(Succ);
+                  total_rs += 1;
+
+                }
+
+              }
+
+            }
+
+          }
 
         }
 
@@ -402,33 +448,24 @@ struct InsTrim : public ModulePass {
 
       for (BasicBlock &BB : F) {
 
-        auto PI = pred_begin(&BB);
-        auto PE = pred_end(&BB);
         if (MarkSetOpt && MS.find(&BB) == MS.end()) { continue; }
 
         IRBuilder<> IRB(&*BB.getFirstInsertionPt());
         Value *     L = NULL;
-        if (PI == PE) {
 
-          L = ConstantInt::get(Int32Ty, genLabel());
+        auto *PN = PHINode::Create(Int32Ty, 0, "", &*BB.begin());
+        DenseMap<BasicBlock *, unsigned> PredMap;
+        for (auto PI = pred_begin(&BB), PE = pred_end(&BB); PI != PE; ++PI) {
 
-        } else {
-
-          auto *PN = PHINode::Create(Int32Ty, 0, "", &*BB.begin());
-          DenseMap<BasicBlock *, unsigned> PredMap;
-          for (auto PI = pred_begin(&BB), PE = pred_end(&BB); PI != PE; ++PI) {
-
-            BasicBlock *PBB = *PI;
-            auto        It = PredMap.insert({PBB, genLabel()});
-            unsigned    Label = It.first->second;
-            PN->addIncoming(ConstantInt::get(Int32Ty, Label), PBB);
-
-          }
-
-          L = PN;
+          BasicBlock *PBB = *PI;
+          auto        It = PredMap.insert({PBB, genLabel()});
+          unsigned    Label = It.first->second;
+          PN->addIncoming(ConstantInt::get(Int32Ty, Label), PBB);
 
         }
 
+        L = PN;
+
         /* Load prev_loc */
         LoadInst *PrevLoc = IRB.CreateLoad(OldPrev);
         PrevLoc->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
@@ -475,6 +512,11 @@ struct InsTrim : public ModulePass {
         IRB.CreateStore(Incr, MapPtrIdx)
             ->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
 
+        // Bugfix #3: save the actually location ID to OldPrev
+        Value *Shr = IRB.CreateLShr(L, One32);
+        IRB.CreateStore(Shr, OldPrev)
+            ->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
+
         total_instr++;
 
       }
diff --git a/llvm_mode/README.md b/llvm_mode/README.md
index ee6e51b5..e6c47c9c 100644
--- a/llvm_mode/README.md
+++ b/llvm_mode/README.md
@@ -2,7 +2,7 @@
 
   (See [../README](../README.md) for the general instruction manual.)
 
-  (See [../gcc_plugin/README.gcc](../gcc_plugin/README.gcc.md) for the GCC-based instrumentation.)
+  (See [../gcc_plugin/README](../gcc_plugin/README.md) for the GCC-based instrumentation.)
 
 ## 1) Introduction
 
diff --git a/llvm_mode/afl-llvm-rt.o.c b/llvm_mode/afl-llvm-rt.o.c
index 53852320..b3561cb2 100644
--- a/llvm_mode/afl-llvm-rt.o.c
+++ b/llvm_mode/afl-llvm-rt.o.c
@@ -129,7 +129,7 @@ static void __afl_map_shm(void) {
     __afl_area_ptr[0] = 1;
 
   }
-  
+
   id_str = getenv(CMPLOG_SHM_ENV_VAR);
 
   if (id_str) {
@@ -260,7 +260,7 @@ static void __afl_start_forkserver(void) {
 
 }
 
-/* A simplified persistent mode handler, used as explained in README.llvm. */
+/* A simplified persistent mode handler, used as explained in llvm_mode/README.md. */
 
 int __afl_persistent_loop(unsigned int max_cnt) {
 
@@ -346,7 +346,7 @@ __attribute__((constructor(CONST_PRIO))) void __afl_auto_init(void) {
 
 /* The following stuff deals with supporting -fsanitize-coverage=trace-pc-guard.
    It remains non-operational in the traditional, plugin-backed LLVM mode.
-   For more info about 'trace-pc-guard', see README.llvm.
+   For more info about 'trace-pc-guard', see llvm_mode/README.md.
 
    The first function (__sanitizer_cov_trace_pc_guard) is called back on every
    edge (as opposed to every basic block). */