From 80eabd6e8a30c2ffc0f084ab34df8b9d582419c3 Mon Sep 17 00:00:00 2001
From: vanhauser-thc <vh@thc.org>
Date: Mon, 13 Feb 2023 11:34:14 +0100
Subject: AFL_LLVM_DICT2FILE_NO_MAIN support

---
 instrumentation/afl-llvm-dict2file.so.cc | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

(limited to 'instrumentation/afl-llvm-dict2file.so.cc')

diff --git a/instrumentation/afl-llvm-dict2file.so.cc b/instrumentation/afl-llvm-dict2file.so.cc
index bbbbe32c..97f1d47f 100644
--- a/instrumentation/afl-llvm-dict2file.so.cc
+++ b/instrumentation/afl-llvm-dict2file.so.cc
@@ -182,7 +182,7 @@ bool AFLdict2filePass::runOnModule(Module &M) {
 
   DenseMap<Value *, std::string *> valueMap;
   char                            *ptr;
-  int                              found = 0;
+  int                              found = 0, handle_main = 1;
 
   /* Show a banner */
   setvbuf(stdout, NULL, _IONBF, 0);
@@ -192,10 +192,14 @@ bool AFLdict2filePass::runOnModule(Module &M) {
     SAYF(cCYA "afl-llvm-dict2file" VERSION cRST
               " by Marc \"vanHauser\" Heuse <mh@mh-sec.de>\n");
 
-  } else
+  } else {
 
     be_quiet = 1;
 
+  }
+
+  if (getenv("AFL_LLVM_DICT2FILE_NO_MAIN")) { handle_main = 0; }
+
   scanForDangerousFunctions(&M);
 
   ptr = getenv("AFL_LLVM_DICT2FILE");
@@ -210,7 +214,14 @@ bool AFLdict2filePass::runOnModule(Module &M) {
 
   for (auto &F : M) {
 
-    if (isIgnoreFunction(&F)) continue;
+    if (!handle_main &&
+        (!F.getName().compare("main") || !F.getName().compare("_main"))) {
+
+      continue;
+
+    }
+
+    if (isIgnoreFunction(&F)) { continue; }
     if (!isInInstrumentList(&F, MNAME) || !F.size()) { continue; }
 
     /*  Some implementation notes.
-- 
cgit 1.4.1