AFL_LLVM_DICT2FILE_NO_MAIN support

9 files changed, 41 insertions, 8 deletions

diff --git a/TODO.md b/TODO.md
index 862224f0..187fa191 100644
--- a/TODO.md
+++ b/TODO.md
@@ -9,13 +9,12 @@
  - afl-plot to support multiple plot_data
  - parallel builds for source-only targets
  - get rid of check_binary, replace with more forkserver communication
- - first fuzzer should be a main automatically
+ - first fuzzer should be a main automatically? not sure.
 
 ## Maybe
 
  - forkserver tells afl-fuzz if cmplog is supported and if so enable
    it by default, with AFL_CMPLOG_NO=1 (?) set to skip?
- - afl_custom_fuzz_splice_optin()
  - afl_custom_splice()
  - cmdline option from-to range for mutations
 
diff --git a/docs/Changelog.md b/docs/Changelog.md
index eee88a51..89c37912 100644
--- a/docs/Changelog.md
+++ b/docs/Changelog.md
@@ -10,6 +10,7 @@
     - add CFI sanitizer variant to gcc targets
     - llvm 16 support (thanks to @devnexen!)
     - support llvm 15 native pcguard changes
+  - LTO autoken and llvm_mode: added AFL_LLVM_DICT2FILE_NO_MAIN support
   - better sanitizer default options support for all tools
   - unicorn_mode: updated and minor issues fixed
   - frida_mode: fix issue on MacOS
diff --git a/docs/env_variables.md b/docs/env_variables.md
index 0a57d190..61fb1e2b 100644
--- a/docs/env_variables.md
+++ b/docs/env_variables.md
@@ -129,6 +129,9 @@ subset of the settings discussed in section 1, with the exception of:
     write all constant string comparisons to this file to be used later with
     afl-fuzz' `-x` option.
 
+  - An option to `AFL_LLVM_DICT2FILE` is `AFL_LLVM_DICT2FILE_NO_MAIN=1` which
+    skill not parse `main()`.
+
   - `TMPDIR` and `AFL_KEEP_ASSEMBLY`, since no temporary assembly files are
     created.
 
diff --git a/docs/fuzzing_in_depth.md b/docs/fuzzing_in_depth.md
index 87f31a58..efab0633 100644
--- a/docs/fuzzing_in_depth.md
+++ b/docs/fuzzing_in_depth.md
@@ -534,6 +534,8 @@ dictionaries/FORMAT.dict`.
 * With `afl-clang-fast`, you can set
   `AFL_LLVM_DICT2FILE=/full/path/to/new/file.dic` to automatically generate a
   dictionary during target compilation.
+  Adding `AFL_LLVM_DICT2FILE_NO_MAIN=1` to not parse main (usually command line
+  parameter parsing) is often a good idea too.
 * You also have the option to generate a dictionary yourself during an
   independent run of the target, see
   [utils/libtokencap/README.md](../utils/libtokencap/README.md).
@@ -935,7 +937,7 @@ phase and start fuzzing at once.
 3. Also randomize the afl-fuzz runtime options, e.g.:
     * 65% for `AFL_DISABLE_TRIM`
     * 50% for `AFL_KEEP_TIMEOUTS`
-    * 50% use a dictionary generated by `AFL_LLVM_DICT2FILE`
+    * 50% use a dictionary generated by `AFL_LLVM_DICT2FILE` + `AFL_LLVM_DICT2FILE_NO_MAIN=1`
     * 40% use MOpt (`-L 0`)
     * 40% for `AFL_EXPAND_HAVOC_NOW`
     * 20% for old queue processing (`-Z`)
diff --git a/include/envs.h b/include/envs.h
index 0770f94d..5018b0f8 100644
--- a/include/envs.h
+++ b/include/envs.h
@@ -133,6 +133,7 @@ static char *afl_environment_variables[] = {
     "AFL_LLVM_CTX",
     "AFL_LLVM_CTX_K",
     "AFL_LLVM_DICT2FILE",
+    "AFL_LLVM_DICT2FILE_NO_MAIN",
     "AFL_LLVM_DOCUMENT_IDS",
     "AFL_LLVM_INSTRIM_LOOPHEAD",
     "AFL_LLVM_INSTRUMENT",
diff --git a/instrumentation/README.llvm.md b/instrumentation/README.llvm.md
index 9da1b0f6..c0677474 100644
--- a/instrumentation/README.llvm.md
+++ b/instrumentation/README.llvm.md
@@ -167,6 +167,10 @@ Just specify `AFL_LLVM_DICT2FILE=/absolute/path/file.txt` and during compilation
 all constant string compare parameters will be written to this file to be used
 with afl-fuzz' `-x` option.
 
+Adding `AFL_LLVM_DICT2FILE_NO_MAIN=1` will skip parsing `main()` which often
+does command line parsing which has string comparisons that are not helpful
+for fuzzing.
+
 ## 6) AFL++ Context Sensitive Branch Coverage
 
 ### What is this?
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) {
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.
diff --git a/src/afl-cc.c b/src/afl-cc.c
index 7c3682fb..7b059d40 100644
--- a/src/afl-cc.c
+++ b/src/afl-cc.c
@@ -2041,6 +2041,8 @@ int main(int argc, char **argv, char **envp) {
 
             "  AFL_LLVM_DICT2FILE: generate an afl dictionary based on found "
             "comparisons\n"
+            "  AFL_LLVM_DICT2FILE_NO_MAIN: skip parsing main() for the "
+            "dictionary\n"
             "  AFL_LLVM_LAF_ALL: enables all LAF splits/transforms\n"
             "  AFL_LLVM_LAF_SPLIT_COMPARES: enable cascaded comparisons\n"
             "  AFL_LLVM_LAF_SPLIT_COMPARES_BITW: size limit (default 8)\n"
@@ -2128,7 +2130,8 @@ int main(int argc, char **argv, char **envp) {
         "defaults.\n"
         "Recommended is afl-clang-lto with AFL_LLVM_CMPLOG or afl-clang-fast "
         "with\n"
-        "AFL_LLVM_CMPLOG and AFL_LLVM_DICT2FILE.\n\n");
+        "AFL_LLVM_CMPLOG and "
+        "AFL_LLVM_DICT2FILE+AFL_LLVM_DICT2FILE_NO_MAIN.\n\n");
 
     exit(1);