working ugly version

1 files changed, 330 insertions, 118 deletions

diff --git a/instrumentation/SanitizerCoverageLTO.so.cc b/instrumentation/SanitizerCoverageLTO.so.cc
index c74069e1..54cc1752 100644
--- a/instrumentation/SanitizerCoverageLTO.so.cc
+++ b/instrumentation/SanitizerCoverageLTO.so.cc
@@ -192,12 +192,14 @@ class ModuleSanitizerCoverageLTO
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
 
  private:
-  void            instrumentFunction(Function &F, DomTreeCallback DTCallback,
-                                     PostDomTreeCallback PDTCallback);
-  void            InjectCoverageForIndirectCalls(Function               &F,
-                                                 ArrayRef<Instruction *> IndirCalls);
-  bool            InjectCoverage(Function &F, ArrayRef<BasicBlock *> AllBlocks,
-                                 bool IsLeafFunc = true);
+  void instrumentFunction(Function &F, DomTreeCallback DTCallback,
+                          PostDomTreeCallback PDTCallback);
+  void InjectCoverageForIndirectCalls(Function               &F,
+                                      ArrayRef<Instruction *> IndirCalls);
+  bool InjectCoverage(Function &F, ArrayRef<BasicBlock *> AllBlocks,
+                      bool IsLeafFunc = true);
+  bool Fake_InjectCoverage(Function &F, ArrayRef<BasicBlock *> AllBlocks,
+                           bool IsLeafFunc = true);
   GlobalVariable *CreateFunctionLocalArrayInSection(size_t    NumElements,
                                                     Function &F, Type *Ty,
                                                     const char *Section);
@@ -248,6 +250,7 @@ class ModuleSanitizerCoverageLTO
   uint32_t                         unhandled = 0;
   uint32_t                         select_cnt = 0;
   uint32_t                         instrument_ctx = 0;
+  uint32_t                         extra_ctx_inst = 0;
   uint64_t                         map_addr = 0;
   const char                      *skip_nozero = NULL;
   const char                      *use_threadsafe_counters = nullptr;
@@ -259,6 +262,7 @@ class ModuleSanitizerCoverageLTO
   IntegerType                     *Int64Tyi = NULL;
   ConstantInt                     *Zero = NULL;
   ConstantInt                     *One = NULL;
+  AllocaInst                      *CTX_add = NULL;
   LLVMContext                     *Ct = NULL;
   Module                          *Mo = NULL;
   GlobalVariable                  *AFLMapPtr = NULL;
@@ -1164,9 +1168,10 @@ bool ModuleSanitizerCoverageLTO::instrumentModule(
                getenv("AFL_USE_TSAN") ? ", TSAN" : "",
                getenv("AFL_USE_CFISAN") ? ", CFISAN" : "",
                getenv("AFL_USE_UBSAN") ? ", UBSAN" : "");
-      OKF("Instrumented %u locations (%u selects) without collisions (%llu "
-          "collisions have been avoided) (%s mode).",
-          inst, select_cnt, calculateCollisions(inst), modeline);
+
+      OKF("Instrumented %u locations (%u selects) with %u extra map entries "
+          "for CTX (%s mode).",
+          inst, select_cnt, extra_ctx_inst, modeline);
 
     }
 
@@ -1292,19 +1297,112 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
   const DominatorTree     *DT = DTCallback(F);
   const PostDominatorTree *PDT = PDTCallback(F);
   bool                     IsLeafFunc = true;
