reuse first_pass aux var in persistent loop for record replay mode, keep area ptr and loc logic intact in record replay mode, move replay record example to own dir in utils, update docs, move record compat layer to separate header file

8 files changed, 287 insertions, 136 deletions

diff --git a/include/persistent_replay.h b/include/afl-persistent-replay.h
index 58b22fb4..9e60ff9c 100644
--- a/include/persistent_replay.h
+++ b/include/afl-persistent-replay.h
@@ -10,6 +10,10 @@
 #include <sys/stat.h>
 #include <fcntl.h>
 
+#ifndef PATH_MAX
+  #define PATH_MAX 4096
+#endif
+
 static unsigned short int is_replay_record;
 static unsigned int       replay_record;
 static unsigned int       replay_record_cnt;
@@ -23,7 +27,7 @@ static char **record_arg = NULL;
 
 static int select_files(const struct dirent *dirbuf) {
 
-  char fn[4096];
+  char fn[PATH_MAX];
 
   if (dirbuf->d_name[0] == '.') {
 
@@ -123,85 +127,5 @@ __attribute__((constructor)) static void __afl_record_replay_init(
 
 }
 
-/* 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);
-
-  #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);
-
-    }
-
-  } else {
-
-    if (!inited) {
-
-      cycle_cnt = max_cnt;
-      inited = 1;
-
-    }
-
-  }
-
-  return cycle_cnt--;
-
-}
-
-#endif  // AFL_COMPAT
-
 #endif  // _HAVE_PERSISTENT_REPLAY_H
 
diff --git a/include/afl-record-compat.h b/include/afl-record-compat.h
new file mode 100644
index 00000000..2c79595d
--- /dev/null
+++ b/include/afl-record-compat.h
@@ -0,0 +1,67 @@
+#ifndef _HAVE_AFL_COMPAT_H
+#define _HAVE_AFL_COMPAT_H
+
+#include <afl-persistent-replay.h>
+
+#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);
+
+#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);
+
+    }
+
+  }
+
+  return --cycle_cnt;
+
+}
+
+#endif  // _HAVE_AFL_COMPAT_H
+
diff --git a/include/config.h b/include/config.h
index d44cda9c..a5b6eba1 100644
--- a/include/config.h
+++ b/include/config.h
@@ -83,13 +83,11 @@
    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
-
-/* Adds support in compiler-rt to replay persistent records */
-#define AFL_PERSISTENT_REPLAY
+// #define AFL_PERSISTENT_RECORD
 
 /* 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
diff --git a/instrumentation/README.persistent_mode.md b/instrumentation/README.persistent_mode.md
index b5d982b0..8e4f6ae4 100644
--- a/instrumentation/README.persistent_mode.md
+++ b/instrumentation/README.persistent_mode.md
@@ -201,26 +201,28 @@ And that is all!
 
 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`).
+When `AFL_PERSISTENT_RECORD` is enabled, replay functionality is also included in the compiler-rt library. To replay a specific record, assign the record number to the AFL_PERSISTENT_REPLAY environment variable (i.e., `RECORD:XXXXX`` -> `AFL_PERSISTENT_REPLAY=XXXXX`), and run the test binary as you would normally do.
 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.
+If your harness reads the input files from arguments using the special `@@` argument you will need to include support by enabling `AFL_PERSISTENT_ARGPARSE` in  `config.h`.
+
 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
+## 7) Drop-in persistent loop replay replacement
 
-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.
+To use the replay functionality without having to use `afl-cc`, include the [include/record_compat.h](../include/afl-record_compat.h) header file. Together with the [include/afl-persistent-replay.h](../include/afl-persistent-replay.h) header included in it, `afl-record-compat.h` 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"
+  #include "afl-record-compat.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
+A simple example is provided in [persistent_demo_replay.c](../utils/replay_record/persistent_demo_replay.c).
+
+Be aware that the [afl-record-compat.h](../include/afl-record-compat.h) header should only be included in a single compilation unit, or you will end up with clobbered functions and variables.
+
+If you need a cleaner solution, you'll have to move the functions and variables defined in [include/record_compat.h](../include/afl-record-compat.h) and [include/afl-persistent-replay.h](../include/afl-persistent-replay.h) in a C file, and add the relevant declarations to a header file. After including the new header file, the compilation unit resulting from compiling the C file can then be linked with your program.
\ No newline at end of file
diff --git a/instrumentation/afl-compiler-rt.o.c b/instrumentation/afl-compiler-rt.o.c
index 037caaf0..4c5d4e79 100644
--- a/instrumentation/afl-compiler-rt.o.c
+++ b/instrumentation/afl-compiler-rt.o.c
@@ -83,8 +83,8 @@
 #include <sys/mman.h>
 #include <fcntl.h>
 
-#ifdef AFL_PERSISTENT_REPLAY
-  #include "persistent_replay.h"
+#ifdef AFL_PERSISTENT_RECORD
+  #include "afl-persistent-replay.h"
 #endif
 
 /* Globals needed by the injected instrumentation. The __afl_area_initial region
@@ -1342,68 +1342,73 @@ int __afl_persistent_loop(unsigned int max_cnt) {
   static u8  first_pass = 1;
   static u32 cycle_cnt;
 
-#ifdef AFL_PERSISTENT_REPLAY
+#ifdef AFL_PERSISTENT_RECORD
+  char tcase[PATH_MAX];
+#endif
 
-  #ifndef PATH_MAX
-    #define PATH_MAX 4096
-  #endif
+  if (first_pass) {
 
-  static u8 inited = 0;
-  char      tcase[PATH_MAX];
+    /* 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
+       iteration, it's our job to erase any trace of whatever happened
+       before the loop. */
 
-  if (unlikely(is_replay_record)) {
+    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));
 
-    if (!inited) {
+    first_pass = 0;
+    __afl_selective_coverage_temp = 1;
+
+#ifdef AFL_PERSISTENT_RECORD
+    if (unlikely(is_replay_record)) {
 
       cycle_cnt = replay_record_cnt;
-      inited = 1;
+      goto persistent_record;
 
-    }
+    } else
 
-    snprintf(tcase, PATH_MAX, "%s/%s",
-             replay_record_dir ? replay_record_dir : "./",
-             record_list[replay_record_cnt - cycle_cnt]->d_name);
+#endif
+    {
 
-  #ifdef AFL_PERSISTENT_REPLAY_ARGPARSE
-    if (record_arg) {
+      cycle_cnt = max_cnt;
 
-      *record_arg = tcase;
+    }
 
-    } else
+    return 1;
 
-  #endif  // AFL_PERSISTENT_REPLAY_ARGPARSE
-    {
+  } else if (--cycle_cnt) {
 
-      int fd = open(tcase, O_RDONLY);
-      dup2(fd, 0);
-      close(fd);
+#ifdef AFL_PERSISTENT_RECORD
+    if (unlikely(is_replay_record)) {
 
-    }
+    persistent_record:
 
-    return cycle_cnt--;
+      snprintf(tcase, PATH_MAX, "%s/%s",
+               replay_record_dir ? replay_record_dir : "./",
+               record_list[replay_record_cnt - cycle_cnt]->d_name);
 
-  } else
+  #ifdef AFL_PERSISTENT_REPLAY_ARGPARSE
+      if (unlikely(record_arg)) {
 
-#endif
+        *record_arg = tcase;
 
-      if (first_pass) {
+      } else
 
-    /* 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
-       iteration, it's our job to erase any trace of whatever happened
-       before the loop. */
+  #endif  // AFL_PERSISTENT_REPLAY_ARGPARSE
+      {
 
-    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));
+        int fd = open(tcase, O_RDONLY);
+        dup2(fd, 0);
+        close(fd);
 
-    cycle_cnt = max_cnt;
-    first_pass = 0;
-    __afl_selective_coverage_temp = 1;
+      }
 
