initial commit

5 files changed, 571 insertions, 6 deletions

diff --git a/llvm_mode/LLVMInsTrim.so.cc b/llvm_mode/LLVMInsTrim.so.cc
new file mode 100644
index 00000000..2a13981d
--- /dev/null
+++ b/llvm_mode/LLVMInsTrim.so.cc
@@ -0,0 +1,197 @@
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/IR/CFG.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/LegacyPassManager.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/IPO/PassManagerBuilder.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include <unordered_set>
+#include <random>
+
+#include "MarkNodes.h"
+
+using namespace llvm;
+
+static cl::opt<bool> MarkSetOpt("markset", cl::desc("MarkSet"),
+                                cl::init(false));
+static cl::opt<bool> LoopHeadOpt("loophead", cl::desc("LoopHead"),
+                                 cl::init(false));
+
+namespace {
+  struct InsTrim : public ModulePass {
+  private:
+    std::mt19937 generator;
+    int total_instr = 0;
+
+    unsigned genLabel() {
+      return generator() % 65536;
+    }
+
+  public:
+    static char ID;
+    InsTrim() : ModulePass(ID), generator(0) {}
+
+    void getAnalysisUsage(AnalysisUsage &AU) const override {
+      AU.addRequired<DominatorTreeWrapperPass>();
+    }
+
+    StringRef getPassName() const override {
+      return "InstTrim Instrumentation";
+    }
+
+    bool runOnModule(Module &M) override {
+      if (getenv("LOOPHEAD")) {
+        LoopHeadOpt = true;
+        MarkSetOpt = true;
+      } else if (getenv("MARKSET")) {
+        MarkSetOpt = true;
+      }
+
+      LLVMContext &C = M.getContext();
+      IntegerType *Int8Ty  = IntegerType::getInt8Ty(C);
+      IntegerType *Int32Ty = IntegerType::getInt32Ty(C);
+
+      GlobalVariable *CovMapPtr = new GlobalVariable(
+        M, PointerType::getUnqual(Int8Ty), false, GlobalValue::ExternalLinkage,
+        nullptr, "__afl_area_ptr");
+
+      GlobalVariable *OldPrev = new GlobalVariable(
+        M, Int32Ty, false, GlobalValue::ExternalLinkage, 0, "__afl_prev_loc",
+        0, GlobalVariable::GeneralDynamicTLSModel, 0, false);
+
+      unsigned total_rs = 0;
+      unsigned total_hs = 0;
+
+      for (Function &F : M) {
+        if (!F.size()) {
+          continue;
+        }
+
+        std::unordered_set<BasicBlock *> MS;
+        if (!MarkSetOpt) {
+          for (auto &BB : F) {
+            MS.insert(&BB);
+          }
+          total_rs += F.size();
+        } else {
+          auto Result = markNodes(&F);
+          auto RS = Result.first;
+          auto HS = Result.second;
+
+          MS.insert(RS.begin(), RS.end());
+          if (!LoopHeadOpt) {
+            MS.insert(HS.begin(), HS.end());
+            total_rs += MS.size();
+          } else {
+            DenseSet<std::pair<BasicBlock *, BasicBlock *>> EdgeSet;
+            DominatorTreeWrapperPass *DTWP =
+              &getAnalysis<DominatorTreeWrapperPass>(F);
+            auto DT = &DTWP->getDomTree();
+
+            total_rs += RS.size();
+            total_hs += HS.size();
+
+            for (BasicBlock *BB : HS) {
+              bool Inserted = false;
+              for (auto BI = pred_begin(BB), BE = pred_end(BB);
+                   BI != BE; ++BI
+              ) {
+                auto Edge = BasicBlockEdge(*BI, BB);
+                if (Edge.isSingleEdge() && DT->dominates(Edge, BB)) {
+                  EdgeSet.insert({*BI, BB});
+                  Inserted = true;
+                  break;
+                }
+              }
+              if (!Inserted) {
+                MS.insert(BB);
+                total_rs += 1;
+                total_hs -= 1;
+              }
+            }
+            for (auto I = EdgeSet.begin(), E = EdgeSet.end(); I != E; ++I) {
+              auto PredBB = I->first;
+              auto SuccBB = I->second;
+              auto NewBB = SplitBlockPredecessors(SuccBB, {PredBB}, ".split",
+                                                  DT, nullptr, false);
+              MS.insert(NewBB);
+            }
+          }
+
+          auto *EBB = &F.getEntryBlock();
+          if (succ_begin(EBB) == succ_end(EBB)) {
+            MS.insert(EBB);
+            total_rs += 1;
+          }
+
+          for (BasicBlock &BB : F) {
+            if (MS.find(&BB) == MS.end()) {
+              continue;
+            }
+            IRBuilder<> IRB(&*BB.getFirstInsertionPt());
+            IRB.CreateStore(ConstantInt::get(Int32Ty, genLabel()), OldPrev);
+          }
+        }
+
+        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());
+          } 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;
+          }
+
+          LoadInst *PrevLoc = IRB.CreateLoad(OldPrev);
+          Value *PrevLocCasted = IRB.CreateZExt(PrevLoc, IRB.getInt32Ty());
+
+          LoadInst *MapPtr = IRB.CreateLoad(CovMapPtr);
+          Value *MapPtrIdx = IRB.CreateGEP(MapPtr,
+                                           IRB.CreateXor(PrevLocCasted, L));
+
+          LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
+          Value *Incr = IRB.CreateAdd(Counter, ConstantInt::get(Int8Ty, 1));
+          IRB.CreateStore(Incr, MapPtrIdx);
+          total_instr++;
+        }
+      }
+
+      errs() << total_instr << " locations instrumented ("<< total_rs << "," << total_hs << ")\n";
+      return false;
+    }
+  }; // end of struct InsTrim
+}  // end of anonymous namespace
+
+char InsTrim::ID = 0;
+
+static void registerAFLPass(const PassManagerBuilder &,
+                            legacy::PassManagerBase &PM) {
+  PM.add(new InsTrim());
+}
+
+static RegisterStandardPasses RegisterAFLPass(
+    PassManagerBuilder::EP_OptimizerLast, registerAFLPass);
+
+static RegisterStandardPasses RegisterAFLPass0(
+    PassManagerBuilder::EP_EnabledOnOptLevel0, registerAFLPass);
diff --git a/llvm_mode/Makefile b/llvm_mode/Makefile
index b6ab0c61..a66f18ab 100644
--- a/llvm_mode/Makefile
+++ b/llvm_mode/Makefile
@@ -67,7 +67,7 @@ ifeq "$(origin CC)" "default"
 endif
 
 ifndef AFL_TRACE_PC