-  uint32_t                 skip_next = 0, call_counter = 0;
-  Value                   *PrevCtx = NULL;
+  uint32_t                 skip_next = 0;
+  uint32_t                 call_counter = 0;
+  uint32_t                 inst_save = inst;
+  uint32_t                 inst_in_this_func = 0;
+  LLVMContext             &Context = F.getContext();
+
+  MDNode *N = MDNode::get(Context, MDString::get(Context, "nosanitize"));
+  CTX_add = NULL;
+
+  fprintf(stderr, "Function: %s\n", F.getName().str().c_str());
+
+  // Fake instrumentation so we can count how much there will be in this
+  // function
+
+  for (auto &BB : F) {
+
+    for (auto &IN : BB) {
+
+      CallInst *callInst = nullptr;
+
+      if ((callInst = dyn_cast<CallInst>(&IN))) {
+
+        Function *Callee = callInst->getCalledFunction();
+        if (!Callee) continue;
+        if (callInst->getCallingConv() != llvm::CallingConv::C) continue;
+        StringRef FuncName = Callee->getName();
+
+        if (FuncName.compare(StringRef("__afl_coverage_interesting"))) continue;
+
+        ++inst;
+
+      }
+
+      SelectInst *selectInst = nullptr;
+
+      if (!skip_next && (selectInst = dyn_cast<SelectInst>(&IN)) && 1 == 0) {
+
+        uint32_t vector_cnt = 0;
+        Value   *condition = selectInst->getCondition();
+        auto     t = condition->getType();
+
+        if (t->getTypeID() == llvm::Type::IntegerTyID) {
+
+          skip_next = 1;
+          inst += 2;
+
+        } else
+
+#if LLVM_VERSION_MAJOR >= 14
+            if (t->getTypeID() == llvm::Type::FixedVectorTyID) {
+
+          FixedVectorType *tt = dyn_cast<FixedVectorType>(t);
+          if (tt) {
+
+            uint32_t elements = tt->getElementCount().getFixedValue();
+            vector_cnt = elements;
+            inst += vector_cnt * 2;
 
-  MDNode *N =
-      MDNode::get(F.getContext(), MDString::get(F.getContext(), "nosanitize"));
+          }
+
+        } else
+
+#endif
+        {
+
+          continue;
+
+        }
+
+        skip_next = 1;
+
+      } else {
+
+        skip_next = 0;
+
+      }
+
+    }
+
+    if (shouldInstrumentBlock(F, &BB, DT, PDT, Options))
+      BlocksToInstrument.push_back(&BB);
+    /*
+        if (Options.IndirectCalls)
+          for (auto &Inst : BB) {
+
+    fprintf(stderr, "FUCK TODO\n");
+            CallBase *CB = dyn_cast<CallBase>(&Inst);
+            if (CB && !CB->getCalledFunction()) IndirCalls.push_back(&Inst);
+    /// XXX TODO!!!! set 0 !
+
+          }
+
+    */
+
+  }
+
+  Fake_InjectCoverage(F, BlocksToInstrument, IsLeafFunc);
+
+  // CTX init
 
   for (auto &BB : F) {
 
-    if (/*F.size() > 1 &&*/ instrument_ctx && &BB == &F.getEntryBlock()) {
+    if (instrument_ctx && &BB == &F.getEntryBlock()) {
 
       // we insert a CTX value in all our callers:
-      LLVMContext &Context = F.getContext();
-      IRBuilder<>  Builder(Context);
+      IRBuilder<> Builder(Context);
       for (auto *U : F.users()) {
 
         if (auto *CI = dyn_cast<CallInst>(U)) {
@@ -1323,24 +1421,60 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
       }
 
       // We read the CTX for this call
+      Value               *CTX_offset;
       BasicBlock::iterator IP = BB.getFirstInsertionPt();
       IRBuilder<>          IRB(&(*IP));
       LoadInst            *PrevCtxLoad = IRB.CreateLoad(
 #if LLVM_VERSION_MAJOR >= 14
-          Builder.getInt32Ty(),
+          IRB.getInt32Ty(),
 #endif
           AFLContext);
+
       PrevCtxLoad->setMetadata("nosanitize", N);
-      PrevCtx = PrevCtxLoad;
+
+      if (inst == inst_save || call_counter < 2) {
+
+        fprintf(stderr, "%s: ZERO!\n", F.getName().str().c_str());
+        CTX_offset = Zero;
+
+      } else {
+
+        fprintf(stderr, "%s: %u * %u\n", F.getName().str().c_str(),
+                inst - inst_save, call_counter);
+        CTX_offset = IRB.CreateMul(
+            ConstantInt::get(Type::getInt32Ty(Context), inst - inst_save),
+            PrevCtxLoad, "CTXmul", false, true);
+
+      }
+
+      CTX_add = IRB.CreateAlloca(Type::getInt32Ty(Context), nullptr, "CTX_add");
+      auto nosan = IRB.CreateStore(CTX_offset, CTX_add);
+      nosan->setMetadata("nosanitize", N);
 
     }
 
   }
 
+  inst_in_this_func = inst - inst_save;
+  inst = inst_save;
+
+  // Now the real instrumentation
+
+  IsLeafFunc = true;
+  skip_next = 0;
+
   for (auto &BB : F) {
 
+    // fprintf(stderr, "BB: %s\n", BB.getName().str().c_str());
+
     for (auto &IN : BB) {
 
+      /*
+            std::string              Str;
+            llvm::raw_string_ostream RSO(Str);
+            IN.print(RSO);
+            errs() << RSO.str() << "\n";
+      */
       CallInst *callInst = nullptr;
 
       if ((callInst = dyn_cast<CallInst>(&IN))) {
@@ -1363,8 +1497,16 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
 
         if (FuncName.compare(StringRef("__afl_coverage_interesting"))) continue;
 
-        Value *val = ConstantInt::get(Int32Ty, ++afl_global_id);
-        callInst->setOperand(1, val);
+        IRBuilder<> Builder(Context);
+        Value      *val = ConstantInt::get(Int32Ty, ++afl_global_id);
+        LoadInst   *CTX_load = Builder.CreateLoad(
+#if LLVM_VERSION_MAJOR >= 14
+            Builder.getInt32Ty(),
+#endif
+            CTX_add);
+        ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(CTX_load);
+        Value *val2 = Builder.CreateAdd(val, CTX_load);
+        callInst->setOperand(1, val2);
         ++inst;
 
       }
@@ -1377,158 +1519,193 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
             IN.print(os);
             fprintf(stderr, "X(%u): %s\n", skip_next, os.str().c_str());
       */
-      if (!skip_next && (selectInst = dyn_cast<SelectInst>(&IN))) {
+      if ((selectInst = dyn_cast<SelectInst>(&IN))) {
 
-        uint32_t    vector_cnt = 0;
-        Value      *condition = selectInst->getCondition();
-        Value      *result;
-        auto        t = condition->getType();
-        IRBuilder<> IRB(selectInst->getNextNode());
-
-        ++select_cnt;
-
-        if (t->getTypeID() == llvm::Type::IntegerTyID) {
-
-          Value *val1 = ConstantInt::get(Int32Ty, ++afl_global_id);
-          Value *val2 = ConstantInt::get(Int32Ty, ++afl_global_id);
-          result = IRB.CreateSelect(condition, val1, val2);
-          skip_next = 1;
-          inst += 2;
+        if (!skip_next) {
 
-        } else
+          uint32_t    vector_cnt = 0;
+          Value      *condition = selectInst->getCondition();
+          Value      *result;
+          auto        t = condition->getType();
+          IRBuilder<> IRB(selectInst->getNextNode());
 
+          // fprintf(stderr, "ADDING!!!\n");
+          LoadInst *CTX_load = IRB.CreateLoad(
 #if LLVM_VERSION_MAJOR >= 14
-            if (t->getTypeID() == llvm::Type::FixedVectorTyID) {
+              IRB.getInt32Ty(),
+#endif
+              CTX_add);
+          ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(CTX_load);
 
-          FixedVectorType *tt = dyn_cast<FixedVectorType>(t);
-          if (tt) {
+          ++select_cnt;
 
-            uint32_t elements = tt->getElementCount().getFixedValue();
-            vector_cnt = elements;
-            inst += vector_cnt * 2;
-            if (elements) {
+          if (t->getTypeID() == llvm::Type::IntegerTyID) {
 
-              FixedVectorType *GuardPtr1 =
-                  FixedVectorType::get(Int32Ty, elements);
-              FixedVectorType *GuardPtr2 =
-                  FixedVectorType::get(Int32Ty, elements);
-              Value *x, *y;
+            Value *val1 = ConstantInt::get(Int32Ty, ++afl_global_id);
+            Value *val11 = IRB.CreateAdd(val1, CTX_load);
+            Value *val2 = ConstantInt::get(Int32Ty, ++afl_global_id);
+            Value *val22 = IRB.CreateAdd(val2, CTX_load);
+            result = IRB.CreateSelect(condition, val11, val22);
+            skip_next = 1;
+            inst += 2;
 
-              Value *val1 = ConstantInt::get(Int32Ty, ++afl_global_id);
-              Value *val2 = ConstantInt::get(Int32Ty, ++afl_global_id);
-              x = IRB.CreateInsertElement(GuardPtr1, val1, (uint64_t)0);
-              y = IRB.CreateInsertElement(GuardPtr2, val2, (uint64_t)0);
+          } else
+
+#if LLVM_VERSION_MAJOR >= 14
+              if (t->getTypeID() == llvm::Type::FixedVectorTyID) {
+
+            FixedVectorType *tt = dyn_cast<FixedVectorType>(t);
+            if (tt) {
+
+              uint32_t elements = tt->getElementCount().getFixedValue();
+              vector_cnt = elements;
+              inst += vector_cnt * 2;
+              if (elements) {
+
+                FixedVectorType *GuardPtr1 =
+                    FixedVectorType::get(Int32Ty, elements);
+                FixedVectorType *GuardPtr2 =
+                    FixedVectorType::get(Int32Ty, elements);
+                Value *x, *y;
+
+                Value *val1 = ConstantInt::get(Int32Ty, ++afl_global_id);
+                Value *val11 = IRB.CreateAdd(val1, CTX_add);
+                Value *val2 = ConstantInt::get(Int32Ty, ++afl_global_id);
+                Value *val22 = IRB.CreateAdd(val2, CTX_add);
+                x = IRB.CreateInsertElement(GuardPtr1, val11, (uint64_t)0);
+                y = IRB.CreateInsertElement(GuardPtr2, val22, (uint64_t)0);
+
+                for (uint64_t i = 1; i < elements; i++) {
+
+                  val1 = ConstantInt::get(Int32Ty, ++afl_global_id);
+                  val11 = IRB.CreateAdd(val1, CTX_add);
+                  val2 = ConstantInt::get(Int32Ty, ++afl_global_id);
+                  val11 = IRB.CreateAdd(val1, CTX_add);
+                  x = IRB.CreateInsertElement(GuardPtr1, val11, i);
+                  y = IRB.CreateInsertElement(GuardPtr2, val22, i);
 
-              for (uint64_t i = 1; i < elements; i++) {
+                }
 
-                val1 = ConstantInt::get(Int32Ty, ++afl_global_id);
-                val2 = ConstantInt::get(Int32Ty, ++afl_global_id);
-                x = IRB.CreateInsertElement(GuardPtr1, val1, i);
-                y = IRB.CreateInsertElement(GuardPtr2, val2, i);
+                result = IRB.CreateSelect(condition, x, y);
+                skip_next = 1;
 
               }
 
-              result = IRB.CreateSelect(condition, x, y);
-              skip_next = 1;
-
             }
 
-          }
-
-        } else
+          } else
 
 #endif
-        {
+          {
 
-          unhandled++;
-          continue;
+            unhandled++;
+            continue;
 
-        }
+          }
 
-        uint32_t vector_cur = 0;
-        /* Load SHM pointer */
-        LoadInst *MapPtr =
-            IRB.CreateLoad(PointerType::get(Int8Ty, 0), AFLMapPtr);
-        ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(MapPtr);
+          uint32_t vector_cur = 0;
+          /* Load SHM pointer */
+          LoadInst *MapPtr =
+              IRB.CreateLoad(PointerType::get(Int8Ty, 0), AFLMapPtr);
+          ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(MapPtr);
 
-        while (1) {
+          while (1) {
 
-          /* Get CurLoc */
-          Value *MapPtrIdx = nullptr;
+            /* Get CurLoc */
+            Value *MapPtrIdx = nullptr;
 
-          /* Load counter for CurLoc */
-          if (!vector_cnt) {
+            /* Load counter for CurLoc */
+            if (!vector_cnt) {
 
-            MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, result);
+              MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, result);
 
-          } else {
+            } else {
 
-            auto element = IRB.CreateExtractElement(result, vector_cur++);
-            MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, element);
+              auto element = IRB.CreateExtractElement(result, vector_cur++);
+              MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, element);
 
-          }
+            }
 
-          if (use_threadsafe_counters) {
+            if (use_threadsafe_counters) {
 
-            IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One,
+              IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx,
+                                  One,
 #if LLVM_VERSION_MAJOR >= 13
-                                llvm::MaybeAlign(1),
+                                  llvm::MaybeAlign(1),
 #endif
-                                llvm::AtomicOrdering::Monotonic);
+                                  llvm::AtomicOrdering::Monotonic);
 
-          } else {
+            } else {
 
-            LoadInst *Counter = IRB.CreateLoad(IRB.getInt8Ty(), MapPtrIdx);
-            ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(Counter);
+              LoadInst *Counter = IRB.CreateLoad(IRB.getInt8Ty(), MapPtrIdx);
+              ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(Counter);
 
-            /* Update bitmap */
+              /* Update bitmap */
 
-            Value *Incr = IRB.CreateAdd(Counter, One);
+              Value *Incr = IRB.CreateAdd(Counter, One);
 
-            if (skip_nozero == NULL) {
+              if (skip_nozero == NULL) {
 
-              auto cf = IRB.CreateICmpEQ(Incr, Zero);
-              auto carry = IRB.CreateZExt(cf, Int8Ty);
-              Incr = IRB.CreateAdd(Incr, carry);
+                auto cf = IRB.CreateICmpEQ(Incr, Zero);
+                auto carry = IRB.CreateZExt(cf, Int8Ty);
+                Incr = IRB.CreateAdd(Incr, carry);
+
+              }
+
+              auto nosan = IRB.CreateStore(Incr, MapPtrIdx);
+              ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(nosan);
 
             }
 
-            auto nosan = IRB.CreateStore(Incr, MapPtrIdx);
-            ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(nosan);
+            if (!vector_cnt || vector_cnt == vector_cur) { break; }
 
           }
 
-          if (!vector_cnt || vector_cnt == vector_cur) { break; }
+          skip_next = 1;
 
-        }
+        } else {
 
-        skip_next = 1;
+          skip_next = 0;
 
-      } else {
-
-        skip_next = 0;
+        }
 
       }
 
     }
 
-    if (shouldInstrumentBlock(F, &BB, DT, PDT, Options))
-      BlocksToInstrument.push_back(&BB);
-    for (auto &Inst : BB) {
+    //    if (shouldInstrumentBlock(F, &BB, DT, PDT, Options))
+    //      BlocksToInstrument.push_back(&BB);
 
-      if (Options.IndirectCalls) {
+  }
 
-        CallBase *CB = dyn_cast<CallBase>(&Inst);
-        if (CB && !CB->getCalledFunction()) IndirCalls.push_back(&Inst);
+  InjectCoverage(F, BlocksToInstrument, IsLeafFunc);
+  // InjectCoverageForIndirectCalls(F, IndirCalls);
 
-      }
+  if (inst_in_this_func && call_counter > 1) {
 
-    }
+    extra_ctx_inst += inst_in_this_func * (call_counter - 1);
+    afl_global_id += inst_in_this_func * (call_counter - 1);
 
   }
 
-  InjectCoverage(F, BlocksToInstrument, IsLeafFunc);
-  InjectCoverageForIndirectCalls(F, IndirCalls);
+  /*
+    fprintf(stderr, "FUNCTION: %s\n", F.getName().str().c_str());
+    int n = 0;
+    for (auto &BB : F) {
+
+      fprintf(stderr, "BB %d\n", n++);
+      for (auto &IN : BB) {
+
+        std::string              Str;
+        llvm::raw_string_ostream RSO(Str);
+        IN.print(RSO);
+        errs() << RSO.str() << "\n";
+
+      }
+
+    }
+
+  */
 
 }
 
