From 1db3b81d2eb855167dcf65734f8833a2329609da Mon Sep 17 00:00:00 2001
From: vanhauser-thc <vh@thc.org>
Date: Wed, 15 May 2024 14:16:44 +0200
Subject: dump cc

---
 instrumentation/SanitizerCoveragePCGUARD.so.cc |  6 +++-
 instrumentation/afl-llvm-common.cc             | 45 ++++++++++++++++++++++++++
 instrumentation/afl-llvm-common.h              |  1 +
 3 files changed, 51 insertions(+), 1 deletion(-)

(limited to 'instrumentation')

diff --git a/instrumentation/SanitizerCoveragePCGUARD.so.cc b/instrumentation/SanitizerCoveragePCGUARD.so.cc
index f88ce126..01881f28 100644
--- a/instrumentation/SanitizerCoveragePCGUARD.so.cc
+++ b/instrumentation/SanitizerCoveragePCGUARD.so.cc
@@ -195,7 +195,7 @@ class ModuleSanitizerCoverageAFL
 
   SanitizerCoverageOptions Options;
 
-  uint32_t        instr = 0, selects = 0, unhandled = 0;
+  uint32_t        instr = 0, selects = 0, unhandled = 0, dump_cc = 0;
   GlobalVariable *AFLMapPtr = NULL;
   ConstantInt    *One = NULL;
   ConstantInt    *Zero = NULL;
@@ -330,6 +330,8 @@ bool ModuleSanitizerCoverageAFL::instrumentModule(
 
   if (getenv("AFL_DEBUG")) { debug = 1; }
 
+  if (getenv("AFL_DUMP_CYCLOMATIC_COMPLEXITY")) { dump_cc = 1; }
+
   if ((isatty(2) && !getenv("AFL_QUIET")) || debug) {
 
     SAYF(cCYA "SanitizerCoveragePCGUARD" VERSION cRST "\n");
@@ -638,6 +640,8 @@ void ModuleSanitizerCoverageAFL::instrumentFunction(
   // InjectTraceForCmp(F, CmpTraceTargets);
   // InjectTraceForSwitch(F, SwitchTraceTargets);
 
+  if (dump_cc) { calcCyclomaticComplexity(&F); }
+
 }
 
 GlobalVariable *ModuleSanitizerCoverageAFL::CreateFunctionLocalArrayInSection(
diff --git a/instrumentation/afl-llvm-common.cc b/instrumentation/afl-llvm-common.cc
index 8e9e7800..ed9268dc 100644
--- a/instrumentation/afl-llvm-common.cc
+++ b/instrumentation/afl-llvm-common.cc
@@ -26,6 +26,51 @@ static std::list<std::string> allowListFunctions;
 static std::list<std::string> denyListFiles;
 static std::list<std::string> denyListFunctions;
 
+unsigned int calcCyclomaticComplexity(llvm::Function *F) {
+
+  unsigned int numBlocks = 0;
+  unsigned int numEdges = 0;
+  unsigned int numCalls = 0;
+
+  // Iterate through each basic block in the function
+  for (BasicBlock &BB : *F) {
+
+    // count all nodes == basic blocks
+    numBlocks++;
+    // Count the number of successors (outgoing edges)
+    for (BasicBlock *Succ : successors(&BB)) {
+
+      // count edges for CC
+      numEdges++;
+      (void)(Succ);
+
+    }
+
+    for (Instruction &I : BB) {
+
+      // every call is also an edge, so we need to count the calls too
+      if (isa<CallInst>(&I) || isa<InvokeInst>(&I)) { numCalls++; }
+
+    }
+
+  }
+
+  // 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;
+
+  // if (debug) {
+
+  fprintf(stderr, "CyclomaticComplexity for %s: %u\n",
+          F->getName().str().c_str(), CC);
+
+  //}
+
+  return CC;
+
+}
+
 char *getBBName(const llvm::BasicBlock *BB) {
 
   static char *name;
diff --git a/instrumentation/afl-llvm-common.h b/instrumentation/afl-llvm-common.h
index 23f67179..6b628d64 100644
--- a/instrumentation/afl-llvm-common.h
+++ b/instrumentation/afl-llvm-common.h
@@ -55,6 +55,7 @@ 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);
 
 #ifndef IS_EXTERN
   #define IS_EXTERN
-- 
cgit 1.4.1