xref: /aosp_15_r20/external/AFLplusplus/instrumentation/SanitizerCoveragePCGUARD.so.cc (revision 08b48e0b10e97b33e7b60c5b6e2243bd915777f2)
1 //===-- SanitizerCoverage.cpp - coverage instrumentation for sanitizers ---===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // Coverage instrumentation done on LLVM IR level, works with Sanitizers.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "llvm/Transforms/Instrumentation/SanitizerCoverage.h"
14 #include "llvm/ADT/ArrayRef.h"
15 #include "llvm/ADT/SmallVector.h"
16 #if LLVM_VERSION_MAJOR >= 15
17   #if LLVM_VERSION_MAJOR < 17
18     #include "llvm/ADT/Triple.h"
19   #endif
20 #endif
21 #include "llvm/Analysis/PostDominators.h"
22 #if LLVM_VERSION_MAJOR < 15
23   #include "llvm/IR/CFG.h"
24 #endif
25 #include "llvm/IR/Constant.h"
26 #include "llvm/IR/DataLayout.h"
27 #if LLVM_VERSION_MAJOR < 15
28   #include "llvm/IR/DebugInfo.h"
29 #endif
30 #include "llvm/IR/Dominators.h"
31 #if LLVM_VERSION_MAJOR >= 17
32   #include "llvm/IR/EHPersonalities.h"
33 #else
34   #include "llvm/Analysis/EHPersonalities.h"
35 #endif
36 #include "llvm/IR/Function.h"
37 #if LLVM_VERSION_MAJOR >= 16
38   #include "llvm/IR/GlobalVariable.h"
39 #endif
40 #include "llvm/IR/IRBuilder.h"
41 #if LLVM_VERSION_MAJOR < 15
42   #include "llvm/IR/InlineAsm.h"
43 #endif
44 #include "llvm/IR/IntrinsicInst.h"
45 #include "llvm/IR/Intrinsics.h"
46 #include "llvm/IR/LLVMContext.h"
47 #if LLVM_VERSION_MAJOR < 15
48   #include "llvm/IR/MDBuilder.h"
49   #include "llvm/IR/Mangler.h"
50 #endif
51 #include "llvm/IR/Module.h"
52 #include "llvm/IR/PassManager.h"
53 #include "llvm/Passes/PassBuilder.h"
54 #include "llvm/Passes/PassPlugin.h"
55 #include "llvm/IR/Type.h"
56 #if LLVM_VERSION_MAJOR < 17
57   #include "llvm/InitializePasses.h"
58 #endif
59 #include "llvm/Support/CommandLine.h"
60 #include "llvm/Support/Debug.h"
61 #include "llvm/Support/SpecialCaseList.h"
62 #include "llvm/Support/VirtualFileSystem.h"
63 #if LLVM_VERSION_MAJOR < 15
64   #include "llvm/Support/raw_ostream.h"
65 #endif
66 #if LLVM_VERSION_MAJOR < 17
67   #include "llvm/Transforms/Instrumentation.h"
68 #else
69   #include "llvm/TargetParser/Triple.h"
70 #endif
71 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
72 #include "llvm/Transforms/Utils/ModuleUtils.h"
73 
74 #include "config.h"
75 #include "debug.h"
76 #include "afl-llvm-common.h"
77 
78 using namespace llvm;
79 
80 #define DEBUG_TYPE "sancov"
81 
82 static const uint64_t SanCtorAndDtorPriority = 2;
83 
84 const char SanCovTracePCName[] = "__sanitizer_cov_trace_pc";
85 const char SanCovTraceCmp1[] = "__sanitizer_cov_trace_cmp1";
86 const char SanCovTraceCmp2[] = "__sanitizer_cov_trace_cmp2";
87 const char SanCovTraceCmp4[] = "__sanitizer_cov_trace_cmp4";
88 const char SanCovTraceCmp8[] = "__sanitizer_cov_trace_cmp8";
89 const char SanCovTraceConstCmp1[] = "__sanitizer_cov_trace_const_cmp1";
90 const char SanCovTraceConstCmp2[] = "__sanitizer_cov_trace_const_cmp2";
91 const char SanCovTraceConstCmp4[] = "__sanitizer_cov_trace_const_cmp4";
92 const char SanCovTraceConstCmp8[] = "__sanitizer_cov_trace_const_cmp8";
93 const char SanCovTraceSwitchName[] = "__sanitizer_cov_trace_switch";
94 
95 const char SanCovModuleCtorTracePcGuardName[] =
96     "sancov.module_ctor_trace_pc_guard";
97 const char SanCovTracePCGuardInitName[] = "__sanitizer_cov_trace_pc_guard_init";
98 
99 const char SanCovTracePCGuardName[] = "__sanitizer_cov_trace_pc_guard";
100 
101 const char SanCovGuardsSectionName[] = "sancov_guards";
102 const char SanCovCountersSectionName[] = "sancov_cntrs";
103 const char SanCovBoolFlagSectionName[] = "sancov_bools";
104 const char SanCovPCsSectionName[] = "sancov_pcs";
105 
106 const char SanCovLowestStackName[] = "__sancov_lowest_stack";
107 
108 static const char *skip_nozero;
109 static const char *use_threadsafe_counters;
110 
111 namespace {
112 
OverrideFromCL(SanitizerCoverageOptions Options)113 SanitizerCoverageOptions OverrideFromCL(SanitizerCoverageOptions Options) {
114 
115   Options.CoverageType = SanitizerCoverageOptions::SCK_Edge;
116   // Options.NoPrune = true;
117   Options.TracePCGuard = true;  // TracePCGuard is default.
118   return Options;
119 
120 }
121 
122 using DomTreeCallback = function_ref<const DominatorTree *(Function &F)>;
123 using PostDomTreeCallback =
124     function_ref<const PostDominatorTree *(Function &F)>;
125 
126 class ModuleSanitizerCoverageAFL
127     : public PassInfoMixin<ModuleSanitizerCoverageAFL> {
128 
129  public:
ModuleSanitizerCoverageAFL(const SanitizerCoverageOptions & Options=SanitizerCoverageOptions ())130   ModuleSanitizerCoverageAFL(
131       const SanitizerCoverageOptions &Options = SanitizerCoverageOptions())
132       : Options(OverrideFromCL(Options)) {
133 
134   }
135 
136   PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
137   bool              instrumentModule(Module &M, DomTreeCallback DTCallback,
138                                      PostDomTreeCallback PDTCallback);
139 
140  private:
141   void instrumentFunction(Function &F, DomTreeCallback DTCallback,
142                           PostDomTreeCallback PDTCallback);
143   void InjectTraceForCmp(Function &F, ArrayRef<Instruction *> CmpTraceTargets);
144   void InjectTraceForSwitch(Function               &F,
145                             ArrayRef<Instruction *> SwitchTraceTargets);
146   bool InjectCoverage(Function &F, ArrayRef<BasicBlock *> AllBlocks,
147                       bool IsLeafFunc = true);
148   GlobalVariable *CreateFunctionLocalArrayInSection(size_t    NumElements,
149                                                     Function &F, Type *Ty,
150                                                     const char *Section);
151   GlobalVariable *CreatePCArray(Function &F, ArrayRef<BasicBlock *> AllBlocks);
152   void CreateFunctionLocalArrays(Function &F, ArrayRef<BasicBlock *> AllBlocks,
153                                  uint32_t special);
154   void InjectCoverageAtBlock(Function &F, BasicBlock &BB, size_t Idx,
155                              bool IsLeafFunc = true);
156   Function *CreateInitCallsForSections(Module &M, const char *CtorName,
157                                        const char *InitFunctionName, Type *Ty,
158                                        const char *Section);
159   std::pair<Value *, Value *> CreateSecStartEnd(Module &M, const char *Section,
160                                                 Type *Ty);
161 
SetNoSanitizeMetadata(Instruction * I)162   void SetNoSanitizeMetadata(Instruction *I) {
163 
164 #if LLVM_VERSION_MAJOR >= 16
165     I->setMetadata(LLVMContext::MD_nosanitize, MDNode::get(*C, std::nullopt));
166 #else
167     I->setMetadata(I->getModule()->getMDKindID("nosanitize"),
168                    MDNode::get(*C, None));
169 #endif
170 
171   }
172 
173   std::string     getSectionName(const std::string &Section) const;
174   std::string     getSectionStart(const std::string &Section) const;
175   std::string     getSectionEnd(const std::string &Section) const;
176   FunctionCallee  SanCovTracePC, SanCovTracePCGuard;
177   FunctionCallee  SanCovTraceCmpFunction[4];
178   FunctionCallee  SanCovTraceConstCmpFunction[4];
179   FunctionCallee  SanCovTraceSwitchFunction;
180   GlobalVariable *SanCovLowestStack;
181   Type *IntptrTy, *IntptrPtrTy, *Int64Ty, *Int64PtrTy, *Int32Ty, *Int32PtrTy,
182       *Int16Ty, *Int8Ty, *Int8PtrTy, *Int1Ty, *Int1PtrTy;
183   Module           *CurModule;
184   std::string       CurModuleUniqueId;
185   Triple            TargetTriple;
186   LLVMContext      *C;
187   const DataLayout *DL;
188 
189   GlobalVariable *FunctionGuardArray;        // for trace-pc-guard.
190   GlobalVariable *Function8bitCounterArray;  // for inline-8bit-counters.
191   GlobalVariable *FunctionBoolArray;         // for inline-bool-flag.
192   GlobalVariable *FunctionPCsArray;          // for pc-table.
193   SmallVector<GlobalValue *, 20> GlobalsToAppendToUsed;
194   SmallVector<GlobalValue *, 20> GlobalsToAppendToCompilerUsed;
195 
196   SanitizerCoverageOptions Options;
197 
198   uint32_t        instr = 0, selects = 0, unhandled = 0;
199   GlobalVariable *AFLMapPtr = NULL;
200   ConstantInt    *One = NULL;
201   ConstantInt    *Zero = NULL;
202 
203 };
204 
205 }  // namespace
206 
207 extern "C" ::llvm::PassPluginLibraryInfo LLVM_ATTRIBUTE_WEAK
llvmGetPassPluginInfo()208 llvmGetPassPluginInfo() {
209 
210   return {LLVM_PLUGIN_API_VERSION, "SanitizerCoveragePCGUARD", "v0.2",
211           /* lambda to insert our pass into the pass pipeline. */
212           [](PassBuilder &PB) {
213 
214 #if LLVM_VERSION_MAJOR == 13
215             using OptimizationLevel = typename PassBuilder::OptimizationLevel;
216 #endif
217 #if LLVM_VERSION_MAJOR >= 16
218             PB.registerOptimizerEarlyEPCallback(
219 #else
220             PB.registerOptimizerLastEPCallback(
221 #endif
222                 [](ModulePassManager &MPM, OptimizationLevel OL) {
223 
224                   MPM.addPass(ModuleSanitizerCoverageAFL());
225 
226                 });
227 
228           }};
229 
230 }
231 
run(Module & M,ModuleAnalysisManager & MAM)232 PreservedAnalyses ModuleSanitizerCoverageAFL::run(Module                &M,
233                                                   ModuleAnalysisManager &MAM) {
234 
235   ModuleSanitizerCoverageAFL ModuleSancov(Options);
236   auto &FAM = MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
237   auto  DTCallback = [&FAM](Function &F) -> const DominatorTree  *{
238 
239     return &FAM.getResult<DominatorTreeAnalysis>(F);
240 
241   };
242 
243   auto PDTCallback = [&FAM](Function &F) -> const PostDominatorTree * {
244 
245     return &FAM.getResult<PostDominatorTreeAnalysis>(F);
246 
247   };
248 
249   if (ModuleSancov.instrumentModule(M, DTCallback, PDTCallback))
250     return PreservedAnalyses::none();
251   return PreservedAnalyses::all();
252 
253 }
254 
CreateSecStartEnd(Module & M,const char * Section,Type * Ty)255 std::pair<Value *, Value *> ModuleSanitizerCoverageAFL::CreateSecStartEnd(
256     Module &M, const char *Section, Type *Ty) {
257 
258   // Use ExternalWeak so that if all sections are discarded due to section
259   // garbage collection, the linker will not report undefined symbol errors.
260   // Windows defines the start/stop symbols in compiler-rt so no need for
261   // ExternalWeak.
262   GlobalValue::LinkageTypes Linkage = TargetTriple.isOSBinFormatCOFF()
263                                           ? GlobalVariable::ExternalLinkage
264                                           : GlobalVariable::ExternalWeakLinkage;
265   GlobalVariable *SecStart = new GlobalVariable(M, Ty, false, Linkage, nullptr,
266                                                 getSectionStart(Section));
267   SecStart->setVisibility(GlobalValue::HiddenVisibility);
268   GlobalVariable *SecEnd = new GlobalVariable(M, Ty, false, Linkage, nullptr,
269                                               getSectionEnd(Section));
270   SecEnd->setVisibility(GlobalValue::HiddenVisibility);
271   IRBuilder<> IRB(M.getContext());
272   if (!TargetTriple.isOSBinFormatCOFF())
273     return std::make_pair(SecStart, SecEnd);
274 
275   // Account for the fact that on windows-msvc __start_* symbols actually
276   // point to a uint64_t before the start of the array.
277   auto SecStartI8Ptr = IRB.CreatePointerCast(SecStart, Int8PtrTy);
278   auto GEP = IRB.CreateGEP(Int8Ty, SecStartI8Ptr,
279                            ConstantInt::get(IntptrTy, sizeof(uint64_t)));
280   return std::make_pair(IRB.CreatePointerCast(GEP, PointerType::getUnqual(Ty)),
281                         SecEnd);
282 
283 }
284 
CreateInitCallsForSections(Module & M,const char * CtorName,const char * InitFunctionName,Type * Ty,const char * Section)285 Function *ModuleSanitizerCoverageAFL::CreateInitCallsForSections(
286     Module &M, const char *CtorName, const char *InitFunctionName, Type *Ty,
287     const char *Section) {
288 
289   auto      SecStartEnd = CreateSecStartEnd(M, Section, Ty);
290   auto      SecStart = SecStartEnd.first;
291   auto      SecEnd = SecStartEnd.second;
292   Function *CtorFunc;
293   Type     *PtrTy = PointerType::getUnqual(Ty);
294   std::tie(CtorFunc, std::ignore) = createSanitizerCtorAndInitFunctions(
295       M, CtorName, InitFunctionName, {PtrTy, PtrTy}, {SecStart, SecEnd});
296   assert(CtorFunc->getName() == CtorName);
297 
298   if (TargetTriple.supportsCOMDAT()) {
299 
300     // Use comdat to dedup CtorFunc.
301     CtorFunc->setComdat(M.getOrInsertComdat(CtorName));
302     appendToGlobalCtors(M, CtorFunc, SanCtorAndDtorPriority, CtorFunc);
303 
304   } else {
305 
306     appendToGlobalCtors(M, CtorFunc, SanCtorAndDtorPriority);
307 
308   }
309 
310   if (TargetTriple.isOSBinFormatCOFF()) {
311 
312     // In COFF files, if the contructors are set as COMDAT (they are because
313     // COFF supports COMDAT) and the linker flag /OPT:REF (strip unreferenced
314     // functions and data) is used, the constructors get stripped. To prevent
315     // this, give the constructors weak ODR linkage and ensure the linker knows
316     // to include the sancov constructor. This way the linker can deduplicate
317     // the constructors but always leave one copy.
318     CtorFunc->setLinkage(GlobalValue::WeakODRLinkage);
319 
320   }
321 
322   return CtorFunc;
323 
324 }
325 
instrumentModule(Module & M,DomTreeCallback DTCallback,PostDomTreeCallback PDTCallback)326 bool ModuleSanitizerCoverageAFL::instrumentModule(
327     Module &M, DomTreeCallback DTCallback, PostDomTreeCallback PDTCallback) {
328 
329   setvbuf(stdout, NULL, _IONBF, 0);
330 
331   if (getenv("AFL_DEBUG")) { debug = 1; }
332 
333   if ((isatty(2) && !getenv("AFL_QUIET")) || debug) {
334 
335     SAYF(cCYA "SanitizerCoveragePCGUARD" VERSION cRST "\n");
336 
337   } else {
338 
339     be_quiet = 1;
340 
341   }
342 
343   skip_nozero = getenv("AFL_LLVM_SKIP_NEVERZERO");
344   use_threadsafe_counters = getenv("AFL_LLVM_THREADSAFE_INST");
345 
346   initInstrumentList();
347   scanForDangerousFunctions(&M);
348 
349   C = &(M.getContext());
350   DL = &M.getDataLayout();
351   CurModule = &M;
352   CurModuleUniqueId = getUniqueModuleId(CurModule);
353   TargetTriple = Triple(M.getTargetTriple());
354   FunctionGuardArray = nullptr;
355   Function8bitCounterArray = nullptr;
356   FunctionBoolArray = nullptr;
357   FunctionPCsArray = nullptr;
358   IntptrTy = Type::getIntNTy(*C, DL->getPointerSizeInBits());
359   IntptrPtrTy = PointerType::getUnqual(IntptrTy);
360   Type       *VoidTy = Type::getVoidTy(*C);
361   IRBuilder<> IRB(*C);
362   Int64PtrTy = PointerType::getUnqual(IRB.getInt64Ty());
363   Int32PtrTy = PointerType::getUnqual(IRB.getInt32Ty());
364   Int8PtrTy = PointerType::getUnqual(IRB.getInt8Ty());
365   Int1PtrTy = PointerType::getUnqual(IRB.getInt1Ty());
366   Int64Ty = IRB.getInt64Ty();
367   Int32Ty = IRB.getInt32Ty();
368   Int16Ty = IRB.getInt16Ty();
369   Int8Ty = IRB.getInt8Ty();
370   Int1Ty = IRB.getInt1Ty();
371 
372   LLVMContext &Ctx = M.getContext();
373   AFLMapPtr =
374       new GlobalVariable(M, PointerType::get(Int8Ty, 0), false,
375                          GlobalValue::ExternalLinkage, 0, "__afl_area_ptr");
376   One = ConstantInt::get(IntegerType::getInt8Ty(Ctx), 1);
377   Zero = ConstantInt::get(IntegerType::getInt8Ty(Ctx), 0);
378 
379   // Make sure smaller parameters are zero-extended to i64 if required by the
380   // target ABI.
381   AttributeList SanCovTraceCmpZeroExtAL;
382   SanCovTraceCmpZeroExtAL =
383       SanCovTraceCmpZeroExtAL.addParamAttribute(*C, 0, Attribute::ZExt);
384   SanCovTraceCmpZeroExtAL =
385       SanCovTraceCmpZeroExtAL.addParamAttribute(*C, 1, Attribute::ZExt);
386 
387   SanCovTraceCmpFunction[0] =
388       M.getOrInsertFunction(SanCovTraceCmp1, SanCovTraceCmpZeroExtAL, VoidTy,
389                             IRB.getInt8Ty(), IRB.getInt8Ty());
390   SanCovTraceCmpFunction[1] =
391       M.getOrInsertFunction(SanCovTraceCmp2, SanCovTraceCmpZeroExtAL, VoidTy,
392                             IRB.getInt16Ty(), IRB.getInt16Ty());
393   SanCovTraceCmpFunction[2] =
394       M.getOrInsertFunction(SanCovTraceCmp4, SanCovTraceCmpZeroExtAL, VoidTy,
395                             IRB.getInt32Ty(), IRB.getInt32Ty());
396   SanCovTraceCmpFunction[3] =
397       M.getOrInsertFunction(SanCovTraceCmp8, VoidTy, Int64Ty, Int64Ty);
398 
399   SanCovTraceConstCmpFunction[0] = M.getOrInsertFunction(
400       SanCovTraceConstCmp1, SanCovTraceCmpZeroExtAL, VoidTy, Int8Ty, Int8Ty);
401   SanCovTraceConstCmpFunction[1] = M.getOrInsertFunction(
402       SanCovTraceConstCmp2, SanCovTraceCmpZeroExtAL, VoidTy, Int16Ty, Int16Ty);
403   SanCovTraceConstCmpFunction[2] = M.getOrInsertFunction(
404       SanCovTraceConstCmp4, SanCovTraceCmpZeroExtAL, VoidTy, Int32Ty, Int32Ty);
405   SanCovTraceConstCmpFunction[3] =
406       M.getOrInsertFunction(SanCovTraceConstCmp8, VoidTy, Int64Ty, Int64Ty);
407 
408   SanCovTraceSwitchFunction =
409       M.getOrInsertFunction(SanCovTraceSwitchName, VoidTy, Int64Ty, Int64PtrTy);
410 
411   Constant *SanCovLowestStackConstant =
412       M.getOrInsertGlobal(SanCovLowestStackName, IntptrTy);
413   SanCovLowestStack = dyn_cast<GlobalVariable>(SanCovLowestStackConstant);
414   if (!SanCovLowestStack || SanCovLowestStack->getValueType() != IntptrTy) {
415 
416     C->emitError(StringRef("'") + SanCovLowestStackName +
417                  "' should not be declared by the user");
418     return true;
419 
420   }
421 
422   SanCovLowestStack->setThreadLocalMode(
423       GlobalValue::ThreadLocalMode::InitialExecTLSModel);
424 
425   SanCovTracePC = M.getOrInsertFunction(SanCovTracePCName, VoidTy);
426   SanCovTracePCGuard =
427       M.getOrInsertFunction(SanCovTracePCGuardName, VoidTy, Int32PtrTy);
428 
429   for (auto &F : M)
430     instrumentFunction(F, DTCallback, PDTCallback);
431 
432   Function *Ctor = nullptr;
433 
434   if (FunctionGuardArray)
435     Ctor = CreateInitCallsForSections(M, SanCovModuleCtorTracePcGuardName,
436                                       SanCovTracePCGuardInitName, Int32PtrTy,
437                                       SanCovGuardsSectionName);
438 
439   if (Ctor && debug) {
440 
441     fprintf(stderr, "SANCOV: installed pcguard_init in ctor\n");
442 
443   }
444 
445   appendToUsed(M, GlobalsToAppendToUsed);
446   appendToCompilerUsed(M, GlobalsToAppendToCompilerUsed);
447 
448   if (!be_quiet) {
449 
450     if (!instr) {
451 
452       WARNF("No instrumentation targets found.");
453 
454     } else {
455 
456       char modeline[128];
457       snprintf(modeline, sizeof(modeline), "%s%s%s%s%s%s",
458                getenv("AFL_HARDEN") ? "hardened" : "non-hardened",
459                getenv("AFL_USE_ASAN") ? ", ASAN" : "",
460                getenv("AFL_USE_MSAN") ? ", MSAN" : "",
461                getenv("AFL_USE_TSAN") ? ", TSAN" : "",
462                getenv("AFL_USE_CFISAN") ? ", CFISAN" : "",
463                getenv("AFL_USE_UBSAN") ? ", UBSAN" : "");
464       OKF("Instrumented %u locations with no collisions (%s mode) of which are "
465           "%u handled and %u unhandled selects.",
466           instr, modeline, selects, unhandled);
467 
468     }
469 
470   }
471 
472   return true;
473 
474 }
475 
476 // True if block has successors and it dominates all of them.
isFullDominator(const BasicBlock * BB,const DominatorTree * DT)477 static bool isFullDominator(const BasicBlock *BB, const DominatorTree *DT) {
478 
479   if (succ_empty(BB)) return false;
480 
481   return llvm::all_of(successors(BB), [&](const BasicBlock *SUCC) {
482 
483     return DT->dominates(BB, SUCC);
484 
485   });
486 
487 }
488 
489 // True if block has predecessors and it postdominates all of them.
isFullPostDominator(const BasicBlock * BB,const PostDominatorTree * PDT)490 static bool isFullPostDominator(const BasicBlock        *BB,
491                                 const PostDominatorTree *PDT) {
492 
493   if (pred_empty(BB)) return false;
494 
495   return llvm::all_of(predecessors(BB), [&](const BasicBlock *PRED) {
496 
497     return PDT->dominates(BB, PRED);
498 
499   });
500 
501 }
502 
shouldInstrumentBlock(const Function & F,const BasicBlock * BB,const DominatorTree * DT,const PostDominatorTree * PDT,const SanitizerCoverageOptions & Options)503 static bool shouldInstrumentBlock(const Function &F, const BasicBlock *BB,
504                                   const DominatorTree            *DT,
505                                   const PostDominatorTree        *PDT,
506                                   const SanitizerCoverageOptions &Options) {
507 
508   // Don't insert coverage for blocks containing nothing but unreachable: we
509   // will never call __sanitizer_cov() for them, so counting them in
510   // NumberOfInstrumentedBlocks() might complicate calculation of code coverage
511   // percentage. Also, unreachable instructions frequently have no debug
512   // locations.
513   if (isa<UnreachableInst>(BB->getFirstNonPHIOrDbgOrLifetime())) return false;
514 
515   // Don't insert coverage into blocks without a valid insertion point
516   // (catchswitch blocks).
517   if (BB->getFirstInsertionPt() == BB->end()) return false;
518 
519   if (Options.NoPrune || &F.getEntryBlock() == BB) return true;
520 
521   // Do not instrument full dominators, or full post-dominators with multiple
522   // predecessors.
523   return !isFullDominator(BB, DT) &&
524          !(isFullPostDominator(BB, PDT) && !BB->getSinglePredecessor());
525 
526 }
527 
528 // Returns true iff From->To is a backedge.
529 // A twist here is that we treat From->To as a backedge if
530 //   * To dominates From or
531 //   * To->UniqueSuccessor dominates From
532 #if 0
533 static bool IsBackEdge(BasicBlock *From, BasicBlock *To,
534                        const DominatorTree *DT) {
535 
536   if (DT->dominates(To, From))
537     return true;
538   if (auto Next = To->getUniqueSuccessor())
539     if (DT->dominates(Next, From))
540       return true;
541   return false;
542 
543 }
544 
545 #endif
546 
547 // Prunes uninteresting Cmp instrumentation:
548 //   * CMP instructions that feed into loop backedge branch.
549 //
550 // Note that Cmp pruning is controlled by the same flag as the
551 // BB pruning.
552 #if 0
553 static bool IsInterestingCmp(ICmpInst *CMP, const DominatorTree *DT,
554                              const SanitizerCoverageOptions &Options) {
555 
556   if (!Options.NoPrune)
557     if (CMP->hasOneUse())
558       if (auto BR = dyn_cast<BranchInst>(CMP->user_back()))
559         for (BasicBlock *B : BR->successors())
560           if (IsBackEdge(BR->getParent(), B, DT))
561             return false;
562   return true;
563 
564 }
565 
566 #endif
567 
instrumentFunction(Function & F,DomTreeCallback DTCallback,PostDomTreeCallback PDTCallback)568 void ModuleSanitizerCoverageAFL::instrumentFunction(
569     Function &F, DomTreeCallback DTCallback, PostDomTreeCallback PDTCallback) {
570 
571   if (F.empty()) return;
572   if (!isInInstrumentList(&F, FMNAME)) return;
573   if (F.getName().find(".module_ctor") != std::string::npos)
574     return;  // Should not instrument sanitizer init functions.
575 #if LLVM_VERSION_MAJOR >= 18
576   if (F.getName().starts_with("__sanitizer_"))
577 #else
578   if (F.getName().startswith("__sanitizer_"))
579 #endif
580     return;  // Don't instrument __sanitizer_* callbacks.
581   // Don't touch available_externally functions, their actual body is elewhere.
582   if (F.getLinkage() == GlobalValue::AvailableExternallyLinkage) return;
583   // Don't instrument MSVC CRT configuration helpers. They may run before normal
584   // initialization.
585   if (F.getName() == "__local_stdio_printf_options" ||
586       F.getName() == "__local_stdio_scanf_options")
587     return;
588   if (isa<UnreachableInst>(F.getEntryBlock().getTerminator())) return;
589   // Don't instrument functions using SEH for now. Splitting basic blocks like
590   // we do for coverage breaks WinEHPrepare.
591   // FIXME: Remove this when SEH no longer uses landingpad pattern matching.
592   if (F.hasPersonalityFn() &&
593       isAsynchronousEHPersonality(classifyEHPersonality(F.getPersonalityFn())))
594     return;
595   if (F.hasFnAttribute(Attribute::NoSanitizeCoverage)) return;
596   if (Options.CoverageType >= SanitizerCoverageOptions::SCK_Edge)
597     SplitAllCriticalEdges(
598         F, CriticalEdgeSplittingOptions().setIgnoreUnreachableDests());
599   SmallVector<BasicBlock *, 16> BlocksToInstrument;
600   SmallVector<Instruction *, 8> CmpTraceTargets;
601   SmallVector<Instruction *, 8> SwitchTraceTargets;
602 
603   const DominatorTree     *DT = DTCallback(F);
604   const PostDominatorTree *PDT = PDTCallback(F);
605   bool                     IsLeafFunc = true;
606 
607   for (auto &BB : F) {
608 
609     if (shouldInstrumentBlock(F, &BB, DT, PDT, Options))
610       BlocksToInstrument.push_back(&BB);
611     /*
612         for (auto &Inst : BB) {
613 
614           if (Options.TraceCmp) {
615 
616             if (ICmpInst *CMP = dyn_cast<ICmpInst>(&Inst))
617               if (IsInterestingCmp(CMP, DT, Options))
618                 CmpTraceTargets.push_back(&Inst);
619             if (isa<SwitchInst>(&Inst))
620               SwitchTraceTargets.push_back(&Inst);
621 
622           }
623 
624         }
625 
626     */
627 
628   }
629 
630   if (debug) {
631 
632     fprintf(stderr, "SanitizerCoveragePCGUARD: instrumenting %s in %s\n",
633             F.getName().str().c_str(), F.getParent()->getName().str().c_str());
634 
635   }
636 
637   InjectCoverage(F, BlocksToInstrument, IsLeafFunc);
638   // InjectTraceForCmp(F, CmpTraceTargets);
639   // InjectTraceForSwitch(F, SwitchTraceTargets);
640 
641 }
642 
CreateFunctionLocalArrayInSection(size_t NumElements,Function & F,Type * Ty,const char * Section)643 GlobalVariable *ModuleSanitizerCoverageAFL::CreateFunctionLocalArrayInSection(
644     size_t NumElements, Function &F, Type *Ty, const char *Section) {
645 
646   ArrayType *ArrayTy = ArrayType::get(Ty, NumElements);
647   auto       Array = new GlobalVariable(
648       *CurModule, ArrayTy, false, GlobalVariable::PrivateLinkage,
649       Constant::getNullValue(ArrayTy), "__sancov_gen_");
650 
651   if (TargetTriple.supportsCOMDAT() &&
652       (TargetTriple.isOSBinFormatELF() || !F.isInterposable()))
653     if (auto Comdat = getOrCreateFunctionComdat(F, TargetTriple))
654       Array->setComdat(Comdat);
655   Array->setSection(getSectionName(Section));
656 #if LLVM_VERSION_MAJOR >= 16
657   Array->setAlignment(Align(DL->getTypeStoreSize(Ty).getFixedValue()));
658 #else
659   Array->setAlignment(Align(DL->getTypeStoreSize(Ty).getFixedSize()));
660 #endif
661 
662   // sancov_pcs parallels the other metadata section(s). Optimizers (e.g.
663   // GlobalOpt/ConstantMerge) may not discard sancov_pcs and the other
664   // section(s) as a unit, so we conservatively retain all unconditionally in
665   // the compiler.
666   //
667   // With comdat (COFF/ELF), the linker can guarantee the associated sections
668   // will be retained or discarded as a unit, so llvm.compiler.used is
669   // sufficient. Otherwise, conservatively make all of them retained by the
670   // linker.
671   if (Array->hasComdat())
672     GlobalsToAppendToCompilerUsed.push_back(Array);
673   else
674     GlobalsToAppendToUsed.push_back(Array);
675 
676   return Array;
677 
678 }
679 
CreatePCArray(Function & F,ArrayRef<BasicBlock * > AllBlocks)680 GlobalVariable *ModuleSanitizerCoverageAFL::CreatePCArray(
681     Function &F, ArrayRef<BasicBlock *> AllBlocks) {
682 
683   size_t N = AllBlocks.size();
684   assert(N);
685   SmallVector<Constant *, 32> PCs;
686   IRBuilder<>                 IRB(&*F.getEntryBlock().getFirstInsertionPt());
687   for (size_t i = 0; i < N; i++) {
688 
689     if (&F.getEntryBlock() == AllBlocks[i]) {
690 
691       PCs.push_back((Constant *)IRB.CreatePointerCast(&F, IntptrPtrTy));
692       PCs.push_back((Constant *)IRB.CreateIntToPtr(
693           ConstantInt::get(IntptrTy, 1), IntptrPtrTy));
694 
695     } else {
696 
697       PCs.push_back((Constant *)IRB.CreatePointerCast(
698           BlockAddress::get(AllBlocks[i]), IntptrPtrTy));
699 #if LLVM_VERSION_MAJOR >= 16
700       PCs.push_back(Constant::getNullValue(IntptrPtrTy));
701 #else
702       PCs.push_back((Constant *)IRB.CreateIntToPtr(
703           ConstantInt::get(IntptrTy, 0), IntptrPtrTy));
704 #endif
705 
706     }
707 
708   }
709 
710   auto *PCArray = CreateFunctionLocalArrayInSection(N * 2, F, IntptrPtrTy,
711                                                     SanCovPCsSectionName);
712   PCArray->setInitializer(
713       ConstantArray::get(ArrayType::get(IntptrPtrTy, N * 2), PCs));
714   PCArray->setConstant(true);
715 
716   return PCArray;
717 
718 }
719 
CreateFunctionLocalArrays(Function & F,ArrayRef<BasicBlock * > AllBlocks,uint32_t special)720 void ModuleSanitizerCoverageAFL::CreateFunctionLocalArrays(
721     Function &F, ArrayRef<BasicBlock *> AllBlocks, uint32_t special) {
722 
723   if (Options.TracePCGuard)
724     FunctionGuardArray = CreateFunctionLocalArrayInSection(
725         AllBlocks.size() + special, F, Int32Ty, SanCovGuardsSectionName);
726 
727 }
728 
InjectCoverage(Function & F,ArrayRef<BasicBlock * > AllBlocks,bool IsLeafFunc)729 bool ModuleSanitizerCoverageAFL::InjectCoverage(
730     Function &F, ArrayRef<BasicBlock *> AllBlocks, bool IsLeafFunc) {
731 
732   if (AllBlocks.empty()) return false;
733 
734   uint32_t        cnt_cov = 0, cnt_sel = 0, cnt_sel_inc = 0;
735   static uint32_t first = 1;
736 
737   for (auto &BB : F) {
738 
739     for (auto &IN : BB) {
740 
741       CallInst *callInst = nullptr;
742 
743       if ((callInst = dyn_cast<CallInst>(&IN))) {
744 
745         Function *Callee = callInst->getCalledFunction();
746         if (!Callee) continue;
747         if (callInst->getCallingConv() != llvm::CallingConv::C) continue;
748         StringRef FuncName = Callee->getName();
749         if (!FuncName.compare(StringRef("dlopen")) ||
750             !FuncName.compare(StringRef("_dlopen"))) {
751 
752           fprintf(stderr,
753                   "WARNING: dlopen() detected. To have coverage for a library "
754                   "that your target dlopen()'s this must either happen before "
755                   "__AFL_INIT() or you must use AFL_PRELOAD to preload all "
756                   "dlopen()'ed libraries!\n");
757           continue;
758 
759         }
760 
761         if (!FuncName.compare(StringRef("__afl_coverage_interesting"))) {
762 
763           cnt_cov++;
764 
765         }
766 
767       }
768 
769       SelectInst *selectInst = nullptr;
770 
771       if ((selectInst = dyn_cast<SelectInst>(&IN))) {
772 
773         Value *c = selectInst->getCondition();
774         auto   t = c->getType();
775         if (t->getTypeID() == llvm::Type::IntegerTyID) {
776 
777           cnt_sel++;
778           cnt_sel_inc += 2;
779 
780         }
781 
782         else if (t->getTypeID() == llvm::Type::FixedVectorTyID) {
783 
784           FixedVectorType *tt = dyn_cast<FixedVectorType>(t);
785           if (tt) {
786 
787             cnt_sel++;
788             cnt_sel_inc += (tt->getElementCount().getKnownMinValue() * 2);
789 
790           }
791 
792         }
793 
794       }
795 
796     }
797 
798   }
799 
800   CreateFunctionLocalArrays(F, AllBlocks, first + cnt_cov + cnt_sel_inc);
801 
802   if (first) { first = 0; }
803   selects += cnt_sel;
804 
805   uint32_t special = 0, local_selects = 0, skip_next = 0;
806 
807   for (auto &BB : F) {
808 
809     for (auto &IN : BB) {
810 
811       CallInst *callInst = nullptr;
812 
813       if ((callInst = dyn_cast<CallInst>(&IN))) {
814 
815         Function *Callee = callInst->getCalledFunction();
816         if (!Callee) continue;
817         if (callInst->getCallingConv() != llvm::CallingConv::C) continue;
818         StringRef FuncName = Callee->getName();
819         if (FuncName.compare(StringRef("__afl_coverage_interesting"))) continue;
820 
821         IRBuilder<> IRB(callInst);
822 
823         if (!FunctionGuardArray) {
824 
825           fprintf(stderr,
826                   "SANCOV: FunctionGuardArray is NULL, failed to emit "
827                   "instrumentation.");
828           continue;
829 
830         }
831 
832         Value *GuardPtr = IRB.CreateIntToPtr(
833             IRB.CreateAdd(
834                 IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
835                 ConstantInt::get(IntptrTy, (++special + AllBlocks.size()) * 4)),
836             Int32PtrTy);
837 
838         LoadInst *Idx = IRB.CreateLoad(IRB.getInt32Ty(), GuardPtr);
839         ModuleSanitizerCoverageAFL::SetNoSanitizeMetadata(Idx);
840 
841         callInst->setOperand(1, Idx);
842 
843       }
844 
845       SelectInst *selectInst = nullptr;
846 
847       if (!skip_next && (selectInst = dyn_cast<SelectInst>(&IN))) {
848 
849         uint32_t    vector_cnt = 0;
850         Value      *condition = selectInst->getCondition();
851         Value      *result;
852         auto        t = condition->getType();
853         IRBuilder<> IRB(selectInst->getNextNode());
854 
855         if (t->getTypeID() == llvm::Type::IntegerTyID) {
856 
857           if (!FunctionGuardArray) {
858 
859             fprintf(stderr,
860                     "SANCOV: FunctionGuardArray is NULL, failed to emit "
861                     "instrumentation.");
862             continue;
863 
864           }
865 
866           auto GuardPtr1 = IRB.CreateIntToPtr(
867               IRB.CreateAdd(
868                   IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
869                   ConstantInt::get(
870                       IntptrTy,
871                       (cnt_cov + local_selects++ + AllBlocks.size()) * 4)),
872               Int32PtrTy);
873 
874           auto GuardPtr2 = IRB.CreateIntToPtr(
875               IRB.CreateAdd(
876                   IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
877                   ConstantInt::get(
878                       IntptrTy,
879                       (cnt_cov + local_selects++ + AllBlocks.size()) * 4)),
880               Int32PtrTy);
881 
882           result = IRB.CreateSelect(condition, GuardPtr1, GuardPtr2);
883 
884         } else
885 
886 #if LLVM_VERSION_MAJOR >= 14
887             if (t->getTypeID() == llvm::Type::FixedVectorTyID) {
888 
889           FixedVectorType *tt = dyn_cast<FixedVectorType>(t);
890           if (tt) {
891 
892             uint32_t elements = tt->getElementCount().getFixedValue();
893             vector_cnt = elements;
894             if (elements) {
895 
896               FixedVectorType *GuardPtr1 =
897                   FixedVectorType::get(Int32PtrTy, elements);
898               FixedVectorType *GuardPtr2 =
899                   FixedVectorType::get(Int32PtrTy, elements);
900               Value *x, *y;
901 
902               if (!FunctionGuardArray) {
903 
904                 fprintf(stderr,
905                         "SANCOV: FunctionGuardArray is NULL, failed to emit "
906                         "instrumentation.");
907                 continue;
908 
909               }
910 
911               Value *val1 = IRB.CreateIntToPtr(
912                   IRB.CreateAdd(
913                       IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
914                       ConstantInt::get(
915                           IntptrTy,
916                           (cnt_cov + local_selects++ + AllBlocks.size()) * 4)),
917                   Int32PtrTy);
918               x = IRB.CreateInsertElement(GuardPtr1, val1, (uint64_t)0);
919 
920               Value *val2 = IRB.CreateIntToPtr(
921                   IRB.CreateAdd(
922                       IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
923                       ConstantInt::get(
924                           IntptrTy,
925                           (cnt_cov + local_selects++ + AllBlocks.size()) * 4)),
926                   Int32PtrTy);
927               y = IRB.CreateInsertElement(GuardPtr2, val2, (uint64_t)0);
928 
929               for (uint64_t i = 1; i < elements; i++) {
930 
931                 val1 = IRB.CreateIntToPtr(
932                     IRB.CreateAdd(
933                         IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
934                         ConstantInt::get(IntptrTy, (cnt_cov + local_selects++ +
935                                                     AllBlocks.size()) *
936                                                        4)),
937                     Int32PtrTy);
938                 x = IRB.CreateInsertElement(x, val1, i);
939 
940                 val2 = IRB.CreateIntToPtr(
941                     IRB.CreateAdd(
942                         IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
943                         ConstantInt::get(IntptrTy, (cnt_cov + local_selects++ +
944                                                     AllBlocks.size()) *
945                                                        4)),
946                     Int32PtrTy);
947                 y = IRB.CreateInsertElement(y, val2, i);
948 
949               }
950 
951               result = IRB.CreateSelect(condition, x, y);
952 
953             }
954 
955           }
956 
957         } else
958 
959 #endif
960         {
961 
962           // fprintf(stderr, "UNHANDLED: %u\n", t->getTypeID());
963           unhandled++;
964           continue;
965 
966         }
967 
968         uint32_t vector_cur = 0;
969 
970         /* Load SHM pointer */
971 
972         LoadInst *MapPtr =
973             IRB.CreateLoad(PointerType::get(Int8Ty, 0), AFLMapPtr);
974         ModuleSanitizerCoverageAFL::SetNoSanitizeMetadata(MapPtr);
975 
976         while (1) {
977 
978           /* Get CurLoc */
979           LoadInst *CurLoc = nullptr;
980           Value    *MapPtrIdx = nullptr;
981 
982           /* Load counter for CurLoc */
983           if (!vector_cnt) {
984 
985             CurLoc = IRB.CreateLoad(IRB.getInt32Ty(), result);
986             ModuleSanitizerCoverageAFL::SetNoSanitizeMetadata(CurLoc);
987             MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, CurLoc);
988 
989           } else {
990 
991             auto element = IRB.CreateExtractElement(result, vector_cur++);
992             auto elementptr = IRB.CreateIntToPtr(element, Int32PtrTy);
993             auto elementld = IRB.CreateLoad(IRB.getInt32Ty(), elementptr);
994             ModuleSanitizerCoverageAFL::SetNoSanitizeMetadata(elementld);
995             MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, elementld);
996 
997           }
998 
999           if (use_threadsafe_counters) {
1000 
1001             IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One,
1002 #if LLVM_VERSION_MAJOR >= 13
1003                                 llvm::MaybeAlign(1),
1004 #endif
1005                                 llvm::AtomicOrdering::Monotonic);
1006 
1007           } else {
1008 
1009             LoadInst *Counter = IRB.CreateLoad(IRB.getInt8Ty(), MapPtrIdx);
1010             ModuleSanitizerCoverageAFL::SetNoSanitizeMetadata(Counter);
1011 
1012             /* Update bitmap */
1013 
1014             Value *Incr = IRB.CreateAdd(Counter, One);
1015 
1016             if (skip_nozero == NULL) {
1017 
1018               auto cf = IRB.CreateICmpEQ(Incr, Zero);
1019               auto carry = IRB.CreateZExt(cf, Int8Ty);
1020               Incr = IRB.CreateAdd(Incr, carry);
1021 
1022             }
1023 
1024             StoreInst *StoreCtx = IRB.CreateStore(Incr, MapPtrIdx);
1025             ModuleSanitizerCoverageAFL::SetNoSanitizeMetadata(StoreCtx);
1026 
1027           }
1028 
1029           if (!vector_cnt) {
1030 
1031             vector_cnt = 2;
1032             break;
1033 
1034           } else if (vector_cnt == vector_cur) {
1035 
1036             break;
1037 
1038           }
1039 
1040         }
1041 
1042         skip_next = 1;
1043         instr += vector_cnt;
1044 
1045       } else {
1046 
1047         skip_next = 0;
1048 
1049       }
1050 
1051     }
1052 
1053   }
1054 
1055   if (AllBlocks.empty() && !special && !local_selects) return false;
1056 
1057   if (!AllBlocks.empty())
1058     for (size_t i = 0, N = AllBlocks.size(); i < N; i++)
1059       InjectCoverageAtBlock(F, *AllBlocks[i], i, IsLeafFunc);
1060 
1061   return true;
1062 
1063 }
1064 
1065 // For every switch statement we insert a call:
1066 // __sanitizer_cov_trace_switch(CondValue,
1067 //      {NumCases, ValueSizeInBits, Case0Value, Case1Value, Case2Value, ... })
1068 
InjectTraceForSwitch(Function &,ArrayRef<Instruction * > SwitchTraceTargets)1069 void ModuleSanitizerCoverageAFL::InjectTraceForSwitch(
1070     Function &, ArrayRef<Instruction *> SwitchTraceTargets) {
1071 
1072   for (auto I : SwitchTraceTargets) {
1073 
1074     if (SwitchInst *SI = dyn_cast<SwitchInst>(I)) {
1075 
1076       IRBuilder<>                 IRB(I);
1077       SmallVector<Constant *, 16> Initializers;
1078       Value                      *Cond = SI->getCondition();
1079       if (Cond->getType()->getScalarSizeInBits() >
1080           Int64Ty->getScalarSizeInBits())
1081         continue;
1082       Initializers.push_back(ConstantInt::get(Int64Ty, SI->getNumCases()));
1083       Initializers.push_back(
1084           ConstantInt::get(Int64Ty, Cond->getType()->getScalarSizeInBits()));
1085       if (Cond->getType()->getScalarSizeInBits() <
1086           Int64Ty->getScalarSizeInBits())
1087         Cond = IRB.CreateIntCast(Cond, Int64Ty, false);
1088       for (auto It : SI->cases()) {
1089 
1090         Constant *C = It.getCaseValue();
1091         if (C->getType()->getScalarSizeInBits() <
1092             Int64Ty->getScalarSizeInBits())
1093           C = ConstantExpr::getCast(CastInst::ZExt, It.getCaseValue(), Int64Ty);
1094         Initializers.push_back(C);
1095 
1096       }
1097 
1098       llvm::sort(drop_begin(Initializers, 2),
1099                  [](const Constant *A, const Constant *B) {
1100 
1101                    return cast<ConstantInt>(A)->getLimitedValue() <
1102                           cast<ConstantInt>(B)->getLimitedValue();
1103 
1104                  });
1105 
1106       ArrayType *ArrayOfInt64Ty = ArrayType::get(Int64Ty, Initializers.size());
1107       GlobalVariable *GV = new GlobalVariable(
1108           *CurModule, ArrayOfInt64Ty, false, GlobalVariable::InternalLinkage,
1109           ConstantArray::get(ArrayOfInt64Ty, Initializers),
1110           "__sancov_gen_cov_switch_values");
1111       IRB.CreateCall(SanCovTraceSwitchFunction,
1112                      {Cond, IRB.CreatePointerCast(GV, Int64PtrTy)});
1113 
1114     }
1115 
1116   }
1117 
1118 }
1119 
InjectTraceForCmp(Function &,ArrayRef<Instruction * > CmpTraceTargets)1120 void ModuleSanitizerCoverageAFL::InjectTraceForCmp(
1121     Function &, ArrayRef<Instruction *> CmpTraceTargets) {
1122 
1123   for (auto I : CmpTraceTargets) {
1124 
1125     if (ICmpInst *ICMP = dyn_cast<ICmpInst>(I)) {
1126 
1127       IRBuilder<> IRB(ICMP);
1128       Value      *A0 = ICMP->getOperand(0);
1129       Value      *A1 = ICMP->getOperand(1);
1130       if (!A0->getType()->isIntegerTy()) continue;
1131       uint64_t TypeSize = DL->getTypeStoreSizeInBits(A0->getType());
1132       int      CallbackIdx = TypeSize == 8    ? 0
1133                              : TypeSize == 16 ? 1
1134                              : TypeSize == 32 ? 2
1135                              : TypeSize == 64 ? 3
1136                                               : -1;
1137       if (CallbackIdx < 0) continue;
1138       // __sanitizer_cov_trace_cmp((type_size << 32) | predicate, A0, A1);
1139       auto CallbackFunc = SanCovTraceCmpFunction[CallbackIdx];
1140       bool FirstIsConst = isa<ConstantInt>(A0);
1141       bool SecondIsConst = isa<ConstantInt>(A1);
1142       // If both are const, then we don't need such a comparison.
1143       if (FirstIsConst && SecondIsConst) continue;
1144       // If only one is const, then make it the first callback argument.
1145       if (FirstIsConst || SecondIsConst) {
1146 
1147         CallbackFunc = SanCovTraceConstCmpFunction[CallbackIdx];
1148         if (SecondIsConst) std::swap(A0, A1);
1149 
1150       }
1151 
1152       auto Ty = Type::getIntNTy(*C, TypeSize);
1153       IRB.CreateCall(CallbackFunc, {IRB.CreateIntCast(A0, Ty, true),
1154                                     IRB.CreateIntCast(A1, Ty, true)});
1155 
1156     }
1157 
1158   }
1159 
1160 }
1161 
InjectCoverageAtBlock(Function & F,BasicBlock & BB,size_t Idx,bool IsLeafFunc)1162 void ModuleSanitizerCoverageAFL::InjectCoverageAtBlock(Function   &F,
1163                                                        BasicBlock &BB,
1164                                                        size_t      Idx,
1165                                                        bool        IsLeafFunc) {
1166 
1167   BasicBlock::iterator IP = BB.getFirstInsertionPt();
1168   bool                 IsEntryBB = &BB == &F.getEntryBlock();
1169   DebugLoc             EntryLoc;
1170 
1171   if (IsEntryBB) {
1172 
1173     if (auto SP = F.getSubprogram())
1174       EntryLoc = DILocation::get(SP->getContext(), SP->getScopeLine(), 0, SP);
1175     // Keep static allocas and llvm.localescape calls in the entry block.  Even
1176     // if we aren't splitting the block, it's nice for allocas to be before
1177     // calls.
1178     IP = PrepareToSplitEntryBlock(BB, IP);
1179 #if LLVM_VERSION_MAJOR < 15
1180 
1181   } else {
1182 
1183     EntryLoc = IP->getDebugLoc();
1184     if (!EntryLoc)
1185       if (auto *SP = F.getSubprogram())
1186         EntryLoc = DILocation::get(SP->getContext(), 0, 0, SP);
1187 #endif
1188 
1189   }
1190 
1191 #if LLVM_VERSION_MAJOR >= 16
1192   InstrumentationIRBuilder IRB(&*IP);
1193 #else
1194   IRBuilder<> IRB(&*IP);
1195 #endif
1196   if (EntryLoc) IRB.SetCurrentDebugLocation(EntryLoc);
1197   if (Options.TracePCGuard) {
1198 
1199     /*
1200       auto GuardPtr = IRB.CreateIntToPtr(
1201           IRB.CreateAdd(IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
1202                         ConstantInt::get(IntptrTy, Idx * 4)),
1203           Int32PtrTy);
1204       IRB.CreateCall(SanCovTracePCGuard, GuardPtr)->setCannotMerge();
1205     */
1206 
1207     /* Get CurLoc */
1208 
1209     Value *GuardPtr = IRB.CreateIntToPtr(
1210         IRB.CreateAdd(IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
1211                       ConstantInt::get(IntptrTy, Idx * 4)),
1212         Int32PtrTy);
1213 
1214     LoadInst *CurLoc = IRB.CreateLoad(IRB.getInt32Ty(), GuardPtr);
1215     ModuleSanitizerCoverageAFL::SetNoSanitizeMetadata(CurLoc);
1216 
1217     /* Load SHM pointer */
1218 
1219     LoadInst *MapPtr = IRB.CreateLoad(PointerType::get(Int8Ty, 0), AFLMapPtr);
1220     ModuleSanitizerCoverageAFL::SetNoSanitizeMetadata(MapPtr);
1221 
1222     /* Load counter for CurLoc */
1223 
1224     Value *MapPtrIdx = IRB.CreateGEP(Int8Ty, MapPtr, CurLoc);
1225 
1226     if (use_threadsafe_counters) {
1227 
1228       IRB.CreateAtomicRMW(llvm::AtomicRMWInst::BinOp::Add, MapPtrIdx, One,
1229 #if LLVM_VERSION_MAJOR >= 13
1230                           llvm::MaybeAlign(1),
1231 #endif
1232                           llvm::AtomicOrdering::Monotonic);
1233 
1234     } else {
1235 
1236       LoadInst *Counter = IRB.CreateLoad(IRB.getInt8Ty(), MapPtrIdx);
1237       ModuleSanitizerCoverageAFL::SetNoSanitizeMetadata(Counter);
1238 
1239       /* Update bitmap */
1240 
1241       Value *Incr = IRB.CreateAdd(Counter, One);
1242 
1243       if (skip_nozero == NULL) {
1244 
1245         auto cf = IRB.CreateICmpEQ(Incr, Zero);
1246         auto carry = IRB.CreateZExt(cf, Int8Ty);
1247         Incr = IRB.CreateAdd(Incr, carry);
1248 
1249       }
1250 
1251       StoreInst *StoreCtx = IRB.CreateStore(Incr, MapPtrIdx);
1252       ModuleSanitizerCoverageAFL::SetNoSanitizeMetadata(StoreCtx);
1253 
1254     }
1255 
1256     // done :)
1257 
1258     //    IRB.CreateCall(SanCovTracePCGuard, Offset)->setCannotMerge();
1259     //    IRB.CreateCall(SanCovTracePCGuard, GuardPtr)->setCannotMerge();
1260     ++instr;
1261 
1262   }
1263 
1264 }
1265 
getSectionName(const std::string & Section) const1266 std::string ModuleSanitizerCoverageAFL::getSectionName(
1267     const std::string &Section) const {
1268 
1269   if (TargetTriple.isOSBinFormatCOFF()) {
1270 
1271     if (Section == SanCovCountersSectionName) return ".SCOV$CM";
1272     if (Section == SanCovBoolFlagSectionName) return ".SCOV$BM";
1273     if (Section == SanCovPCsSectionName) return ".SCOVP$M";
1274     return ".SCOV$GM";  // For SanCovGuardsSectionName.
1275 
1276   }
1277 
1278   if (TargetTriple.isOSBinFormatMachO()) return "__DATA,__" + Section;
1279   return "__" + Section;
1280 
1281 }
1282 
getSectionStart(const std::string & Section) const1283 std::string ModuleSanitizerCoverageAFL::getSectionStart(
1284     const std::string &Section) const {
1285 
1286   if (TargetTriple.isOSBinFormatMachO())
1287     return "\1section$start$__DATA$__" + Section;
1288   return "__start___" + Section;
1289 
1290 }
1291 
getSectionEnd(const std::string & Section) const1292 std::string ModuleSanitizerCoverageAFL::getSectionEnd(
1293     const std::string &Section) const {
1294 
1295   if (TargetTriple.isOSBinFormatMachO())
1296     return "\1section$end$__DATA$__" + Section;
1297   return "__stop___" + Section;
1298 
1299 }
1300 
1301