add loop analysis to CC

vanhauser-thc · vanhauser-thc · commit 31a7ff2ba236 · 2024-05-15T15:28:03.000+02:00
diff --git a/instrumentation/SanitizerCoveragePCGUARD.so.cc b/instrumentation/SanitizerCoveragePCGUARD.so.cc
@@ -70,6 +70,8 @@
 #endif
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/LoopPass.h"
 
 #include "config.h"
 #include "debug.h"
@@ -119,6 +121,7 @@ SanitizerCoverageOptions OverrideFromCL(SanitizerCoverageOptions Options) {
 
 }
 
+using LoopInfoCallback = function_ref<const LoopInfo *(Function &F)>;
 using DomTreeCallback = function_ref<const DominatorTree *(Function &F)>;
 using PostDomTreeCallback =
     function_ref<const PostDominatorTree *(Function &F)>;
@@ -135,11 +138,13 @@ class ModuleSanitizerCoverageAFL
 
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
   bool              instrumentModule(Module &M, DomTreeCallback DTCallback,
-                                     PostDomTreeCallback PDTCallback);
+                                     PostDomTreeCallback PDTCallback,
+                                     LoopInfoCallback    LCallback);
 
  private:
   void instrumentFunction(Function &F, DomTreeCallback DTCallback,
-                          PostDomTreeCallback PDTCallback);
+                          PostDomTreeCallback PDTCallback,
+                          LoopInfoCallback    LCallback);
   void InjectTraceForCmp(Function &F, ArrayRef<Instruction *> CmpTraceTargets);
   void InjectTraceForSwitch(Function               &F,
                             ArrayRef<Instruction *> SwitchTraceTargets);
