2 files changed, 57 insertions, 50 deletions
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);