-  PROGS      = ../afl-clang-fast ../afl-llvm-pass.so ../afl-llvm-rt.o ../afl-llvm-rt-32.o ../afl-llvm-rt-64.o ../compare-transform-pass.so ../split-compares-pass.so ../split-switches-pass.so
+  PROGS      = ../afl-clang-fast ../libLLVMInsTrim.so ../afl-llvm-rt.o ../afl-llvm-rt-32.o ../afl-llvm-rt-64.o ../compare-transform-pass.so ../split-compares-pass.so ../split-switches-pass.so
 else
   PROGS      = ../afl-clang-fast ../afl-llvm-rt.o ../afl-llvm-rt-32.o ../afl-llvm-rt-64.o ../compare-transform-pass.so ../split-compares-pass.so ../split-switches-pass.so
 endif
@@ -91,8 +91,8 @@ endif
 	$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
 	ln -sf afl-clang-fast ../afl-clang-fast++
 
-../afl-llvm-pass.so: afl-llvm-pass.so.cc | test_deps
-	$(CXX) $(CLANG_CFL) -shared $< -o $@ $(CLANG_LFL)
+../libLLVMInsTrim.so: LLVMInsTrim.so.cc MarkNodes.cc | test_deps
+	$(CXX) $(CLANG_CFL) -DLLVMInsTrim_EXPORTS -fno-rtti -fPIC -std=gnu++11 -shared $< -o $@ $(CLANG_LFL)
 
 # laf
 ../split-switches-pass.so:	split-switches-pass.so.cc | test_deps
