better CTX instrumentation

instrumentation/SanitizerCoverageLTO.so.cc

@@ -194,8 +194,9 @@ class ModuleSanitizerCoverageLTO
  private:
   void instrumentFunction(Function &F, DomTreeCallback DTCallback,
                           PostDomTreeCallback PDTCallback);
-  void InjectCoverageForIndirectCalls(Function               &F,
-                                      ArrayRef<Instruction *> IndirCalls);
+  /*  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,
@@ -263,12 +264,12 @@ class ModuleSanitizerCoverageLTO
   ConstantInt                     *Zero = NULL;
   ConstantInt                     *Zero32 = NULL;
   ConstantInt                     *One = NULL;
-  AllocaInst                      *CTX_add = NULL;
   LLVMContext                     *Ct = NULL;
   Module                          *Mo = NULL;
-  GlobalVariable                  *AFLMapPtr = NULL;
   GlobalVariable                  *AFLContext = NULL;
+  GlobalVariable                  *AFLMapPtr = NULL;
   Value                           *MapPtrFixed = NULL;
+  AllocaInst                      *CTX_add = NULL;
   std::ofstream                    dFile;
   size_t                           found = 0;
   // AFL++ END
@@ -427,7 +428,11 @@ bool ModuleSanitizerCoverageLTO::instrumentModule(
   setvbuf(stdout, NULL, _IONBF, 0);
   if (getenv("AFL_DEBUG")) { debug = 1; }
   if (getenv("AFL_LLVM_DICT2FILE_NO_MAIN")) { autodictionary_no_main = 1; }
-  if (getenv("AFL_LLVM_CALLER")) { instrument_ctx = 1; }
+  if (getenv("AFL_LLVM_CALLER") || getenv("AFL_LLVM_CTX")) {
+
+    instrument_ctx = 1;
+
+  }
 
   if ((isatty(2) && !getenv("AFL_QUIET")) || debug) {
 
@@ -435,10 +440,12 @@ bool ModuleSanitizerCoverageLTO::instrumentModule(
               "%s by Marc \"vanHauser\" Heuse <mh@mh-sec.de>\n",
          instrument_ctx ? " (CTX mode)" : "");
 
-  } else
+  } else {
 
     be_quiet = 1;
 
+  }
+
   skip_nozero = getenv("AFL_LLVM_SKIP_NEVERZERO");
   use_threadsafe_counters = getenv("AFL_LLVM_THREADSAFE_INST");
 
@@ -611,12 +618,12 @@ bool ModuleSanitizerCoverageLTO::instrumentModule(
                 }
 
                 dictionary.push_back(std::string((char *)&val, len));
-                found++;
+                ++found;
 
                 if (val2) {
 
                   dictionary.push_back(std::string((char *)&val2, len));
-                  found++;
+                  ++found;
 
                 }
 
@@ -943,7 +950,7 @@ bool ModuleSanitizerCoverageLTO::instrumentModule(
                  '\0') {
 
                               thestring.append("\0", 1);  // add null byte
-                              optLen++;
+                              ++optLen;
 
                             }
 
@@ -1094,7 +1101,7 @@ bool ModuleSanitizerCoverageLTO::instrumentModule(
       for (auto token : dictionary) {
 
         memlen += token.length();
-        count++;
+        ++count;
 
       }
 
@@ -1115,7 +1122,7 @@ bool ModuleSanitizerCoverageLTO::instrumentModule(
             ptrhld.get()[offset++] = (uint8_t)token.length();
             memcpy(ptrhld.get() + offset, token.c_str(), token.length());
             offset += token.length();
-            count++;
+            ++count;
 
           }
 
@@ -1162,7 +1169,7 @@ bool ModuleSanitizerCoverageLTO::instrumentModule(
       WARNF("No instrumentation targets found.");
     else {
 
-      char modeline[100];
+      char modeline[128];
       snprintf(modeline, sizeof(modeline), "%s%s%s%s%s%s",
                getenv("AFL_HARDEN") ? "hardened" : "non-hardened",
                getenv("AFL_USE_ASAN") ? ", ASAN" : "",
@@ -1170,10 +1177,16 @@ bool ModuleSanitizerCoverageLTO::instrumentModule(
                getenv("AFL_USE_TSAN") ? ", TSAN" : "",
                getenv("AFL_USE_CFISAN") ? ", CFISAN" : "",
                getenv("AFL_USE_UBSAN") ? ", UBSAN" : "");
+      char buf[64] = {};
+      if (instrument_ctx) {
 
-      OKF("Instrumented %u locations (%u selects) with %u extra map entries "
-          "for CTX (%s mode).",
-          inst, select_cnt, extra_ctx_inst, modeline);
+        snprintf(buf, sizeof(buf), " with %u extra map entries for CTX",
+                 extra_ctx_inst);
+
+      }
+
+      OKF("Instrumented %u locations (%u selects)%s (%s mode).", inst,
+          select_cnt, buf, modeline);
 
     }
 
@@ -1254,6 +1267,52 @@ static bool shouldInstrumentBlock(const Function &F, const BasicBlock *BB,
 
 }
 
+/// return the number of calls to this function
+u32 countCallers(Function *F) {
+
+  u32 callers = 0;
+
+  if (!F) { return 0; }
+
+  for (auto *U : F->users()) {
+
+    if (auto *CI = dyn_cast<CallInst>(U)) { ++callers; }
+
+  }
+
+  return callers;
+
+}
+
+/// return the calling function of a function - only if there is a single caller
+Function *returnOnlyCaller(Function *F) {
+
+  Function *caller = NULL;
+
+  if (!F) { return NULL; }
+
+  for (auto *U : F->users()) {
+
+    if (auto *CI = dyn_cast<CallInst>(U)) {
+
+      if (caller == NULL) {
+
+        caller = CI->getParent()->getParent();
+
+      } else {
+
+        return NULL;
+
+      }
+
+    }
+
+  }
+
+  return caller;
+
+}
+
 void ModuleSanitizerCoverageLTO::instrumentFunction(
     Function &F, DomTreeCallback DTCallback, PostDomTreeCallback PDTCallback) {
 
@@ -1287,6 +1346,37 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
 
   // AFL++ START
   if (!F.size()) return;
+
+  LLVMContext &Context = F.getContext();
+  MDNode      *N = MDNode::get(Context, MDString::get(Context, "nosanitize"));
+
+  if (instrument_ctx) {
+
+    // we have to set __afl_ctx 0 for all indirect calls in all functions, even
+    // those not to be instrumented.
+    for (auto &BB : F) {
+
+      for (auto &IN : BB) {
+
+        if (auto *Call = dyn_cast<CallInst>(&IN)) {
+
+          if (Call->isIndirectCall()) {
+
+            IRBuilder<> Builder(IN.getContext());
+            Builder.SetInsertPoint(IN.getParent(), IN.getIterator());
+            StoreInst *StoreCtx = Builder.CreateStore(Zero32, AFLContext);
+            StoreCtx->setMetadata("nosanitize", N);
+
+          }
+
+        }
+
+      }
+
+    }
+
+  }
+
   if (!isInInstrumentList(&F, FMNAME)) return;
   // AFL++ END
 
@@ -1300,19 +1390,74 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
   const PostDominatorTree *PDT = PDTCallback(F);
   bool                     IsLeafFunc = true;
   uint32_t                 skip_next = 0;
-  uint32_t                 call_counter = 0;
+  uint32_t                 call_counter = 0, call_depth = 0;
   uint32_t                 inst_save = inst, save_global = afl_global_id;
   uint32_t                 inst_in_this_func = 0;
-  LLVMContext             &Context = F.getContext();
+  Function                *caller = NULL;
 
-  MDNode *N = MDNode::get(Context, MDString::get(Context, "nosanitize"));
   CTX_add = NULL;
 
-  fprintf(stderr, "Function: %s\n", F.getName().str().c_str());
+  if (debug) fprintf(stderr, "Function: %s\n", F.getName().str().c_str());
 
-  // Fake instrumentation so we can count how much there will be in this
-  // function
+  if (instrument_ctx) {
 
+    caller = &F;
+    call_counter = countCallers(caller);
+    Function *callee = caller;
+
+    if (call_counter == 1) {
+
+      ++call_depth;
+
+      while (((caller = returnOnlyCaller(callee)) || 1 == 1) &&
+             (call_counter = countCallers(callee)) == 1) {
+
+        if (debug && caller && callee)
+          fprintf(stderr, "DEBUG: another depth: %s <- %s [%u]\n",
+                  callee->getName().str().c_str(),
+                  caller->getName().str().c_str(), call_depth);
+        ++call_depth;
+        callee = caller;
+
+      }
+
+      if (!caller && callee) {
+
+        caller = callee;
+        if (debug)
+          fprintf(stderr, "DEBUG: depth found: %s <- %s [count=%u, depth=%u]\n",
+                  caller->getName().str().c_str(), F.getName().str().c_str(),
+                  call_counter, call_depth);
+
+      }
+
+    }
+
+    if (debug && call_counter < 2) {
+
+      fprintf(stderr, "Function %s only %u (%s)\n", F.getName().str().c_str(),
+              call_counter, caller->getName().str().c_str());
+
+    }
+
+    if (call_counter == 1) {
+
+      call_counter = 0;
+      caller = NULL;
+
+    }
+
+    if (debug) {
+
+      fprintf(stderr, "DEBUG: result: Function=%s callers=%u depth=%u\n",
+              F.getName().str().c_str(), call_counter, call_depth);
+
+    }
+
+    if (call_counter > 1) {
+
+      // Fake instrumentation so we can count how many instrumentations there
+      // will be in this function
       for (auto &BB : F) {
 
         for (auto &IN : BB) {
@@ -1326,7 +1471,8 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
             if (callInst->getCallingConv() != llvm::CallingConv::C) continue;
             StringRef FuncName = Callee->getName();
 
-        if (FuncName.compare(StringRef("__afl_coverage_interesting"))) continue;
+            if (FuncName.compare(StringRef("__afl_coverage_interesting")))
+              continue;
 
             ++inst;
 
@@ -1371,41 +1517,47 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
 
         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
+      if (debug)
+        fprintf(stderr, "DEBUG: CTX: %u instrumentations\n", inst - inst_save);
 
+      // we only instrument functions that have more than one instrumented block
+      if (inst > inst_save + 1) {
+
+        inst_in_this_func = inst - inst_save;
+        bool done = false;
+
+        // in rare occasions there can be multiple entry points per function
         for (auto &BB : F) {
 
-    if (instrument_ctx && &BB == &F.getEntryBlock()) {
+          if (&BB == &F.getEntryBlock() && done == false) {
 
             // we insert a CTX value in all our callers:
             IRBuilder<> Builder(Context);
-      for (auto *U : F.users()) {
+            CallInst   *CI = NULL;
+            Function   *F2 = NULL;
+            uint32_t    instrumented_calls = 0;
 
-        if (auto *CI = dyn_cast<CallInst>(U)) {
+            for (auto *U : caller->users()) {
+
+              if ((CI = dyn_cast<CallInst>(U))) {
+
+                F2 = CI->getParent()->getParent();
+                if (debug)
+                  fprintf(stderr,
+                          "DEBUG: CTX call insert %s [%u/%u] -> %s/%s\n",
+                          F2->getName().str().c_str(), instrumented_calls + 1,
+                          call_counter, caller->getName().str().c_str(),
+                          F.getName().str().c_str());
 
-          fprintf(stderr, "Insert %s [%u] -> %s\n",
-                  CI->getParent()->getParent()->getName().str().c_str(),
-                  call_counter, F.getName().str().c_str());
                 Builder.SetInsertPoint(CI);
                 StoreInst *StoreCtx = Builder.CreateStore(
-              ConstantInt::get(Type::getInt32Ty(Context), call_counter++),
+                    ConstantInt::get(Type::getInt32Ty(Context),
+                                     instrumented_calls++),
                     AFLContext);
                 StoreCtx->setMetadata("nosanitize", N);
 
@@ -1413,7 +1565,23 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
 
             }
 
-      // We read the CTX for this call
+            if (instrumented_calls != call_counter) {
+
+              fprintf(stderr, "BUG! %s/%s <=> %u vs %u\n",
+                      caller->getName().str().c_str(),
+                      F.getName().str().c_str(), instrumented_calls,
+                      call_counter);
+              exit(-1);
+
+            }
+
+            done = true;
+
+          }
+
+          // in all entrypoints we have to load the CTX value
+          if (&BB == &F.getEntryBlock()) {
+
             Value               *CTX_offset;
             BasicBlock::iterator IP = BB.getFirstInsertionPt();
             IRBuilder<>          IRB(&(*IP));
@@ -1422,85 +1590,38 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
                 IRB.getInt32Ty(),
 #endif
                 AFLContext);
-
             PrevCtxLoad->setMetadata("nosanitize", N);
 
-      if (inst == inst_save || call_counter < 2) {
-
-        fprintf(stderr, "%s: ZERO!\n", F.getName().str().c_str());
-        CTX_offset = Zero32;
-
-      } 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),
+                ConstantInt::get(Type::getInt32Ty(Context), inst_in_this_func),
                 PrevCtxLoad, "CTXmul", false, true);
 
-      }
-
-      CTX_add = IRB.CreateAlloca(Type::getInt32Ty(Context), nullptr, "CTX_add");
+            CTX_add =
+                IRB.CreateAlloca(Type::getInt32Ty(Context), nullptr, "CTX_add");
             auto nosan = IRB.CreateStore(CTX_offset, CTX_add);
             nosan->setMetadata("nosanitize", N);
 
-    }
-
-    // we have to set __afl_ctx 0 for all indirect calls
-    for (auto &IN : BB) {
-
-      if (auto *Call = dyn_cast<CallInst>(&IN)) {
-
-        if (Call->isIndirectCall()) {
-
-          IRBuilder<> Builder(IN.getContext());
-          Builder.SetInsertPoint(IN.getParent(), IN.getIterator());
-          StoreInst *StoreCtx = Builder.CreateStore(Zero32, AFLContext);
-          StoreCtx->setMetadata("nosanitize", N);
+            if (debug)
+              fprintf(
+                  stderr, "DEBUG: extra CTX instrumentations for %s: %u * %u\n",
+                  F.getName().str().c_str(), inst - inst_save, call_counter);
 
           }
 
         }
 
-    }
-
-  }
-
-  inst_in_this_func = inst - inst_save;
         inst = inst_save;
 
-  // Now the real instrumentation
-
-  if (CTX_add == NULL) {
-
-    auto BB = &F.getEntryBlock();
-    if (!BB) {
-
-      fprintf(stderr, "NULL entry %s %p\n", F.getName().str().c_str(), BB);
-      exit(-1);
-
       }
 
-    BasicBlock::iterator IP = BB->getFirstInsertionPt();
-    IRBuilder<>          IRB(&(*IP));
-    CTX_add = IRB.CreateAlloca(Type::getInt32Ty(Context), nullptr, "CTX_add");
-    auto nosan = IRB.CreateStore(Zero32, CTX_add);
-    nosan->setMetadata("nosanitize", N);
+    }
 
   }
 
   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))) {
@@ -1523,16 +1644,21 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
 
         if (FuncName.compare(StringRef("__afl_coverage_interesting"))) continue;
 
-        IRBuilder<> Builder(Context);
         Value *val = ConstantInt::get(Int32Ty, ++afl_global_id);
+        if (CTX_add) {
+
+          IRBuilder<> Builder(Context);
           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);
+          val = Builder.CreateAdd(val, CTX_load);
+
+        }
+
+        callInst->setOperand(1, val);
         ++inst;
 
       }
@@ -1545,9 +1671,7 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
             IN.print(os);
             fprintf(stderr, "X(%u): %s\n", skip_next, os.str().c_str());
       */
