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/SanitizerCoverageLTO.so.cc | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

(limited to 'instrumentation/SanitizerCoverageLTO.so.cc')

diff --git a/instrumentation/SanitizerCoverageLTO.so.cc b/instrumentation/SanitizerCoverageLTO.so.cc
index 231151f5..f82224ed 100644
--- a/instrumentation/SanitizerCoverageLTO.so.cc
+++ b/instrumentation/SanitizerCoverageLTO.so.cc
@@ -236,6 +236,7 @@ class ModuleSanitizerCoverageLTO
   // const SpecialCaseList *          Allowlist;
   // const SpecialCaseList *          Blocklist;
   uint32_t                         autodictionary = 1;
+  uint32_t                         autodictionary_no_main = 0;
   uint32_t                         inst = 0;
   uint32_t                         afl_global_id = 0;
   uint32_t                         unhandled = 0;
@@ -411,7 +412,8 @@ bool ModuleSanitizerCoverageLTO::instrumentModule(
 
   /* Show a banner */
   setvbuf(stdout, NULL, _IONBF, 0);
-  if (getenv("AFL_DEBUG")) debug = 1;
+  if (getenv("AFL_DEBUG")) { debug = 1; }
+  if (getenv("AFL_LLVM_DICT2FILE_NO_MAIN")) { autodictionary_no_main = 1; }
 
   if ((isatty(2) && !getenv("AFL_QUIET")) || debug) {
 
@@ -503,6 +505,13 @@ bool ModuleSanitizerCoverageLTO::instrumentModule(
 
       if (!isInInstrumentList(&F, MNAME) || !F.size()) { continue; }
 
+      if (autodictionary_no_main &&
+          (!F.getName().compare("main") || !F.getName().compare("_main"))) {
+
+        continue;
+
+      }
+
       for (auto &BB : F) {
 
         for (auto &IN : BB) {
-- 
cgit 1.4.1