fixed AFL_POST_PROCESS_KEEP_ORIGINAL for version 4.07a

1 files changed, 20 insertions, 10 deletions

diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index 2d53de93..ac4fb4a9 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -133,14 +133,22 @@ write_to_testcase(afl_state_t *afl, void **mem, u32 len, u32 fix) {
 
     }
 
-    if (new_mem != *mem && new_mem != NULL && new_size > 0
-        && !afl->afl_env.afl_post_process_keep_original) {
+    if (new_mem != *mem && new_mem != NULL && new_size > 0) {
 
-        u8 *new_buf = afl_realloc(AFL_BUF_PARAM(out_scratch), new_size);
-        if (unlikely(!new_buf)) { PFATAL("alloc"); }
-        *mem = new_buf;
-        memcpy(*mem, new_mem, new_size);
-        afl_swap_bufs(AFL_BUF_PARAM(out), AFL_BUF_PARAM(out_scratch));
+      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;
+      afl_swap_bufs(AFL_BUF_PARAM(out), AFL_BUF_PARAM(out_scratch));
 
     }
 
@@ -162,17 +170,19 @@ write_to_testcase(afl_state_t *afl, void **mem, u32 len, u32 fix) {
     if (likely(!sent)) {
 
       /* everything as planned. use the potentially new data. */
+      afl_fsrv_write_to_testcase(&afl->fsrv, *mem, new_size);
 
       if (likely(!afl->afl_env.afl_post_process_keep_original)) {
 
-        afl_fsrv_write_to_testcase(&afl->fsrv, *mem, new_size);
+        len = new_size;
 
       } else {
 
-        afl_fsrv_write_to_testcase(&afl->fsrv, new_mem, new_size);
+        /* 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));
 
       }
-      len = new_size;
 
     }