-      if ((selectInst = dyn_cast<SelectInst>(&IN))) {
-
-        if (!skip_next) {
+      if (!skip_next && (selectInst = dyn_cast<SelectInst>(&IN))) {
 
         uint32_t    vector_cnt = 0;
         Value      *condition = selectInst->getCondition();
@@ -1555,23 +1679,25 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
         auto        t = condition->getType();
         IRBuilder<> IRB(selectInst->getNextNode());
 
-          // fprintf(stderr, "ADDING!!!\n");
-          LoadInst *CTX_load = IRB.CreateLoad(
-#if LLVM_VERSION_MAJOR >= 14
-              IRB.getInt32Ty(),
-#endif
-              CTX_add);
-          ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(CTX_load);
-
         ++select_cnt;
 
         if (t->getTypeID() == llvm::Type::IntegerTyID) {
 
           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);
+          if (CTX_add) {
+
+            LoadInst *CTX_load = IRB.CreateLoad(
+#if LLVM_VERSION_MAJOR >= 14
+                IRB.getInt32Ty(),
+#endif
+                CTX_add);
+            val1 = IRB.CreateAdd(val1, CTX_load);
+            val2 = IRB.CreateAdd(val2, CTX_load);
+
+          }
+
+          result = IRB.CreateSelect(condition, val1, val2);
           skip_next = 1;
           inst += 2;
 
@@ -1586,7 +1712,6 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
             uint32_t elements = tt->getElementCount().getFixedValue();
             vector_cnt = elements;
             inst += vector_cnt * 2;
-
             if (elements) {
 
               FixedVectorType *GuardPtr1 =
@@ -1596,20 +1721,40 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
               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);
-                x = IRB.CreateInsertElement(GuardPtr1, val11, (uint64_t)0);
-                y = IRB.CreateInsertElement(GuardPtr2, val22, (uint64_t)0);
+              if (CTX_add) {
+
+                LoadInst *CTX_load = IRB.CreateLoad(
+  #if LLVM_VERSION_MAJOR >= 14
+                    IRB.getInt32Ty(),
+  #endif
+                    CTX_add);
+                val1 = IRB.CreateAdd(val1, CTX_load);
+                val2 = IRB.CreateAdd(val2, CTX_load);
+
+              }
+
+              x = IRB.CreateInsertElement(GuardPtr1, val1, (uint64_t)0);
+              y = IRB.CreateInsertElement(GuardPtr2, val2, (uint64_t)0);
 
               for (uint64_t i = 1; i < elements; i++) {
 
                 val1 = ConstantInt::get(Int32Ty, ++afl_global_id);
-                  val11 = IRB.CreateAdd(val1, CTX_load);
                 val2 = ConstantInt::get(Int32Ty, ++afl_global_id);
-                  val11 = IRB.CreateAdd(val1, CTX_load);
-                  x = IRB.CreateInsertElement(GuardPtr1, val11, i);
-                  y = IRB.CreateInsertElement(GuardPtr2, val22, i);
+                /*if (CTX_add) { // already loaded I guess
+
+                  LoadInst *CTX_load = IRB.CreateLoad(
+  #if LLVM_VERSION_MAJOR >= 14
+                      IRB.getInt32Ty(),
+  #endif
+                      CTX_add);
+                  val1 = IRB.CreateAdd(val1, CTX_load);
+                  val2 = IRB.CreateAdd(val2, CTX_load);
+
+                }*/
+
+                x = IRB.CreateInsertElement(GuardPtr1, val1, i);
+                y = IRB.CreateInsertElement(GuardPtr2, val2, i);
 
               }
 
@@ -1625,7 +1770,7 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
 #endif
         {
 
-            unhandled++;
+          ++unhandled;
           continue;
 
         }
@@ -1655,13 +1800,11 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
 
           if (use_threadsafe_counters) {
 
-              auto nosan = IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add,
-                                               MapPtrIdx, One,
+            IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One,
 #if LLVM_VERSION_MAJOR >= 13
                                 llvm::MaybeAlign(1),
 #endif
                                 llvm::AtomicOrdering::Monotonic);
-              ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(nosan);
 
           } else {
 
@@ -1699,11 +1842,21 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
 
     }
 
-    }
-
     // if (shouldInstrumentBlock(F, &BB, DT, PDT, Options))
     //   BlocksToInstrument.push_back(&BB);
 
+    /*
+        for (auto &Inst : BB) {
+
+          if (Options.IndirectCalls) {
+
+            CallBase *CB = dyn_cast<CallBase>(&Inst);
+            if (CB && !CB->getCalledFunction()) IndirCalls.push_back(&Inst);
+
+          }
+
+        }*/
+
   }
 
   InjectCoverage(F, BlocksToInstrument, IsLeafFunc);