@@ -1654,6 +1831,36 @@ bool ModuleSanitizerCoverageLTO::InjectCoverage(
 
 }
 
+// Fake InjectCoverage to count the instrumentations
+bool ModuleSanitizerCoverageLTO::Fake_InjectCoverage(
+    Function &F, ArrayRef<BasicBlock *> AllBlocks, bool IsLeafFunc) {
+
+  if (AllBlocks.empty()) return false;
+
+  for (size_t i = 0, N = AllBlocks.size(); i < N; i++) {
+
+    // AFL++ START
+    if (BlockList.size()) {
+
+      int skip = 0;
+      for (uint32_t k = 0; k < BlockList.size(); k++) {
+
+        if (AllBlocks[i] == BlockList[k]) { skip = 1; }
+
+      }
+
+      if (skip) continue;
+
+    }
+
+    inst++;  // InjectCoverageAtBlock()
+
+  }
+
+  return true;
+
+}
+
 // On every indirect call we call a run-time function
 // __sanitizer_cov_indir_call* with two parameters:
 //   - callee address,
@@ -1725,6 +1932,13 @@ void ModuleSanitizerCoverageLTO::InjectCoverageAtBlock(Function   &F,
     /* Set the ID of the inserted basic block */
 
     ConstantInt *CurLoc = ConstantInt::get(Int32Tyi, afl_global_id);
+    LoadInst    *CTX_load = IRB.CreateLoad(
+#if LLVM_VERSION_MAJOR >= 14
+        IRB.getInt32Ty(),
+#endif
+        CTX_add);
+    ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(CTX_load);
+    Value *offset = IRB.CreateAdd(CurLoc, CTX_load);
 
     /* Load SHM pointer */
 
@@ -1732,13 +1946,13 @@ void ModuleSanitizerCoverageLTO::InjectCoverageAtBlock(Function   &F,
 
     if (map_addr) {
 
-      MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtrFixed, CurLoc);
+      MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtrFixed, offset);
 
     } else {
 
       LoadInst *MapPtr = IRB.CreateLoad(PointerType::get(Int8Ty, 0), AFLMapPtr);
       ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(MapPtr);
-      MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, CurLoc);
+      MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, offset);
 
     }
 
@@ -1777,8 +1991,6 @@ void ModuleSanitizerCoverageLTO::InjectCoverageAtBlock(Function   &F,
     // AFL++ END
 
     /*
-    XXXXXXXXXXXXXXXXXXX
-
         auto GuardPtr = IRB.CreateIntToPtr(
             IRB.CreateAdd(IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
                           ConstantInt::get(IntptrTy, Idx * 4)),