Fix memory errors when trim causes testcase growth (#881) (#903)

* Revert "fixed potential double free in custom trim (#881)"

This reverts commit e9d2f72382cab75832721d859c3e731da071435d.

* Revert "fix custom trim for increasing data"

This reverts commit 86a8ef168dda766d2f25f15c15c4d3ecf21d0667.

* Fix memory errors when trim causes testcase growth

Modify trim_case_custom to avoid writing into in_buf because
some custom mutators can cause the testcase to grow rather than
shrink.

Instead of modifying in_buf directly, we write the update out
to the disk when trimming is complete, and then the caller is
responsible for refreshing the in-memory buffer from the file.

This is still a bit sketchy because it does need to modify q->len in
order to notify the upper layers that something changed, and it could
end up telling upper layer code that the q->len is *bigger* than
the buffer (q->testcase_buf) that contains it, which is asking
for trouble down the line somewhere...

* Fix an unlikely situation

Put back some `unlikely()` calls that were in
the e9d2f72382cab75832721d859c3e731da071435d commit that was
reverted.

4 files changed, 37 insertions, 44 deletions

diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index 040d7ae9..f201782a 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -1003,7 +1003,7 @@ void read_afl_environment(afl_state_t *, char **);
 /* Custom mutators */
 void setup_custom_mutators(afl_state_t *);
 void destroy_custom_mutators(afl_state_t *);
-u8   trim_case_custom(afl_state_t *, struct queue_entry *q, u8 **in_buf,
+u8   trim_case_custom(afl_state_t *, struct queue_entry *q, u8 *in_buf,
                       struct custom_mutator *mutator);
 
 /* Python */
@@ -1093,7 +1093,7 @@ fsrv_run_result_t fuzz_run_target(afl_state_t *, afl_forkserver_t *fsrv, u32);
 void              write_to_testcase(afl_state_t *, void *, u32);
 u8   calibrate_case(afl_state_t *, struct queue_entry *, u8 *, u32, u8);
 void sync_fuzzers(afl_state_t *);
-u8   trim_case(afl_state_t *, struct queue_entry *, u8 **);
+u8   trim_case(afl_state_t *, struct queue_entry *, u8 *);
 u8   common_fuzz_stuff(afl_state_t *, u8 *, u32);
 
 /* Fuzz one */
diff --git a/src/afl-fuzz-mutators.c b/src/afl-fuzz-mutators.c
index d8db8676..3bb37a89 100644
--- a/src/afl-fuzz-mutators.c
+++ b/src/afl-fuzz-mutators.c
@@ -305,16 +305,14 @@ struct custom_mutator *load_custom_mutator(afl_state_t *afl, const char *fn) {
 
 }
 
-// Custom testcase trimming.
-u8 trim_case_custom(afl_state_t *afl, struct queue_entry *q, u8 **in_buf_p,
+u8 trim_case_custom(afl_state_t *afl, struct queue_entry *q, u8 *in_buf,
                     struct custom_mutator *mutator) {
 
-  // We need to pass pointers around, as growing testcases may need to realloc.
-  u8 *in_buf = *in_buf_p;
-
-  u8  needs_write = 0, fault = 0;
+  u8  fault = 0;
   u32 trim_exec = 0;
   u32 orig_len = q->len;
+  u32 out_len = 0;
+  u8* out_buf = NULL;
 
   u8 val_buf[STRINGIFY_VAL_SIZE_MAX];
 
@@ -401,40 +399,33 @@ u8 trim_case_custom(afl_state_t *afl, struct queue_entry *q, u8 **in_buf_p,
 
     if (likely(retlen && cksum == q->exec_cksum)) {
 
-      // Check if we got a new retbuf and to memcpy our buf.
-      if (in_buf != retbuf) {
-
-        if (afl_realloc((void **)in_buf_p, retlen) == NULL) {
-
-          FATAL("can not allocate memory for trim");
-
-        }
+      /* Let's save a clean trace, which will be needed by
+         update_bitmap_score once we're done with the trimming stuff.
+         Use out_buf NULL check to make this only happen once per trim. */
 
-        in_buf = *in_buf_p;
+      if (!out_buf) {
 
-        memcpy(in_buf, retbuf, retlen);
-        q->len = retlen;
+        memcpy(afl->clean_trace_custom, afl->fsrv.trace_bits,
+               afl->fsrv.map_size);
 
       }
 
-      /* Let's save a clean trace, which will be needed by
-         update_bitmap_score once we're done with the trimming stuff. */
-
-      if (!needs_write) {
+      if (afl_realloc((void **)&out_buf, retlen) == NULL) {
 
-        needs_write = 1;
-        memcpy(afl->clean_trace_custom, afl->fsrv.trace_bits,
-               afl->fsrv.map_size);
+        FATAL("can not allocate memory for trim");
 
       }
 
+      out_len = retlen;
+      memcpy(out_buf, retbuf, retlen);
+
       /* Tell the custom mutator that the trimming was successful */
       afl->stage_cur = mutator->afl_custom_post_trim(mutator->data, 1);
 
       if (afl->not_on_tty && afl->debug) {
 
         SAYF("[Custom Trimming] SUCCESS: %u/%u iterations (now at %u bytes)",
-             afl->stage_cur, afl->stage_max, q->len);
+             afl->stage_cur, afl->stage_max, out_len);
 
       }
 
@@ -467,16 +458,10 @@ u8 trim_case_custom(afl_state_t *afl, struct queue_entry *q, u8 **in_buf_p,
 
   }
 
-  if (afl->not_on_tty && afl->debug) {
-
-    SAYF("[Custom Trimming] DONE: %u bytes -> %u bytes", orig_len, q->len);
-
-  }
-
-  /* If we have made changes to in_buf, we also need to update the on-disk
+  /* If we have made changes, we also need to update the on-disk
      version of the test case. */
 
-  if (needs_write) {
+  if (out_buf) {
 
     s32 fd;
 
@@ -486,16 +471,28 @@ u8 trim_case_custom(afl_state_t *afl, struct queue_entry *q, u8 **in_buf_p,
 
     if (fd < 0) { PFATAL("Unable to create '%s'", q->fname); }
 
-    ck_write(fd, in_buf, q->len, q->fname);
+    ck_write(fd, out_buf, out_len, q->fname);
     close(fd);
 
+    /* Update the queue's knowledge of length as soon as we write the file.
+       We do this here so that exit/error cases that *don't* update the file also
+       don't update q->len. */
+    q->len = out_len;
+
     memcpy(afl->fsrv.trace_bits, afl->clean_trace_custom, afl->fsrv.map_size);
     update_bitmap_score(afl, q);
 
   }
 
+  if (afl->not_on_tty && afl->debug) {
+
+    SAYF("[Custom Trimming] DONE: %u bytes -> %u bytes", orig_len, q->len);
+
+  }
+
 abort_trimming:
 
+  if (out_buf) afl_free(out_buf);
   afl->bytes_trim_out += q->len;
   return fault;
 
diff --git a/src/afl-fuzz-one.c b/src/afl-fuzz-one.c
index ed815cb4..4eeb93de 100644
--- a/src/afl-fuzz-one.c
+++ b/src/afl-fuzz-one.c
@@ -508,7 +508,7 @@ u8 fuzz_one_original(afl_state_t *afl) {
 
     u32 old_len = afl->queue_cur->len;
 
-    u8 res = trim_case(afl, afl->queue_cur, &in_buf);
+    u8 res = trim_case(afl, afl->queue_cur, in_buf);
     orig_in = in_buf = queue_testcase_get(afl, afl->queue_cur);
 
     if (unlikely(res == FSRV_RUN_ERROR)) {
@@ -3007,7 +3007,7 @@ static u8 mopt_common_fuzzing(afl_state_t *afl, MOpt_globals_t MOpt_globals) {
 
     u32 old_len = afl->queue_cur->len;
 
-    u8 res = trim_case(afl, afl->queue_cur, &in_buf);
+    u8 res = trim_case(afl, afl->queue_cur, in_buf);
     orig_in = in_buf = queue_testcase_get(afl, afl->queue_cur);
 
     if (unlikely(res == FSRV_RUN_ERROR)) {
diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index 397d62bf..6e5210b8 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -718,10 +718,7 @@ void sync_fuzzers(afl_state_t *afl) {
    trimmer uses power-of-two increments somewhere between 1/16 and 1/1024 of
    file size, to keep the stage short and sweet. */
 
-u8 trim_case(afl_state_t *afl, struct queue_entry *q, u8 **in_buf_p) {
-
-  // We need to pass pointers around, as growing testcases may need to realloc.
-  u8 *in_buf = *in_buf_p;
+u8 trim_case(afl_state_t *afl, struct queue_entry *q, u8 *in_buf) {
 
   u32 orig_len = q->len;
 
@@ -735,8 +732,7 @@ u8 trim_case(afl_state_t *afl, struct queue_entry *q, u8 **in_buf_p) {
 
       if (el->afl_custom_trim) {
 
-        trimmed_case = trim_case_custom(afl, q, in_buf_p, el);
-        in_buf = *in_buf_p;
+        trimmed_case = trim_case_custom(afl, q, in_buf, el);
         custom_trimmed = true;
 
       }