diff --git a/llvm_mode/MarkNodes.cc b/llvm_mode/MarkNodes.cc
new file mode 100644
index 00000000..3c2129ef
--- /dev/null
+++ b/llvm_mode/MarkNodes.cc
@@ -0,0 +1,355 @@
+#include <algorithm>
+#include <map>
+#include <queue>
+#include <set>
+#include <vector>
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/CFG.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace llvm;
+
+DenseMap<BasicBlock *, uint32_t> LMap;
+std::vector<BasicBlock *> Blocks;
+std::set<uint32_t> Marked , Markabove;
+std::vector< std::vector<uint32_t> > Succs , Preds;
+
+void reset(){
+  LMap.clear();
+  Blocks.clear();
+  Marked.clear();
+  Markabove.clear();
+}
+
+uint32_t start_point;
+
+void labelEachBlock(Function *F) {
+  // Fake single endpoint;
+  LMap[NULL] = Blocks.size();
+  Blocks.push_back(NULL);
+ 
+  // Assign the unique LabelID to each block;
+  for (auto I = F->begin(), E = F->end(); I != E; ++I) {
+    BasicBlock *BB = &*I;
+    LMap[BB] = Blocks.size();
+    Blocks.push_back(BB);
+  }
+  
+  start_point = LMap[&F->getEntryBlock()];
+}
+
+void buildCFG(Function *F) {
+  Succs.resize( Blocks.size() );
+  Preds.resize( Blocks.size() );
+  for( size_t i = 0 ; i < Succs.size() ; i ++ ){
+    Succs[ i ].clear();
+    Preds[ i ].clear();
+  }
+
+  uint32_t FakeID = 0;
+  for (auto S = F->begin(), E = F->end(); S != E; ++S) {
+    BasicBlock *BB = &*S;
+    uint32_t MyID = LMap[BB];
+    //if (succ_begin(BB) == succ_end(BB)) {
+      //Succs[MyID].push_back(FakeID);
+      //Marked.insert(MyID);
+    //}
+    for (auto I = succ_begin(BB), E = succ_end(BB); I != E; ++I) {
+      Succs[MyID].push_back(LMap[*I]);
+    }
+  }
+}
+
+std::vector< std::vector<uint32_t> > tSuccs;
+std::vector<bool> tag , indfs;
+
+void DFStree(size_t now_id) {
+  if(tag[now_id]) return;
+  tag[now_id]=true;
+  indfs[now_id]=true;
+  for (auto succ: tSuccs[now_id]) {
+    if(tag[succ] and indfs[succ]) {
+      Marked.insert(succ);
+      Markabove.insert(succ);
+      continue;
+    }
+    Succs[now_id].push_back(succ);
+    Preds[succ].push_back(now_id);
+    DFStree(succ);
+  }
+  indfs[now_id]=false;
+}
+void turnCFGintoDAG(Function *F) {
+  tSuccs = Succs;
+  tag.resize(Blocks.size());
+  indfs.resize(Blocks.size());
+  for (size_t i = 0; i < Blocks.size(); ++ i) {
+    Succs[i].clear();
+    tag[i]=false;
+    indfs[i]=false;
+  }
+  DFStree(start_point);
+  for (size_t i = 0; i < Blocks.size(); ++ i) 
+    if( Succs[i].empty() ){
+      Succs[i].push_back(0);
+      Preds[0].push_back(i);
+    }
+}
+
+uint32_t timeStamp;
+namespace DominatorTree{
+  std::vector< std::vector<uint32_t> > cov;
+  std::vector<uint32_t> dfn, nfd, par, sdom, idom, mom, mn;
+
+  bool Compare(uint32_t u, uint32_t v) {
+    return dfn[u] < dfn[v];
+  }
+  uint32_t eval(uint32_t u) {
+    if( mom[u] == u ) return u;
+    uint32_t res = eval( mom[u] );
+    if(Compare(sdom[mn[mom[u]]] , sdom[mn[u]])) {
+      mn[u] = mn[mom[u]];
+    }
+    return mom[u] = res;
+  }
+
+  void DFS(uint32_t now) {
+    timeStamp += 1;
+    dfn[now] = timeStamp;
+    nfd[timeStamp - 1] = now;
+    for( auto succ : Succs[now] ) {
+      if( dfn[succ] == 0 ) {
+        par[succ] = now;
+        DFS(succ);
+      }
+    }
+  }
+
+  void DominatorTree(Function *F) {
+    if( Blocks.empty() ) return;
+    uint32_t s = start_point;
+
+    // Initialization
+    mn.resize(Blocks.size());
+    cov.resize(Blocks.size());
+    dfn.resize(Blocks.size());
+    nfd.resize(Blocks.size());
+    par.resize(Blocks.size());
+    mom.resize(Blocks.size());
+    sdom.resize(Blocks.size());
+    idom.resize(Blocks.size());
+
+    for( uint32_t i = 0 ; i < Blocks.size() ; i ++ ) {
+      dfn[i] = 0;
+      nfd[i] = Blocks.size();
+      cov[i].clear();
+      idom[i] = mom[i] = mn[i] = sdom[i] = i;
+    }
+
+    timeStamp = 0;
+    DFS(s);
+
+    for( uint32_t i = Blocks.size() - 1 ; i >= 1u ; i -- ) {
+      uint32_t now = nfd[i];
+      if( now == Blocks.size() ) {
+        continue;
+      }
+      for( uint32_t pre : Preds[ now ] ) {
+        if( dfn[ pre ] ) {
+          eval(pre);
+          if( Compare(sdom[mn[pre]], sdom[now]) ) {
+            sdom[now] = sdom[mn[pre]];
+          }
+        }
+      }
+      cov[sdom[now]].push_back(now);
+      mom[now] = par[now];
+      for( uint32_t x : cov[par[now]] ) {
+        eval(x);
+        if( Compare(sdom[mn[x]], par[now]) ) {
+          idom[x] = mn[x];
+        } else {
+          idom[x] = par[now];
+        }
+      }
+    }
+
+    for( uint32_t i = 1 ; i < Blocks.size() ; i += 1 ) {
+      uint32_t now = nfd[i];
+      if( now == Blocks.size() ) {
+        continue;
+      }
+      if(idom[now] != sdom[now])
+        idom[now] = idom[idom[now]];
+    }
+  }
+}; // End of DominatorTree
+
+std::vector<uint32_t> Visited, InStack;
+std::vector<uint32_t> TopoOrder, InDeg;
+std::vector< std::vector<uint32_t> > t_Succ , t_Pred;
+
+void Go(uint32_t now, uint32_t tt) {
+  if( now == tt ) return;
+  Visited[now] = InStack[now] = timeStamp;
+
+  for(uint32_t nxt : Succs[now]) {
+    if(Visited[nxt] == timeStamp and InStack[nxt] == timeStamp) {
+      Marked.insert(nxt);
+    }
+    t_Succ[now].push_back(nxt);
+    t_Pred[nxt].push_back(now);
+    InDeg[nxt] += 1;
+    if(Visited[nxt] == timeStamp) {
+      continue;
+    }
+    Go(nxt, tt);
+  }
+
+  InStack[now] = 0;
+}
+
+void TopologicalSort(uint32_t ss, uint32_t tt) {
+  timeStamp += 1;
+
+  Go(ss, tt);
+
+  TopoOrder.clear();
+  std::queue<uint32_t> wait;
+  wait.push(ss);
+  while( not wait.empty() ) {
+    uint32_t now = wait.front(); wait.pop();
+    TopoOrder.push_back(now);
+    for(uint32_t nxt : t_Succ[now]) {
+      InDeg[nxt] -= 1;
+      if(InDeg[nxt] == 0u) {
+        wait.push(nxt);
+      }
+    }
+  }
+}
+
+std::vector< std::set<uint32_t> > NextMarked;
+bool Indistinguish(uint32_t node1, uint32_t node2) {
+  if(NextMarked[node1].size() > NextMarked[node2].size()){
+    uint32_t _swap = node1;
+    node1 = node2;
+    node2 = _swap;
+  }
+  for(uint32_t x : NextMarked[node1]) {
+    if( NextMarked[node2].find(x) != NextMarked[node2].end() ) {
+      return true;
+    }
+  }
+  return false;
+}
+
+void MakeUniq(uint32_t now) {
+  bool StopFlag = false;
+  if (Marked.find(now) == Marked.end()) {
+    for(uint32_t pred1 : t_Pred[now]) {
+      for(uint32_t pred2 : t_Pred[now]) {
+        if(pred1 == pred2) continue;
+        if(Indistinguish(pred1, pred2)) {
+          Marked.insert(now);
+          StopFlag = true;
+          break;
+        }
+      }
+      if (StopFlag) {
+        break;
+      }
+    }
+  }
+  if(Marked.find(now) != Marked.end()) {
+    NextMarked[now].insert(now);
+  } else {
+    for(uint32_t pred : t_Pred[now]) {
+      for(uint32_t x : NextMarked[pred]) {
+        NextMarked[now].insert(x);
+      }
+    }
+  }
+}
+
+void MarkSubGraph(uint32_t ss, uint32_t tt) {
+  TopologicalSort(ss, tt);
+  if(TopoOrder.empty()) return;
+
+  for(uint32_t i : TopoOrder) {
+    NextMarked[i].clear();
+  }
+
+  NextMarked[TopoOrder[0]].insert(TopoOrder[0]);
+  for(uint32_t i = 1 ; i < TopoOrder.size() ; i += 1) {
+    MakeUniq(TopoOrder[i]);
+  }
+}
+
+void MarkVertice(Function *F) {
+  uint32_t s = start_point;
+
+  InDeg.resize(Blocks.size());
+  Visited.resize(Blocks.size());
+  InStack.resize(Blocks.size());
+  t_Succ.resize(Blocks.size());
+  t_Pred.resize(Blocks.size());
+  NextMarked.resize(Blocks.size());
+
+  for( uint32_t i = 0 ; i < Blocks.size() ; i += 1 ) {
+    Visited[i] = InStack[i] = InDeg[i] = 0;
+    t_Succ[i].clear();
+    t_Pred[i].clear();
+  }
+  timeStamp = 0;
+  uint32_t t = 0;
+  //MarkSubGraph(s, t);
+  //return;
+
+  while( s != t ) {
+    MarkSubGraph(DominatorTree::idom[t], t);
+    t = DominatorTree::idom[t];
+  }
+
+}
+
+// return {marked nodes}
+std::pair<std::vector<BasicBlock *>,
+          std::vector<BasicBlock *> >markNodes(Function *F) {
+  assert(F->size() > 0 && "Function can not be empty");
+
+  reset();
+  labelEachBlock(F);
+  buildCFG(F);
+  turnCFGintoDAG(F);
+  DominatorTree::DominatorTree(F);
+  MarkVertice(F);
+
+  std::vector<BasicBlock *> Result , ResultAbove;
+  for( uint32_t x : Markabove ) {
+    auto it = Marked.find( x );
+    if( it != Marked.end() )
+      Marked.erase( it );
+    if( x )
+      ResultAbove.push_back(Blocks[x]);
+  }
+  for( uint32_t x : Marked ) {
+    if (x == 0) {
+      continue;
+    } else {
+      Result.push_back(Blocks[x]);
+    }
+  }
+
+  return { Result , ResultAbove };
+}
diff --git a/llvm_mode/MarkNodes.h b/llvm_mode/MarkNodes.h
new file mode 100644
index 00000000..e3bf3ce5
--- /dev/null
+++ b/llvm_mode/MarkNodes.h
@@ -0,0 +1,11 @@
+#ifndef __MARK_NODES__
+#define __MARK_NODES__
+
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Function.h"
+#include<vector>
+
+std::pair<std::vector<llvm::BasicBlock *>,
+          std::vector<llvm::BasicBlock *>> markNodes(llvm::Function *F);
+
+#endif
diff --git a/llvm_mode/afl-clang-fast.c b/llvm_mode/afl-clang-fast.c
index 5bc4ae8c..2034f10a 100644
--- a/llvm_mode/afl-clang-fast.c
+++ b/llvm_mode/afl-clang-fast.c
@@ -32,6 +32,7 @@
 #include <unistd.h>
 #include <stdlib.h>
 #include <string.h>