@@ -233,8 +238,10 @@ PreservedAnalyses ModuleSanitizerCoverageAFL::run(Module                &M,
                                                   ModuleAnalysisManager &MAM) {
 
   ModuleSanitizerCoverageAFL ModuleSancov(Options);
+
   auto &FAM = MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
-  auto  DTCallback = [&FAM](Function &F) -> const DominatorTree  *{
+
+  auto DTCallback = [&FAM](Function &F) -> const DominatorTree * {
 
     return &FAM.getResult<DominatorTreeAnalysis>(F);
 
@@ -246,9 +253,21 @@ PreservedAnalyses ModuleSanitizerCoverageAFL::run(Module                &M,
 
   };
 
-  if (ModuleSancov.instrumentModule(M, DTCallback, PDTCallback))
+  auto LoopCallback = [&FAM](Function &F) -> const LoopInfo * {
+
+    return &FAM.getResult<LoopAnalysis>(F);
+
+  };
+
+  if (ModuleSancov.instrumentModule(M, DTCallback, PDTCallback, LoopCallback)) {
+
     return PreservedAnalyses::none();
-  return PreservedAnalyses::all();
+
+  } else {
+
+    return PreservedAnalyses::all();
+
+  }
 
 }
 
@@ -324,7 +343,8 @@ Function *ModuleSanitizerCoverageAFL::CreateInitCallsForSections(
 }
 
 bool ModuleSanitizerCoverageAFL::instrumentModule(
-    Module &M, DomTreeCallback DTCallback, PostDomTreeCallback PDTCallback) {
+    Module &M, DomTreeCallback DTCallback, PostDomTreeCallback PDTCallback,
+    LoopInfoCallback LCallback) {
 
   setvbuf(stdout, NULL, _IONBF, 0);
 
@@ -429,7 +449,7 @@ bool ModuleSanitizerCoverageAFL::instrumentModule(
       M.getOrInsertFunction(SanCovTracePCGuardName, VoidTy, Int32PtrTy);
 
   for (auto &F : M)
-    instrumentFunction(F, DTCallback, PDTCallback);
+    instrumentFunction(F, DTCallback, PDTCallback, LCallback);
 
   Function *Ctor = nullptr;
 
@@ -568,7 +588,8 @@ static bool IsInterestingCmp(ICmpInst *CMP, const DominatorTree *DT,
 #endif
 
 void ModuleSanitizerCoverageAFL::instrumentFunction(
-    Function &F, DomTreeCallback DTCallback, PostDomTreeCallback PDTCallback) {
+    Function &F, DomTreeCallback DTCallback, PostDomTreeCallback PDTCallback,
+    LoopInfoCallback LCallback) {
 
   if (F.empty()) return;
   if (!isInInstrumentList(&F, FMNAME)) return;
@@ -604,6 +625,7 @@ void ModuleSanitizerCoverageAFL::instrumentFunction(
 
   const DominatorTree     *DT = DTCallback(F);
   const PostDominatorTree *PDT = PDTCallback(F);
+  const LoopInfo          *LI = LCallback(F);
   bool                     IsLeafFunc = true;
 
   for (auto &BB : F) {
@@ -640,7 +662,7 @@ void ModuleSanitizerCoverageAFL::instrumentFunction(
   // InjectTraceForCmp(F, CmpTraceTargets);
   // InjectTraceForSwitch(F, SwitchTraceTargets);
 
-  if (dump_cc) { calcCyclomaticComplexity(&F); }
+  if (dump_cc) { calcCyclomaticComplexity(&F, LI); }
 
 }
 
diff --git a/instrumentation/afl-llvm-common.cc b/instrumentation/afl-llvm-common.cc
@@ -15,6 +15,8 @@
 #include <cmath>
 
 #include <llvm/Support/raw_ostream.h>
+#include <llvm/Analysis/LoopInfo.h>
+#include <llvm/Analysis/LoopPass.h>
 
 #define IS_EXTERN extern
 #include "afl-llvm-common.h"
@@ -26,11 +28,39 @@ static std::list<std::string> allowListFunctions;
 static std::list<std::string> denyListFiles;
 static std::list<std::string> denyListFunctions;
 
-unsigned int calcCyclomaticComplexity(llvm::Function *F) {
+static void countNestedLoops(Loop *L, int depth, unsigned int &loopCount,
+                             unsigned int &nestedLoopCount,
+                             unsigned int &maxNestingLevel) {
+
+  loopCount++;
+  if (!L->getSubLoops().empty()) {
+
+    // Increment nested loop count by the number of sub-loops
+    nestedLoopCount += L->getSubLoops().size();
+    // Update maximum nesting level
+    if (depth > maxNestingLevel) { maxNestingLevel = depth; }
+
+    // Recursively count sub-loops
+    for (Loop *SubLoop : L->getSubLoops()) {
+
+      countNestedLoops(SubLoop, depth + 1, loopCount, nestedLoopCount,
+                       maxNestingLevel);
+
+    }
+
+  }
+
+}
+
+unsigned int calcCyclomaticComplexity(llvm::Function       *F,
+                                      const llvm::LoopInfo *LI) {
 
   unsigned int numBlocks = 0;
   unsigned int numEdges = 0;
   unsigned int numCalls = 0;
+  unsigned int numLoops = 0;
+  unsigned int numNestedLoops = 0;
+  unsigned int maxLoopNesting = 0;
 
   // Iterate through each basic block in the function
   for (BasicBlock &BB : *F) {
@@ -55,15 +85,25 @@ unsigned int calcCyclomaticComplexity(llvm::Function *F) {
 
   }
 
+  for (Loop *L : *LI) {
+
+    countNestedLoops(L, 1, numLoops, numNestedLoops, maxLoopNesting);
+
+  }
+
   // Cyclomatic Complexity V(G) = E - N + 2P
   // For a single function, P (number of connected components) is 1
   // Calls are considered to be an edge
-  unsigned int CC = 2 + numCalls + numEdges - numBlocks;
+  unsigned int CC = 2 + numCalls + numEdges - numBlocks + numLoops +
+                    numNestedLoops + maxLoopNesting;
 
   // if (debug) {
 
-  fprintf(stderr, "CyclomaticComplexity for %s: %u\n",
-          F->getName().str().c_str(), CC);
+  fprintf(stderr,
+          "CyclomaticComplexity for %s: %u (calls=%u edges=%u blocks=%u "
+          "loops=%u nested_loops=%u max_loop_nesting_level=%u)\n",
+          F->getName().str().c_str(), CC, numCalls, numEdges, numBlocks,
+          numLoops, numNestedLoops, maxLoopNesting);
 
   //}
 
diff --git a/instrumentation/afl-llvm-common.h b/instrumentation/afl-llvm-common.h
@@ -55,7 +55,8 @@ void  initInstrumentList();
 bool  isInInstrumentList(llvm::Function *F, std::string Filename);
 unsigned long long int calculateCollisions(uint32_t edges);
 void                   scanForDangerousFunctions(llvm::Module *M);
-unsigned int           calcCyclomaticComplexity(llvm::Function *F);
+unsigned int           calcCyclomaticComplexity(llvm::Function       *F,
+                                                const llvm::LoopInfo *LI);
 
 #ifndef IS_EXTERN
   #define IS_EXTERN
