fix custom mutator examples

4 files changed, 31 insertions, 37 deletions

diff --git a/custom_mutators/examples/example.c b/custom_mutators/examples/example.c
index 5c174e10..3f299508 100644
--- a/custom_mutators/examples/example.c
+++ b/custom_mutators/examples/example.c
@@ -352,7 +352,7 @@ uint8_t afl_custom_queue_get(my_mutator_t *data, const uint8_t *filename) {
  * @return if the file contents was modified return 1 (True), 0 (False)
  *         otherwise
  */
-uint8_t afl_custom_queue_new_entry(my_mutator_t * data,
+uint8_t afl_custom_queue_new_entry(my_mutator_t  *data,
                                    const uint8_t *filename_new_queue,
                                    const uint8_t *filename_orig_queue) {
 
diff --git a/custom_mutators/examples/post_library_gif.so.c b/custom_mutators/examples/post_library_gif.so.c
index aec05720..9cd224f4 100644
--- a/custom_mutators/examples/post_library_gif.so.c
+++ b/custom_mutators/examples/post_library_gif.so.c
@@ -72,6 +72,7 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include "alloc-inl.h"
 
 /* Header that must be present at the beginning of every test case: */
 
@@ -127,9 +128,11 @@ size_t afl_custom_post_process(post_state_t *data, unsigned char *in_buf,
   }
 
   /* Allocate memory for new buffer, reusing previous allocation if
-     possible. */
+     possible. Note we have to use afl-fuzz's own realloc!
+     Note that you should only do this if you need to grow the buffer,
+     otherwise work with in_buf, and assign it to *out_buf instead. */
 
-  *out_buf = realloc(data->buf, len);
+  *out_buf = afl_realloc(out_buf, len);
 
   /* If we're out of memory, the most graceful thing to do is to return the
      original buffer and give up on modifying it. Let AFL handle OOM on its
@@ -142,9 +145,9 @@ size_t afl_custom_post_process(post_state_t *data, unsigned char *in_buf,
 
   }
 
-  /* Copy the original data to the new location. */
-
-  memcpy(*out_buf, in_buf, len);
+  if (len > strlen(HEADER))
+    memcpy(*out_buf + strlen(HEADER), in_buf + strlen(HEADER),
+           len - strlen(HEADER));
 
   /* Insert the new header. */
 
diff --git a/custom_mutators/examples/post_library_png.so.c b/custom_mutators/examples/post_library_png.so.c
index 941f7e55..cd65b1bc 100644
--- a/custom_mutators/examples/post_library_png.so.c
+++ b/custom_mutators/examples/post_library_png.so.c
@@ -29,8 +29,8 @@
 #include <stdint.h>
 #include <string.h>
 #include <zlib.h>
-
 #include <arpa/inet.h>
+#include "alloc-inl.h"
 
 /* A macro to round an integer up to 4 kB. */
 
@@ -70,9 +70,6 @@ size_t afl_custom_post_process(post_state_t *data, const unsigned char *in_buf,
                                unsigned int          len,
                                const unsigned char **out_buf) {
 
-  unsigned char *new_buf = (unsigned char *)in_buf;
-  unsigned int   pos = 8;
-
   /* Don't do anything if there's not enough room for the PNG header
      (8 bytes). */
 
@@ -83,6 +80,22 @@ size_t afl_custom_post_process(post_state_t *data, const unsigned char *in_buf,
 
   }
 
+  /* This is not a good way to do it, if you do not need to grow the buffer
+     then just work with in_buf instead for speed reasons.
+     But we want to show how to grow a buffer, so this is how it's done: */
+
+  unsigned int   pos = 8;
+  unsigned char *new_buf = afl_realloc(out_buf, UP4K(len));
+
+  if (!new_buf) {
+
+    *out_buf = in_buf;
+    return len;
+
+  }
+
+  memcpy(new_buf, in_buf, len);
+
   /* Minimum size of a zero-length PNG chunk is 12 bytes; if we
      don't have that, we can bail out. */
 
@@ -111,33 +124,6 @@ size_t afl_custom_post_process(post_state_t *data, const unsigned char *in_buf,
 
     if (real_cksum != file_cksum) {
 
-      /* First modification? Make a copy of the input buffer. Round size
-         up to 4 kB to minimize the number of reallocs needed. */
-
-      if (new_buf == in_buf) {
-
-        if (len <= data->size) {
-
-          new_buf = data->buf;
-
-        } else {
-
-          new_buf = realloc(data->buf, UP4K(len));
-          if (!new_buf) {
-
-            *out_buf = in_buf;
-            return len;
-
-          }
-
-          data->buf = new_buf;
-          data->size = UP4K(len);
-          memcpy(new_buf, in_buf, len);
-
-        }
-
-      }
-
       *(uint32_t *)(new_buf + pos + 8 + chunk_len) = real_cksum;
 
     }
diff --git a/docs/custom_mutators.md b/docs/custom_mutators.md
index d84e4e02..6f3353ec 100644
--- a/docs/custom_mutators.md
+++ b/docs/custom_mutators.md
@@ -38,6 +38,11 @@ performed with the custom mutator.
 
 ## 2) APIs
 
+**IMPORTANT NOTE**: If you use our C/C++ API and you want to increase the size
+of an **out_buf buffer, you have to use `afl_realloc()` for this, so include
+`include/alloc-inl.h` - otherwise afl-fuzz will crash when trying to free
+your buffers.
+
 C/C++:
 
 ```c