+#include <assert.h>
 
 static u8*  obj_path;               /* Path to runtime libraries         */
 static u8** cc_params;              /* Parameters passed to the real CC  */
@@ -87,7 +88,7 @@ static void find_obj(u8* argv0) {
     return;
   }
 
-  FATAL("Unable to find 'afl-llvm-rt.o' or 'afl-llvm-pass.so'. Please set AFL_PATH");
+  FATAL("Unable to find 'afl-llvm-rt.o' or 'libLLVMInsTrim.so'. Please set AFL_PATH");
  
 }
 
@@ -113,7 +114,7 @@ static void edit_params(u32 argc, char** argv) {
   }
 
   /* There are two ways to compile afl-clang-fast. In the traditional mode, we
-     use afl-llvm-pass.so to inject instrumentation. In the experimental
+     use libLLVMInsTrim.so to inject instrumentation. In the experimental
      'trace-pc-guard' mode, we use native LLVM instrumentation callbacks
      instead. The latter is a very recent addition - see:
 
@@ -150,7 +151,8 @@ static void edit_params(u32 argc, char** argv) {
   cc_params[cc_par_cnt++] = "-Xclang";
   cc_params[cc_par_cnt++] = "-load";
   cc_params[cc_par_cnt++] = "-Xclang";
-  cc_params[cc_par_cnt++] = alloc_printf("%s/afl-llvm-pass.so", obj_path);
+  cc_params[cc_par_cnt++] = alloc_printf("%s/libLLVMInsTrim.so", obj_path);
+//  cc_params[cc_par_cnt++] = alloc_printf("%s/afl-llvm-pass.so", obj_path);
 #endif /* ^USE_TRACE_PC */
 
   cc_params[cc_par_cnt++] = "-Qunused-arguments";