replay mode support

10 files changed, 257 insertions, 43 deletions

diff --git a/.gitignore b/.gitignore
index f76a86fc..891ced9f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -103,6 +103,7 @@ utils/optimin/build
 utils/optimin/optimin
 utils/persistent_mode/persistent_demo
 utils/persistent_mode/persistent_demo_new
+utils/persistent_mode/persistent_demo_new_compat
 utils/persistent_mode/test-instr
 utils/plot_ui/afl-plot-ui
 vuln_prog
diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index f1813df6..864bc6b6 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -124,6 +124,10 @@
   #define CASE_PREFIX "id_"
 #endif                                                    /* ^!SIMPLE_FILES */
 
+#ifdef AFL_PERSISTENT_RECORD
+  #define RECORD_PREFIX "RECORD:" 
+#endif
+
 #define STAGE_BUF_SIZE (64)  /* usable size for stage name buf in afl_state */
 
 // Little helper to access the ptr to afl->##name_buf - for use in afl_realloc.
diff --git a/include/config.h b/include/config.h
index 63340650..1649f110 100644
--- a/include/config.h
+++ b/include/config.h
@@ -83,7 +83,10 @@
    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
 
 /* 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
new file mode 100644
index 00000000..b1a55e9f
--- /dev/null
+++ b/include/persistent_replay.h
@@ -0,0 +1,149 @@
+#ifndef _HAVE_PERSISTENT_REPLAY_H
+#define _HAVE_PERSISTENT_REPLAY_H
+
+#include <dirent.h>
+#include <string.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <malloc.h>
+#include <unistd.h>
+#include <sys/stat.h>
+#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 int select_files(const struct dirent *dirbuf) {
+
+  char fn[4096];
+
+  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;
+
+  sscanf((*da)->d_name, "RECORD:%*u,cnt:%06u", &c1);
+  sscanf((*db)->d_name, "RECORD:%*u,cnt:%06u", &c2);
+
+  return c1-c2;
+}
+
+__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) {
+  
+  char **argp;
+
+  /* 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");
+    is_replay_record = 0;
+  }
+
+  argp = argv;
+  while (*argp){
+    if (!strcmp(*argp, "@@")){
+      record_arg = argp;
+      *record_arg = replay_record_path;
+      break;
+    }
+    ++argp;
+  }
+
+}
+
+/* only used if explictly included for compatibility
+   compiling without afl-cc */
+
+#ifdef AFL_COMPAT
+
+#ifndef PATH_MAX
+  #define PATH_MAX 4096
+#endif
+
+#define FUZZ_BUF_SIZE 1024000
+
+// 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[];
+
+#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)
+
+unsigned char fuzz_buf[FUZZ_BUF_SIZE];
+
+int __afl_persistent_loop(unsigned int max_cnt) {
+
+  static unsigned int cycle_cnt = 1;
+  static unsigned short int inited = 0;
+  char tcase[PATH_MAX];
+
+  if( is_replay_record ){
+
+    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);
+
+
+    if (record_arg) {
+      *record_arg = tcase;
+    } else {
+      int fd = open(tcase, O_RDONLY);
+      dup2(fd, 0);
+      close(fd);
+    }
+
+  } else {
+
+    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
diff --git a/instrumentation/afl-compiler-rt.o.c b/instrumentation/afl-compiler-rt.o.c
index 39a762b6..0fa22aee 100644
--- a/instrumentation/afl-compiler-rt.o.c
+++ b/instrumentation/afl-compiler-rt.o.c
@@ -83,6 +83,10 @@
 #include <sys/mman.h>
 #include <fcntl.h>
 
+#ifdef AFL_PERSISTENT_REPLAY
+#include "persistent_replay.h"
+#endif
+
 /* Globals needed by the injected instrumentation. The __afl_area_initial region
    is used for instrumentation output before __afl_map_shm() has a chance to
    run. It will end up as .comm, so it shouldn't be too wasteful. */
@@ -1338,6 +1342,38 @@ int __afl_persistent_loop(unsigned int max_cnt) {
   static u8  first_pass = 1;
   static u32 cycle_cnt;
 
+#ifdef AFL_PERSISTENT_REPLAY
+
+#ifndef PATH_MAX
+  #define PATH_MAX 4096
+#endif
+
+  static u8  inited = 0;
+  char tcase[PATH_MAX];
+
+  if( unlikely(is_replay_record) ){
+
+      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);
+
+      if (record_arg) {
+        *record_arg = tcase;
+      } else {
+        int fd = open(tcase, O_RDONLY);
+        dup2(fd, 0);
+        close(fd);
+      }
+    return cycle_cnt--;
+  } else
+
+#endif  
+
   if (first_pass) {
 
     /* Make sure that every iteration of __AFL_LOOP() starts with a clean slate.
diff --git a/src/afl-forkserver.c b/src/afl-forkserver.c
index 3f9bfa72..f8dd783f 100644
--- a/src/afl-forkserver.c
+++ b/src/afl-forkserver.c
@@ -1591,6 +1591,11 @@ afl_fsrv_run_target(afl_forkserver_t *fsrv, u32 timeout,
   u32 exec_ms;
   u32 write_value = fsrv->last_run_timed_out;
 
+#ifdef AFL_PERSISTENT_RECORD
+  fsrv_run_result_t retval = FSRV_RUN_OK;
+  char *persistent_out_fmt;
+#endif
+
 #ifdef __linux__
   if (fsrv->nyx_mode) {
 
@@ -1684,7 +1689,7 @@ afl_fsrv_run_target(afl_forkserver_t *fsrv, u32 timeout,
 
   }
 
-#ifdef AFL_PERSISTENT_RECORD
+#ifdef AFL_eERSISTENT_RECORD
   // end of persistent loop?
   if (unlikely(fsrv->persistent_record &&
                fsrv->persistent_record_pid != fsrv->child_pid)) {
@@ -1790,8 +1795,14 @@ afl_fsrv_run_target(afl_forkserver_t *fsrv, u32 timeout,
   if (unlikely(fsrv->last_run_timed_out)) {
 
     fsrv->last_kill_signal = fsrv->child_kill_signal;
-    return FSRV_RUN_TMOUT;
 
+#ifndef AFL_PERSISTENT_RECORD
+    return FSRV_RUN_TMOUT;
+#else
+    retval = FSRV_RUN_TMOUT;
+    persistent_out_fmt = "%s/hangs/RECORD:%06u,cnt:%06u";
+    goto store_persistent_record;
+#endif
   }
 
   /* Did we crash?
@@ -1811,48 +1822,58 @@ afl_fsrv_run_target(afl_forkserver_t *fsrv, u32 timeout,
           (fsrv->uses_crash_exitcode &&
            WEXITSTATUS(fsrv->child_status) == fsrv->crash_exitcode))) {
 
-#ifdef AFL_PERSISTENT_RECORD
-    if (unlikely(fsrv->persistent_record)) {
+    /* For a proper crash, set last_kill_signal to WTERMSIG, else set it to 0 */
+    fsrv->last_kill_signal =
+        WIFSIGNALED(fsrv->child_status) ? WTERMSIG(fsrv->child_status) : 0;
 
-      char fn[PATH_MAX];
-      u32  i, writecnt = 0;
-      for (i = 0; i < fsrv->persistent_record; ++i) {
+#ifndef AFL_PERSISTENT_RECORD
+    return FSRV_RUN_CRASH;
+#else
+    retval = FSRV_RUN_CRASH;
+    persistent_out_fmt = "%s/crashes/RECORD:%06u,cnt:%06u";
+    goto store_persistent_record;
+#endif
 
-        u32 entry = (i + fsrv->persistent_record_idx) % fsrv->persistent_record;
-        u8 *data = fsrv->persistent_record_data[entry];
-        u32 len = fsrv->persistent_record_len[entry];
-        if (likely(len && data)) {
+  }
 
-          snprintf(fn, sizeof(fn), "%s/RECORD:%06u,cnt:%06u",
-                   fsrv->persistent_record_dir, fsrv->persistent_record_cnt,
-                   writecnt++);
-          int fd = open(fn, O_CREAT | O_TRUNC | O_WRONLY, 0644);
-          if (fd >= 0) {
+  /* success :) */
+  return FSRV_RUN_OK;
+
+#ifdef AFL_PERSISTENT_RECORD
+store_persistent_record:
+  if (unlikely(retval == FSRV_RUN_CRASH || retval == FSRV_RUN_TMOUT)  &&
+      unlikely(fsrv->persistent_record)) {
 
-            ck_write(fd, data, len, fn);
-            close(fd);
+    char fn[PATH_MAX];
+    u32  i, writecnt = 0;
+    for (i = 0; i < fsrv->persistent_record; ++i) {
 
-          }
+      u32 entry = (i + fsrv->persistent_record_idx) % fsrv->persistent_record;
+      u8 *data = fsrv->persistent_record_data[entry];
+      u32 len = fsrv->persistent_record_len[entry];
+      if (likely(len && data)) {
+
+        snprintf(fn, sizeof(fn), persistent_out_fmt,
+                 fsrv->persistent_record_dir, fsrv->persistent_record_cnt,
+                 writecnt++);
+        int fd = open(fn, O_CREAT | O_TRUNC | O_WRONLY, 0644);
+        if (fd >= 0) {
+
+          ck_write(fd, data, len, fn);
+          close(fd);
 
         }
 
       }
 
-      ++fsrv->persistent_record_cnt;
-
     }
 
-#endif
-
-    /* For a proper crash, set last_kill_signal to WTERMSIG, else set it to 0 */
-    fsrv->last_kill_signal =
-        WIFSIGNALED(fsrv->child_status) ? WTERMSIG(fsrv->child_status) : 0;
-    return FSRV_RUN_CRASH;
+    ++fsrv->persistent_record_cnt;
 
   }
 
-  /* success :) */
-  return FSRV_RUN_OK;
+  return retval;
+#endif
 
 }
 
diff --git a/src/afl-fuzz-init.c b/src/afl-fuzz-init.c
index 35932913..5b7dc4c1 100644
--- a/src/afl-fuzz-init.c
+++ b/src/afl-fuzz-init.c
@@ -1915,6 +1915,9 @@ static void handle_existing_out_dir(afl_state_t *afl) {
 
   }
 
+#ifdef AFL_PERSISTENT_RECORD
+  delete_files(fn, RECORD_PREFIX);
+#endif
   if (delete_files(fn, CASE_PREFIX)) { goto dir_cleanup_failed; }
   ck_free(fn);
 
@@ -1947,6 +1950,9 @@ static void handle_existing_out_dir(afl_state_t *afl) {
 
   }
 
+#ifdef AFL_PERSISTENT_RECORD
+  delete_files(fn, RECORD_PREFIX);
+#endif
   if (delete_files(fn, CASE_PREFIX)) { goto dir_cleanup_failed; }
   ck_free(fn);
 
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index 17949fd7..40c30472 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -2163,7 +2163,7 @@ int main(int argc, char **argv_orig, char **envp) {
 
     }
 
-    afl->fsrv.persistent_record_dir = alloc_printf("%s/crashes", afl->out_dir);
+    afl->fsrv.persistent_record_dir = alloc_printf("%s", afl->out_dir);
 
   }
 
diff --git a/utils/persistent_mode/Makefile b/utils/persistent_mode/Makefile
index e348c46c..64de82a7 100644
--- a/utils/persistent_mode/Makefile
+++ b/utils/persistent_mode/Makefile
@@ -1,10 +1,11 @@
 all:
 	../../afl-clang-fast -o persistent_demo persistent_demo.c
 	../../afl-clang-fast -o persistent_demo_new persistent_demo_new.c
+	gcc -g -I ../../include -o persistent_demo_new_compat persistent_demo_new.c
 	AFL_DONT_OPTIMIZE=1 ../../afl-clang-fast -o test-instr test-instr.c
 
 document:
 	AFL_DONT_OPTIMIZE=1 ../../afl-clang-fast -D_AFL_DOCUMENT_MUTATIONS -o test-instr test-instr.c
 
 clean:
-	rm -f persistent_demo persistent_demo_new test-instr
+	rm -f persistent_demo persistent_demo_new persistent_demo_new_compat test-instr
diff --git a/utils/persistent_mode/persistent_demo_new.c b/utils/persistent_mode/persistent_demo_new.c
index 285f50aa..40ada9e1 100644
--- a/utils/persistent_mode/persistent_demo_new.c
+++ b/utils/persistent_mode/persistent_demo_new.c
@@ -31,17 +31,8 @@
 
 /* this lets the source compile without afl-clang-fast/lto */
 #ifndef __AFL_FUZZ_TESTCASE_LEN
-
-ssize_t       fuzz_len;
-unsigned char fuzz_buf[1024000];
-
-  #define __AFL_FUZZ_TESTCASE_LEN fuzz_len
-  #define __AFL_FUZZ_TESTCASE_BUF fuzz_buf
-  #define __AFL_FUZZ_INIT() void sync(void);
-  #define __AFL_LOOP(x) \
-    ((fuzz_len = read(0, fuzz_buf, sizeof(fuzz_buf))) > 0 ? 1 : 0)
-  #define __AFL_INIT() sync()
-
+#define AFL_COMPAT
+#include "persistent_replay.h"
 #endif
 
 __AFL_FUZZ_INIT();
@@ -95,6 +86,8 @@ int main(int argc, char **argv) {
               if (buf[5] == '!') {
 
                 printf("six\n");
+                char *nullo = NULL+1;
+                *nullo = 'p';
                 abort();
 
               }