@@ -1717,6 +1870,7 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
           stderr,
           "BUG! inst_in_this_func %u != afl_global_id %u - save_global %u\n",
           inst_in_this_func, afl_global_id, save_global);
+      exit(-1);
 
     }
 
@@ -1725,24 +1879,6 @@ void ModuleSanitizerCoverageLTO::instrumentFunction(
 
   }
 
-  /*
-    fprintf(stderr, "FUNCTION: %s [%u]\n", F.getName().str().c_str(),
-    extra_ctx_inst); 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";
-
-      }
-
-    }
-
-  */
-
 }
 
 GlobalVariable *ModuleSanitizerCoverageLTO::CreateFunctionLocalArrayInSection(
@@ -1867,7 +2003,6 @@ bool ModuleSanitizerCoverageLTO::InjectCoverage(
 
 }
 
-// Fake InjectCoverage to count the instrumentations
 bool ModuleSanitizerCoverageLTO::Fake_InjectCoverage(
     Function &F, ArrayRef<BasicBlock *> AllBlocks, bool IsLeafFunc) {
 
@@ -1875,7 +2010,6 @@ bool ModuleSanitizerCoverageLTO::Fake_InjectCoverage(
 
   for (size_t i = 0, N = AllBlocks.size(); i < N; i++) {
 
-    // AFL++ START
     if (BlockList.size()) {
 
       int skip = 0;
@@ -1904,6 +2038,7 @@ bool ModuleSanitizerCoverageLTO::Fake_InjectCoverage(
 //     The cache is used to speed up recording the caller-callee pairs.
 // The address of the caller is passed implicitly via caller PC.
 // CacheSize is encoded in the name of the run-time function.
+/*
 void ModuleSanitizerCoverageLTO::InjectCoverageForIndirectCalls(
     Function &F, ArrayRef<Instruction *> IndirCalls) {
 
@@ -1922,6 +2057,8 @@ void ModuleSanitizerCoverageLTO::InjectCoverageForIndirectCalls(
 
 }
 
+*/
+
 void ModuleSanitizerCoverageLTO::InjectCoverageAtBlock(Function   &F,
                                                        BasicBlock &BB,
                                                        size_t      Idx,
@@ -1968,13 +2105,19 @@ void ModuleSanitizerCoverageLTO::InjectCoverageAtBlock(Function   &F,
     /* Set the ID of the inserted basic block */
 
     ConstantInt *CurLoc = ConstantInt::get(Int32Tyi, afl_global_id);
+    Value       *val = CurLoc;
+
+    if (CTX_add) {
+
       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);
+      val = IRB.CreateAdd(CurLoc, CTX_load);
+
+    }
 
     /* Load SHM pointer */
 
@@ -1982,13 +2125,13 @@ void ModuleSanitizerCoverageLTO::InjectCoverageAtBlock(Function   &F,
 
     if (map_addr) {
 
-      MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtrFixed, offset);
+      MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtrFixed, val);
 
     } else {
 
       LoadInst *MapPtr = IRB.CreateLoad(PointerType::get(Int8Ty, 0), AFLMapPtr);
       ModuleSanitizerCoverageLTO::SetNoSanitizeMetadata(MapPtr);
-      MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, offset);
+      MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, val);
 
     }