Merge pull request #1711 from atnwalk/atnwalk

AFL_POST_PROCESS_KEEP_ORIGINAL env variable for intermediate file formats and ATNwalk custom mutator

3 files changed, 32 insertions, 3 deletions

diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index 4d56f3a7..ac4fb4a9 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -135,10 +135,19 @@ write_to_testcase(afl_state_t *afl, void **mem, u32 len, u32 fix) {
 
     if (new_mem != *mem && new_mem != NULL && new_size > 0) {
 
-      u8 *new_buf = afl_realloc(AFL_BUF_PARAM(out_scratch), new_size);
+      new_buf = afl_realloc(AFL_BUF_PARAM(out_scratch), new_size);
       if (unlikely(!new_buf)) { PFATAL("alloc"); }
+      memcpy(new_buf, new_mem, new_size);
+
+      /* if AFL_POST_PROCESS_KEEP_ORIGINAL is set then save the original memory
+         prior post-processing in new_mem to restore it later */
+      if (unlikely(afl->afl_env.afl_post_process_keep_original)) {
+
+        new_mem = *mem;
+
+      }
+
       *mem = new_buf;
-      memcpy(*mem, new_mem, new_size);
       afl_swap_bufs(AFL_BUF_PARAM(out), AFL_BUF_PARAM(out_scratch));
 
     }
@@ -162,7 +171,18 @@ write_to_testcase(afl_state_t *afl, void **mem, u32 len, u32 fix) {
 
       /* everything as planned. use the potentially new data. */
       afl_fsrv_write_to_testcase(&afl->fsrv, *mem, new_size);
-      len = new_size;
+
+      if (likely(!afl->afl_env.afl_post_process_keep_original)) {
+
+        len = new_size;
+
+      } else {
+
+        /* restore the original memory which was saved in new_mem */
+        *mem = new_mem;
+        afl_swap_bufs(AFL_BUF_PARAM(out), AFL_BUF_PARAM(out_scratch));
+
+      }
 
     }
 
diff --git a/src/afl-fuzz-state.c b/src/afl-fuzz-state.c
index 46b67def..5e736029 100644
--- a/src/afl-fuzz-state.c
+++ b/src/afl-fuzz-state.c
@@ -394,6 +394,13 @@ void read_afl_environment(afl_state_t *afl, char **envp) {
             afl->afl_env.afl_statsd =
                 get_afl_env(afl_environment_variables[i]) ? 1 : 0;
 
+          } else if (!strncmp(env, "AFL_POST_PROCESS_KEEP_ORIGINAL",
+
+                              afl_environment_variable_len)) {
+
+            afl->afl_env.afl_post_process_keep_original =
+                get_afl_env(afl_environment_variables[i]) ? 1 : 0;
+
           } else if (!strncmp(env, "AFL_TMPDIR",
 
                               afl_environment_variable_len)) {
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index 394f33f9..ebdbb3fa 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -292,6 +292,8 @@ static void usage(u8 *argv0, int more_help) {
 
       PERSISTENT_MSG
 
+      "AFL_POST_PROCESS_KEEP_ORIGINAL: save the file as it was prior post-processing to the queue,\n"
+      "                                but execute the post-processed one\n"
       "AFL_PRELOAD: LD_PRELOAD / DYLD_INSERT_LIBRARIES settings for target\n"
       "AFL_TARGET_ENV: pass extra environment variables to target\n"
       "AFL_SHUFFLE_QUEUE: reorder the input queue randomly on startup\n"