-    return 1;
+      return 1;
 
-  } else if (--cycle_cnt) {
+    }
+
+#endif
 
     raise(SIGSTOP);
 
diff --git a/utils/persistent_mode/Makefile b/utils/persistent_mode/Makefile
index 64de82a7..498aa3f8 100644
--- a/utils/persistent_mode/Makefile
+++ b/utils/persistent_mode/Makefile
@@ -1,7 +1,6 @@
 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:
diff --git a/utils/replay_record/Makefile b/utils/replay_record/Makefile
new file mode 100644
index 00000000..0d1cba92
--- /dev/null
+++ b/utils/replay_record/Makefile
@@ -0,0 +1,8 @@
+all:
+	test `grep '//[\s\t ]*#define[\s\t ]*AFL_PERSISTENT_RECORD' ../../include/config.h | wc -l` -eq 0 || (echo "AFL_PERSISTENT_RECORD must be enabled in config.h"; exit 1)
+	../../afl-clang-fast -o persistent_demo_replay persistent_demo_replay.c
+	${CC} -I ../../include -o persistent_demo_replay_compat persistent_demo_replay.c
+	${CC} -g -I ../../include -DAFL_PERSISTENT_REPLAY_ARGPARSE -o persistent_demo_replay_argparse persistent_demo_replay.c
+
+clean:
+	rm -f persistent_demo_replay persistent_demo_replay_argparse persistent_demo_replay_compat
diff --git a/utils/replay_record/persistent_demo_replay.c b/utils/replay_record/persistent_demo_replay.c
new file mode 100644
index 00000000..6f6648f1
--- /dev/null
+++ b/utils/replay_record/persistent_demo_replay.c
@@ -0,0 +1,148 @@
+/*
+   american fuzzy lop++ - persistent mode example
+   --------------------------------------------
+
+   Originally written by Michal Zalewski
+
+   Copyright 2015 Google Inc. All rights reserved.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at:
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+   This file demonstrates the high-performance "persistent mode" that may be
+   suitable for fuzzing certain fast and well-behaved libraries, provided that
+   they are stateless or that their internal state can be easily reset
+   across runs.
+
+   To make this work, the library and this shim need to be compiled in LLVM
+   mode using afl-clang-fast (other compiler wrappers will *not* work).
+
+ */
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <unistd.h>
+#include <signal.h>
+#include <string.h>
+#include <limits.h>
+
+#ifdef AFL_PERSISTENT_REPLAY_ARGPARSE
+  #include <sys/stat.h>
+  #include <fcntl.h>
+#endif
+
+/* this lets the source compile without afl-clang-fast/lto */
+#ifndef __AFL_FUZZ_TESTCASE_LEN
+  #include "afl-record-compat.h"
+#endif
+
+__AFL_FUZZ_INIT();
+
+/* Main entry point. */
+
+/* To ensure checks are not optimized out it is recommended to disable
+   code optimization for the fuzzer harness main() */
+#pragma clang optimize off
+#pragma GCC            optimize("O0")
+
+int main(int argc, char **argv) {
+
+  ssize_t        len;                        /* how much input did we read? */
+  unsigned char *buf;                        /* test case buffer pointer    */
+
+#ifdef AFL_PERSISTENT_REPLAY_ARGPARSE
+  int fd;
+
+  if (argc < 2) { printf("Need an input file!"); }
+#endif
+
+  /* The number passed to __AFL_LOOP() controls the maximum number of
+     iterations before the loop exits and the program is allowed to
+     terminate normally. This limits the impact of accidental memory leaks
+     and similar hiccups. */
+
+  __AFL_INIT();
+
+#ifdef AFL_PERSISTENT_REPLAY_ARGPARSE
+  buf = malloc(1000);
+#else
+  buf = __AFL_FUZZ_TESTCASE_BUF;  // this must be assigned before __AFL_LOOP!
+#endif
+
+  while (__AFL_LOOP(UINT_MAX)) {  // increase if you have good stability
+
+#ifdef AFL_PERSISTENT_REPLAY_ARGPARSE
+    fd = open(argv[1], O_RDONLY);
+    len = read(fd, buf, 1000);
+    close(fd);
+#else
+    len = __AFL_FUZZ_TESTCASE_LEN;  // do not use the macro directly in a call!
+#endif
+
+    // fprintf(stderr, "input: %zd \"%s\"\n", len, buf);
+
+    /* do we have enough data? */
+    if (len < 8) continue;
+
+    if (strcmp((char *)buf, "thisisateststring") == 0) printf("teststring\n");
+
+    if (buf[0] == 'f') {
+
+      printf("one\n");
+      if (buf[1] == 'o') {
+
+        printf("two\n");
+        if (buf[2] == 'o') {
+
+          printf("three\n");
+          if (buf[3] == '!') {
+
+            printf("four\n");
+            if (buf[4] == '!') {
+
+              printf("five\n");
+              if (buf[5] == '!') {
+
+                printf("six\n");
+                abort();
+
+              } else {
+
+                if (buf[5] == 'O') {
+
+                  // hang
+                  while (1) {
+
+                    continue;
+
+                  };
+
+                }
+
+              }
+
+            }
+
+          }
+
+        }
+
+      }
+
+    }
+
+    /*** END PLACEHOLDER CODE ***/
+
+  }
+
+  /* Once the loop is exited, terminate normally - AFL will restart the process
+     when this happens, with a clean slate when it comes to allocated memory,
+     leftover file descriptors, etc. */
+
+  return 0;
+
+}
+