better replay mode error handling, added replay mode documentation, code formatting

7 files changed, 182 insertions, 81 deletions

diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index 864bc6b6..f95dcc20 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -125,7 +125,7 @@
 #endif                                                    /* ^!SIMPLE_FILES */
 
 #ifdef AFL_PERSISTENT_RECORD
-  #define RECORD_PREFIX "RECORD:" 
+  #define RECORD_PREFIX "RECORD:"
 #endif
 
 #define STAGE_BUF_SIZE (64)  /* usable size for stage name buf in afl_state */
diff --git a/include/config.h b/include/config.h
index 1649f110..d44cda9c 100644
--- a/include/config.h
+++ b/include/config.h
@@ -83,10 +83,14 @@
    will be kept and written to the crash/ directory as RECORD:... files.
    Note that every crash will be written, not only unique ones! */
 
-//  #define AFL_PERSISTENT_RECORD
+#define AFL_PERSISTENT_RECORD
 
-/* Builds compiler-rt with support to replay persistent records */
-//  #define AFL_PERSISTENT_REPLAY
+/* Adds support in compiler-rt to replay persistent records */
+#define AFL_PERSISTENT_REPLAY
+
+/* Adds support in compiler-rt to replay persistent records in @@-style
+ * harnesses */
+//  #define AFL_PERSISTENT_REPLAY_ARGPARSE
 
 /* console output colors: There are three ways to configure its behavior
  * 1. default: colored outputs fixed on: defined USE_COLOR && defined
diff --git a/include/persistent_replay.h b/include/persistent_replay.h
index b1a55e9f..58b22fb4 100644
--- a/include/persistent_replay.h
+++ b/include/persistent_replay.h
@@ -11,71 +11,116 @@
 #include <fcntl.h>
 
 static unsigned short int is_replay_record;
-static unsigned int replay_record;
-static unsigned int replay_record_cnt;
-static char replay_record_path[PATH_MAX];
-static char **record_arg;
-static char *replay_record_dir;
-static struct dirent **record_list;
+static unsigned int       replay_record;
+static unsigned int       replay_record_cnt;
+static char               replay_record_path[PATH_MAX];
+static char              *replay_record_dir;
+static struct dirent    **record_list;
+
+#ifdef AFL_PERSISTENT_REPLAY_ARGPARSE
+static char **record_arg = NULL;
+#endif  // AFL_PERSISTENT_REPLAY_ARGPARSE
 
 static int select_files(const struct dirent *dirbuf) {
 
   char fn[4096];
 
-  if (dirbuf->d_name[0] == '.'){
+  if (dirbuf->d_name[0] == '.') {
+
     return 0;
+
   } else {
+
     snprintf(fn, sizeof(fn), "RECORD:%06u", replay_record);
     return !!strstr(dirbuf->d_name, fn);
+
   }
+
 }
-  
+
 static int compare_files(const struct dirent **da, const struct dirent **db) {
-  
-  unsigned int c1=0, c2=0;
+
+  unsigned int c1 = 0, c2 = 0;
 
   sscanf((*da)->d_name, "RECORD:%*u,cnt:%06u", &c1);
   sscanf((*db)->d_name, "RECORD:%*u,cnt:%06u", &c2);
 
-  return c1-c2;
+  return c1 - c2;
+
 }
 
-__attribute__((destructor)) static void __afl_record_replay_destroy(void){
-  for (int i=0; i < replay_record_cnt; i++) {
+__attribute__((destructor)) static void __afl_record_replay_destroy(void) {
+
+  for (int i = 0; i < replay_record_cnt; i++) {
+
     free(record_list[i]);
+
   }
+
   free(record_list);
+
 }
 
-__attribute__((constructor)) static void __afl_record_replay_init(int argc, char **argv) {
-  
+__attribute__((constructor)) static void __afl_record_replay_init(
+#ifdef AFL_PERSISTENT_REPLAY_ARGPARSE
+    int argc, char **argv
+#endif  // AFL_PERSISTENT_REPLAY_ARGPARSE
+) {
+
+#ifdef AFL_PERSISTENT_REPLAY_ARGPARSE
   char **argp;
+#endif  // AFL_PERSISTENT_REPLAY_ARGPARSE
+
+  struct stat sb;
+
+  /* caveat: if harness uses @@ and we don't pass it, it will regardless loop
+   * the number of iterations defined for AFL_LOOP (on the same file)*/
+  if (!(is_replay_record = !!getenv("AFL_PERSISTENT_REPLAY"))) {
 
-  /* caveat: if harness uses @@ and we don't pass it, it will regardless loop the number of iterations defined for AFL_LOOP (on the same file)*/
-  if(!(is_replay_record = !!getenv("AFL_PERSISTENT_REPLAY"))){
     // printf("[warning] AFL_PERSISTENT_REPLAY not set.\n");
     return;
+
   }
 
   replay_record = atoi(getenv("AFL_PERSISTENT_REPLAY"));
   replay_record_dir = getenv("AFL_PERSISTENT_DIR");
-  replay_record_cnt = scandir(replay_record_dir ? replay_record_dir : "./", &record_list, select_files, compare_files);  
 
-  if (!replay_record_cnt){
-    printf("[error] Can't find the requested record!\n");
+  if (!(stat(replay_record_dir, &sb) == 0 && S_ISDIR(sb.st_mode))) {
+
+    fprintf(stderr, "[error] Can't find the requested record directory!\n");
     is_replay_record = 0;
+    return;
+
   }
 
+  replay_record_cnt = scandir(replay_record_dir ? replay_record_dir : "./",
+                              &record_list, select_files, compare_files);
+
+  if (!replay_record_cnt) {
+
+    fprintf(stderr, "[error] Can't find the requested record!\n");
+    is_replay_record = 0;
+
+  }
+
+#ifdef AFL_PERSISTENT_REPLAY_ARGPARSE
   argp = argv;
-  while (*argp){
-    if (!strcmp(*argp, "@@")){
+  while (*argp) {
+
+    if (!strcmp(*argp, "@@")) {
+
       record_arg = argp;
       *record_arg = replay_record_path;
       break;
+
     }
+
     ++argp;
+
   }
 
+#endif  // AFL_PERSISTENT_REPLAY_ARGPARSE
+
 }
 
 /* only used if explictly included for compatibility
@@ -83,67 +128,80 @@ __attribute__((constructor)) static void __afl_record_replay_init(int argc, char
 
 #ifdef AFL_COMPAT
 
-#ifndef PATH_MAX
-  #define PATH_MAX 4096
-#endif
+  #ifndef PATH_MAX
+    #define PATH_MAX 4096
+  #endif
 
-#define FUZZ_BUF_SIZE 1024000
+  #define FUZZ_BUF_SIZE 1024000
 
-// extern ssize_t read(int fildes, void *buf, size_t nbyte);
+  // extern ssize_t read(int fildes, void *buf, size_t nbyte);
 
-//extern int __afl_persistent_loop(unsigned int max_cnt);
-//extern unsigned char fuzz_buf[];
+  // extern int __afl_persistent_loop(unsigned int max_cnt);
+  // extern unsigned char fuzz_buf[];
 
-#ifndef __AFL_HAVE_MANUAL_CONTROL
-  #define __AFL_HAVE_MANUAL_CONTROL
-#endif
+  #ifndef __AFL_HAVE_MANUAL_CONTROL
+    #define __AFL_HAVE_MANUAL_CONTROL
+  #endif
 
-#define __AFL_FUZZ_TESTCASE_LEN (read(0, fuzz_buf, FUZZ_BUF_SIZE))
-#define __AFL_FUZZ_TESTCASE_BUF fuzz_buf
-#define __AFL_FUZZ_INIT() void sync(void); 
-#define __AFL_INIT() sync() 
-#define __AFL_LOOP(x) __afl_persistent_loop(x)
+  #define __AFL_FUZZ_TESTCASE_LEN (read(0, fuzz_buf, FUZZ_BUF_SIZE))
+  #define __AFL_FUZZ_TESTCASE_BUF fuzz_buf
+  #define __AFL_FUZZ_INIT() void sync(void);
+  #define __AFL_INIT() sync()
+  #define __AFL_LOOP(x) __afl_persistent_loop(x)
 
 unsigned char fuzz_buf[FUZZ_BUF_SIZE];
 
 int __afl_persistent_loop(unsigned int max_cnt) {
 
-  static unsigned int cycle_cnt = 1;
+  static unsigned int       cycle_cnt = 1;
   static unsigned short int inited = 0;
-  char tcase[PATH_MAX];
+  char                      tcase[PATH_MAX];
+
+  if (is_replay_record) {
 
-  if( is_replay_record ){
+    if (!inited) {
 
-    if (!inited){
       cycle_cnt = replay_record_cnt;
       inited = 1;
+
     }
 
     snprintf(tcase, PATH_MAX, "%s/%s",
-                  replay_record_dir ? replay_record_dir : "./",
-                  record_list[replay_record_cnt-cycle_cnt]->d_name);
-
+             replay_record_dir ? replay_record_dir : "./",
+             record_list[replay_record_cnt - cycle_cnt]->d_name);
 
+  #ifdef AFL_PERSISTENT_REPLAY_ARGPARSE
     if (record_arg) {
+
       *record_arg = tcase;
-    } else {
+
+    } else
+
+  #endif  // AFL_PERSISTENT_REPLAY_ARGPARSE
+    {
+
       int fd = open(tcase, O_RDONLY);
       dup2(fd, 0);
       close(fd);
+
     }
 
   } else {
 
-    if (!inited){
+    if (!inited) {
+
       cycle_cnt = max_cnt;
       inited = 1;
+
     }
 
   }
 
   return cycle_cnt--;
+
 }
 
 #endif  // AFL_COMPAT
 
-#endif  // _HAVE_PERSISTENT_REPLAY_H
\ No newline at end of file
+#endif  // _HAVE_PERSISTENT_REPLAY_H
+
diff --git a/instrumentation/README.persistent_mode.md b/instrumentation/README.persistent_mode.md
index 14e59f4a..b5d982b0 100644
--- a/instrumentation/README.persistent_mode.md
+++ b/instrumentation/README.persistent_mode.md
@@ -195,4 +195,32 @@ Then as first line after the `__AFL_LOOP` while loop:
   int len = __AFL_FUZZ_TESTCASE_LEN;
 ```
 
-And that is all!
\ No newline at end of file
+And that is all!
+
+## 6) Persistent record, and replay
+
+If your software under test requires keeping a state between persistent loop iterations (i.e., a stateful network stack), you can use the `AFL_PERSISTENT_RECORD` variable as described in the [environment variables documentation](../docs/env_variables.md).
+
+To easily replay a crashing, or hanging record, you can use the persistent replay functionality by compiling AFL++ after uncommenting the `AFL_PERSISTENT_REPLAY` define  in [config.h](../include/config.h).
+
+You can then run the test binary specifying the record number via the AFL_PERSISTENT_REPLAY environment variable (i.e., `RECORD:XXXXX`` -> `AFL_PERSISTENT_REPLAY=XXXXX`).
+The directory where the record files live can be specified via the `AFL_PERSISTENT_DIR` environment varilable, otherwise by default it will be considered the current directory (`./`).
+
+If your harness reads the input files from arguments using the special `@@` argument you will need to define `AFL_PERSISTENT_ARGPARSE` in  `config.h`, or before including the `persistent_replay.h` header file as show before.
+In order to offer transparent support to harnesses using the `@@` command line argument, arguments are parsed by the `__afl_record_replay_init` init function. Since not all systems support passing arguments to initializers, this functionality is disabled by default, it's recommendable to use the `__AFL_FUZZ_TESTCASE_BUF/__AFL_FUZZ_TESTCASE_LEN` shared memory mechanism instead.
+
+### 7) Drop in replay functionality
+
+To use the replay functionality without having to use `afl-cc` you can just define `AFL_COMPAT` and include the [include/persistent_replay.h](../include/persistent_replay.h) self contained header file that provides a drop-in replacement for the persistent loop mechanism.
+
+```c
+#ifndef __AFL_FUZZ_TESTCASE_LEN
+  #define AFL_COMPAT
+  // #define AFL_PERSISTENT_REPLAY_ARGPARSE
+  #include "persistent_replay.h"
+#endif
+
+__AFL_FUZZ_INIT();
+```
+
+A simple example is provided in [persistent_demo_new.c](../utils/persistent_mode/persistent_demo_new.c).
\ No newline at end of file
diff --git a/instrumentation/afl-compiler-rt.o.c b/instrumentation/afl-compiler-rt.o.c
index 0fa22aee..037caaf0 100644
--- a/instrumentation/afl-compiler-rt.o.c
+++ b/instrumentation/afl-compiler-rt.o.c
@@ -84,7 +84,7 @@
 #include <fcntl.h>
 
 #ifdef AFL_PERSISTENT_REPLAY
-#include "persistent_replay.h"
+  #include "persistent_replay.h"
 #endif
 
 /* Globals needed by the injected instrumentation. The __afl_area_initial region
@@ -1344,37 +1344,49 @@ int __afl_persistent_loop(unsigned int max_cnt) {
 
 #ifdef AFL_PERSISTENT_REPLAY
 
-#ifndef PATH_MAX
-  #define PATH_MAX 4096
-#endif
+  #ifndef PATH_MAX
+    #define PATH_MAX 4096
+  #endif
 
-  static u8  inited = 0;
-  char tcase[PATH_MAX];
+  static u8 inited = 0;
+  char      tcase[PATH_MAX];
 
-  if( unlikely(is_replay_record) ){
+  if (unlikely(is_replay_record)) {
 
-      if (!inited){
-        cycle_cnt = replay_record_cnt;
-        inited = 1;
-      }
+    if (!inited) {
 
-      snprintf(tcase, PATH_MAX, "%s/%s",
-                  replay_record_dir ? replay_record_dir : "./",
-                  record_list[replay_record_cnt-cycle_cnt]->d_name);
+      cycle_cnt = replay_record_cnt;
+      inited = 1;
+
+    }
+
+    snprintf(tcase, PATH_MAX, "%s/%s",
+             replay_record_dir ? replay_record_dir : "./",
+             record_list[replay_record_cnt - cycle_cnt]->d_name);
+
+  #ifdef AFL_PERSISTENT_REPLAY_ARGPARSE
+    if (record_arg) {
+
+      *record_arg = tcase;
+
+    } else
+
+  #endif  // AFL_PERSISTENT_REPLAY_ARGPARSE
+    {
+
+      int fd = open(tcase, O_RDONLY);
+      dup2(fd, 0);
+      close(fd);
+
+    }
 
-      if (record_arg) {
-        *record_arg = tcase;
-      } else {
-        int fd = open(tcase, O_RDONLY);
-        dup2(fd, 0);
-        close(fd);
-      }
     return cycle_cnt--;
+
   } else
 
-#endif  
+#endif
 
-  if (first_pass) {
+      if (first_pass) {
 
     /* Make sure that every iteration of __AFL_LOOP() starts with a clean slate.
        On subsequent calls, the parent will take care of that, but on the first
diff --git a/src/afl-forkserver.c b/src/afl-forkserver.c
index f8dd783f..36e46444 100644
--- a/src/afl-forkserver.c
+++ b/src/afl-forkserver.c
@@ -1593,7 +1593,7 @@ afl_fsrv_run_target(afl_forkserver_t *fsrv, u32 timeout,
 
 #ifdef AFL_PERSISTENT_RECORD
   fsrv_run_result_t retval = FSRV_RUN_OK;
-  char *persistent_out_fmt;
+  char             *persistent_out_fmt;
 #endif
 
 #ifdef __linux__
@@ -1803,6 +1803,7 @@ afl_fsrv_run_target(afl_forkserver_t *fsrv, u32 timeout,
     persistent_out_fmt = "%s/hangs/RECORD:%06u,cnt:%06u";
     goto store_persistent_record;
 #endif
+
   }
 
   /* Did we crash?
@@ -1841,7 +1842,7 @@ afl_fsrv_run_target(afl_forkserver_t *fsrv, u32 timeout,
 
 #ifdef AFL_PERSISTENT_RECORD
 store_persistent_record:
-  if (unlikely(retval == FSRV_RUN_CRASH || retval == FSRV_RUN_TMOUT)  &&
+  if (unlikely(retval == FSRV_RUN_CRASH || retval == FSRV_RUN_TMOUT) &&
       unlikely(fsrv->persistent_record)) {
 
     char fn[PATH_MAX];
diff --git a/utils/persistent_mode/persistent_demo_new.c b/utils/persistent_mode/persistent_demo_new.c
index 40ada9e1..3d9d90a6 100644
--- a/utils/persistent_mode/persistent_demo_new.c
+++ b/utils/persistent_mode/persistent_demo_new.c
@@ -31,8 +31,8 @@
 
 /* this lets the source compile without afl-clang-fast/lto */
 #ifndef __AFL_FUZZ_TESTCASE_LEN
-#define AFL_COMPAT
-#include "persistent_replay.h"
+  #define AFL_COMPAT
+  #include "persistent_replay.h"
 #endif
 
 __AFL_FUZZ_INIT();
@@ -86,8 +86,6 @@ int main(int argc, char **argv) {
               if (buf[5] == '!') {
 
                 printf("six\n");
-                char *nullo = NULL+1;
-                *nullo = 'p';
                 abort();
 
               }