Always enable persistent mode, no env/bincheck needed

3 files changed, 30 insertions, 30 deletions

diff --git a/dynamic_list.txt b/dynamic_list.txt
index 7293ae77..d06eae4e 100644
--- a/dynamic_list.txt
+++ b/dynamic_list.txt
@@ -8,6 +8,7 @@
   "__afl_auto_first";
   "__afl_auto_init";
   "__afl_auto_second";
+  "__afl_connected";
   "__afl_coverage_discard";
   "__afl_coverage_interesting";
   "__afl_coverage_off";
diff --git a/instrumentation/afl-compiler-rt.o.c b/instrumentation/afl-compiler-rt.o.c
index 20069824..9dded1dd 100644
--- a/instrumentation/afl-compiler-rt.o.c
+++ b/instrumentation/afl-compiler-rt.o.c
@@ -105,6 +105,9 @@ u32 __afl_dictionary_len;
 u64 __afl_map_addr;
 u32 __afl_first_final_loc;
 
+/* 1 if we are running in afl, and the forkserver was stared, else 0 */
+u32 __afl_connected = 0;
+
 // for the __AFL_COVERAGE_ON/__AFL_COVERAGE_OFF features to work:
 int        __afl_selective_coverage __attribute__((weak));
 int        __afl_selective_coverage_start_off __attribute__((weak));
@@ -1048,7 +1051,11 @@ static void __afl_start_forkserver(void) {
   /* Phone home and tell the parent that we're OK. If parent isn't there,
      assume we're not running in forkserver mode and just execute program. */
 
-  if (write(FORKSRV_FD + 1, tmp, 4) != 4) { return; }
+  if (write(FORKSRV_FD + 1, tmp, 4) != 4) {
+    return;
+  }
+
+  __afl_connected = 1;
 
   if (__afl_sharedmem_fuzzing || (__afl_dictionary_len && __afl_dictionary)) {
 
@@ -1260,13 +1267,9 @@ int __afl_persistent_loop(unsigned int max_cnt) {
        iteration, it's our job to erase any trace of whatever happened
        before the loop. */
 
-    if (is_persistent) {
-
-      memset(__afl_area_ptr, 0, __afl_map_size);
-      __afl_area_ptr[0] = 1;
-      memset(__afl_prev_loc, 0, NGRAM_SIZE_MAX * sizeof(PREV_LOC_T));
-
-    }
+    memset(__afl_area_ptr, 0, __afl_map_size);
+    __afl_area_ptr[0] = 1;
+    memset(__afl_prev_loc, 0, NGRAM_SIZE_MAX * sizeof(PREV_LOC_T));
 
     cycle_cnt = max_cnt;
     first_pass = 0;
@@ -1274,34 +1277,28 @@ int __afl_persistent_loop(unsigned int max_cnt) {
 
     return 1;
 
-  }
-
-  if (is_persistent) {
+  } else if (--cycle_cnt) {
 
-    if (--cycle_cnt) {
+    raise(SIGSTOP);
 
-      raise(SIGSTOP);
-
-      __afl_area_ptr[0] = 1;
-      memset(__afl_prev_loc, 0, NGRAM_SIZE_MAX * sizeof(PREV_LOC_T));
-      __afl_selective_coverage_temp = 1;
+    __afl_area_ptr[0] = 1;
+    memset(__afl_prev_loc, 0, NGRAM_SIZE_MAX * sizeof(PREV_LOC_T));
+    __afl_selective_coverage_temp = 1;
 
-      return 1;
+    return 1;
 
-    } else {
+  } else {
 
-      /* When exiting __AFL_LOOP(), make sure that the subsequent code that
-         follows the loop is not traced. We do that by pivoting back to the
-         dummy output region. */
+    /* When exiting __AFL_LOOP(), make sure that the subsequent code that
+        follows the loop is not traced. We do that by pivoting back to the
+        dummy output region. */
 
-      __afl_area_ptr = __afl_area_ptr_dummy;
+    __afl_area_ptr = __afl_area_ptr_dummy;
 
-    }
+    return 0;
 
   }
 
-  return 0;
-
 }
 
 /* This one can be called from user code when deferred forkserver mode
diff --git a/src/afl-cc.c b/src/afl-cc.c
index 469aa825..4d6b636a 100644
--- a/src/afl-cc.c
+++ b/src/afl-cc.c
@@ -45,7 +45,7 @@
   #define LLVM_MINOR 0
 #endif
 
-static u8  *obj_path;                  /* Path to runtime libraries         */
+static u8 * obj_path;                  /* Path to runtime libraries         */
 static u8 **cc_params;                 /* Parameters passed to the real CC  */
 static u32  cc_par_cnt = 1;            /* Param count, including argv0      */
 static u8   clang_mode;                /* Invoked as afl-clang*?            */
@@ -53,7 +53,7 @@ static u8   llvm_fullpath[PATH_MAX];
 static u8   instrument_mode, instrument_opt_mode, ngram_size, ctx_k, lto_mode;
 static u8   compiler_mode, plusplus_mode, have_instr_env = 0, need_aflpplib = 0;
 static u8   have_gcc, have_llvm, have_gcc_plugin, have_lto, have_instr_list = 0;
-static u8  *lto_flag = AFL_CLANG_FLTO, *argvnull;
+static u8 * lto_flag = AFL_CLANG_FLTO, *argvnull;
 static u8   debug;
 static u8   cwd[4096];
 static u8   cmplog_mode;
@@ -1144,14 +1144,16 @@ static void edit_params(u32 argc, char **argv, char **envp) {
       "({ static volatile char *_B __attribute__((used,unused)); "
       " _B = (char*)\"" PERSIST_SIG
       "\"; "
+      "extern int __afl_connected;"
 #ifdef __APPLE__
       "__attribute__((visibility(\"default\"))) "
       "int _L(unsigned int) __asm__(\"___afl_persistent_loop\"); "
 #else
       "__attribute__((visibility(\"default\"))) "
-      "int _L(unsigned int) __asm__(\"__afl_persistent_loop\"); "
+      "int _L(unsigned int) __asm__(\"___afl_persistent_loop\"); "
 #endif                                                        /* ^__APPLE__ */
-      "_L(_A); })";
+      // if afl is connected, we run _A times, else once.
+      "_L(__afl_connected ? _A : 1); })";
 
   cc_params[cc_par_cnt++] =
       "-D__AFL_INIT()="