edited custom mutator pre_save api

6 files changed, 140 insertions, 27 deletions

diff --git a/docs/custom_mutators.md b/docs/custom_mutators.md
index 14d8f518..ecbd50ba 100644
--- a/docs/custom_mutators.md
+++ b/docs/custom_mutators.md
@@ -30,7 +30,7 @@ C/C++:
 void afl_custom_init(unsigned int seed);
 size_t afl_custom_fuzz(uint8_t** buf, size_t buf_size, uint8_t* add_buf,
                        size_t add_buf_size, size_t max_size);
-size_t afl_custom_pre_save(uint8_t* buf, size_t buf_size, uint8_t** out_buf);
+size_t afl_custom_write_to_testcase(uint8_t* buf, size_t buf_size, uint8_t** out_buf);
 uint32_t afl_custom_init_trim(uint8_t* buf, size_t buf_size);
 void afl_custom_trim(uint8_t** out_buf, size_t* out_buf_size);
 uint32_t afl_custom_post_trim(uint8_t success);
diff --git a/examples/custom_mutators/example.c b/examples/custom_mutators/example.c
index 2df17dec..3e708db8 100644
--- a/examples/custom_mutators/example.c
+++ b/examples/custom_mutators/example.c
@@ -12,6 +12,7 @@
 #include <stdint.h>
 #include <stdlib.h>
 #include <string.h>
+#include <stdio.h>
 
 #define DATA_SIZE (100)
 
@@ -112,21 +113,29 @@ size_t afl_custom_fuzz(my_mutator_t *data, uint8_t **buf, size_t buf_size,
  * @param[in] data pointer returned in afl_custom_init for this fuzz case
  * @param[in] buf Buffer containing the test case to be executed
  * @param[in] buf_size Size of the test case
- * @param[out] out_buf Pointer to the buffer containing the test case after
+ * @param[in] out_buf Pointer to the buffer containing the test case after
  *     processing. External library should allocate memory for out_buf. AFL++
  *     will release the memory after saving the test case.
- * @return Size of the output buffer after processing
+ *     out_buf will always be at least as large as buf.
+ * @param[in] out_buf_size The maximum size we may use.
+ *            In case we need to have this bigger, simply return that.
+ * @return Size of the output buffer after processing or the needed amount.
+ *         return 0 to indicate the original buf should be used.
  */
 size_t afl_custom_pre_save(my_mutator_t *data, uint8_t *buf, size_t buf_size,
-                           uint8_t **out_buf) {
+                           uint8_t *out_buf, size_t out_buf_size) {
 
-  size_t out_buf_size;
+  // In case we need more than out_buf_size, we return that amount and get
+  // called again.
+  if (out_buf_size < 32000) return 32000;
 
+  memcpy(out_buf, buf, buf_size);
   out_buf_size = buf_size;
-
-  // External mutator should allocate memory for `out_buf`
-  *out_buf = malloc(out_buf_size);
-  memcpy(*out_buf, buf, out_buf_size);
+  out_buf[0] = 'A';
+  out_buf[1] = 'F';
+  out_buf[2] = 'L';
+  out_buf[3] = '+';
+  out_buf[4] = '+';
 
   return out_buf_size;
 
diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index d610ac29..8bf66403 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -30,6 +30,9 @@
 #define AFL_MAIN
 #define MESSAGES_TO_STDOUT
 
+/* We preallocate a buffer of this size for afl_custom_pre_save */
+#define PRE_SAVE_BUF_INIT_SIZE (16384)
+
 #ifndef _GNU_SOURCE
 #define _GNU_SOURCE 1
 #endif
@@ -293,6 +296,9 @@ typedef struct py_mutator {
   void *    afl_state;
   void *    py_data;
 
+  PyObject *scratch_buf;
+  size_t    scratch_size;
+
 } py_mutator_t;
 
 #endif
@@ -591,6 +597,8 @@ struct custom_mutator {
 
   const char *name;
   void *      dh;
+  u8 *        pre_save_buf;
+  size_t      pre_save_size;
 
   void *data;                                    /* custom mutator data ptr */
 
@@ -639,7 +647,7 @@ struct custom_mutator {
    * @return Size of the output buffer after processing
    */
   size_t (*afl_custom_pre_save)(void *data, u8 *buf, size_t buf_size,
-                                u8 **out_buf);
+                                u8 *out_buf, size_t out_buf_size);
 
   /**
    * This method is called at the start of each trimming operation and receives
@@ -775,7 +783,7 @@ u8   trim_case_custom(afl_state_t *, struct queue_entry *q, u8 *in_buf);
 
 void finalize_py_module(void *);
 
-size_t pre_save_py(void *, u8 *, size_t, u8 **);
+size_t pre_save_py(void *, u8 *, size_t, u8 *, size_t);
 u32    init_trim_py(void *, u8 *, size_t);
 u32    post_trim_py(void *, u8);
 void   trim_py(void *, u8 **, size_t *);
diff --git a/src/afl-fuzz-mutators.c b/src/afl-fuzz-mutators.c
index 0b0c3394..fac97f8e 100644
--- a/src/afl-fuzz-mutators.c
+++ b/src/afl-fuzz-mutators.c
@@ -80,6 +80,14 @@ void destroy_custom_mutator(afl_state_t *afl) {
 
     if (afl->mutator->dh) dlclose(afl->mutator->dh);
 
+    if (afl->mutator->pre_save_buf) {
+
+      ck_free(afl->mutator->pre_save_buf);
+      afl->mutator->pre_save_buf = NULL;
+      afl->mutator->pre_save_size = 0;
+
+    }
+
     ck_free(afl->mutator);
     afl->mutator = NULL;
 
@@ -91,6 +99,8 @@ void load_custom_mutator(afl_state_t *afl, const char *fn) {
 
   void *dh;
   afl->mutator = ck_alloc(sizeof(struct custom_mutator));
+  afl->mutator->pre_save_buf = NULL;
+  afl->mutator->pre_save_size = 0;
 
   afl->mutator->name = fn;
   ACTF("Loading custom mutator library from '%s'...", fn);
@@ -125,9 +135,18 @@ void load_custom_mutator(afl_state_t *afl, const char *fn) {
 
   /* "afl_custom_pre_save", optional */
   afl->mutator->afl_custom_pre_save = dlsym(dh, "afl_custom_pre_save");
-  if (!afl->mutator->afl_custom_pre_save)
+  if (!afl->mutator->afl_custom_pre_save) {
+
     WARNF("Symbol 'afl_custom_pre_save' not found.");
 
+  } else {
+
+    /* if we have a pre_save hook, prealloc some memory. */
+    afl->mutator->pre_save_buf = ck_alloc(PRE_SAVE_BUF_INIT_SIZE * sizeof(u8));
+    afl->mutator->pre_save_size = PRE_SAVE_BUF_INIT_SIZE;
+
+  }
+
   u8 notrim = 0;
   /* "afl_custom_init_trim", optional */
   afl->mutator->afl_custom_init_trim = dlsym(dh, "afl_custom_init_trim");
diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
index 418497d0..b053e8d5 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -287,6 +287,8 @@ void deinit_py(void *py_mutator) {
 void load_custom_mutator_py(afl_state_t *afl, char *module_name) {
 
   afl->mutator = ck_alloc(sizeof(struct custom_mutator));
+  afl->mutator->pre_save_buf = NULL;
+  afl->mutator->pre_save_size = 0;
 
   afl->mutator->name = module_name;
   ACTF("Loading Python mutator library from '%s'...", module_name);
@@ -305,8 +307,14 @@ void load_custom_mutator_py(afl_state_t *afl, char *module_name) {
      is quite different from the custom mutator. */
   afl->mutator->afl_custom_fuzz = fuzz_py;
 
-  if (py_functions[PY_FUNC_PRE_SAVE])
+  if (py_functions[PY_FUNC_PRE_SAVE]) {
+
     afl->mutator->afl_custom_pre_save = pre_save_py;
+    /* if we have a pre_save hook, prealloc some memory. */
+    afl->mutator->pre_save_buf = ck_alloc(PRE_SAVE_BUF_INIT_SIZE * sizeof(u8));
+    afl->mutator->pre_save_size = PRE_SAVE_BUF_INIT_SIZE;
+
+  }
 
   if (py_functions[PY_FUNC_INIT_TRIM])
     afl->mutator->afl_custom_init_trim = init_trim_py;
@@ -336,10 +344,33 @@ void load_custom_mutator_py(afl_state_t *afl, char *module_name) {
 
 }
 
-size_t pre_save_py(void *py_mutator, u8 *buf, size_t buf_size, u8 **out_buf) {
+size_t pre_save_py(void *py_mutator, u8 *buf, size_t buf_size, u8 *out_buf,
+                   size_t out_buf_size) {
 
-  size_t    out_buf_size;
+  size_t    py_out_buf_size;
   PyObject *py_args, *py_value;
+
+  if (((py_mutator_t *)py_mutator)->scratch_buf) {
+
+    /* We are being recalled from an earlier run
+    where we didn't have enough mem. */
+    if (((py_mutator_t *)py_mutator)->scratch_size < out_buf_size) {
+
+      FATAL("out_buf is still too small after resizing in custom mutator.");
+
+    }
+
+    py_value = ((py_mutator_t *)py_mutator)->scratch_buf;
+    py_out_buf_size = ((py_mutator_t *)py_mutator)->scratch_size;
+    ((py_mutator_t *)py_mutator)->scratch_buf = NULL;
+    py_out_buf_size = 0;
+
+    memcpy(out_buf, PyByteArray_AsString(py_value), py_out_buf_size);
+    Py_DECREF(py_value);
+    return py_out_buf_size;
+
+  }
+
   py_args = PyTuple_New(1);
   py_value = PyByteArray_FromStringAndSize(buf, buf_size);
   if (!py_value) {
@@ -358,11 +389,21 @@ size_t pre_save_py(void *py_mutator, u8 *buf, size_t buf_size, u8 **out_buf) {
 
   if (py_value != NULL) {
 
-    out_buf_size = PyByteArray_Size(py_value);
-    *out_buf = malloc(out_buf_size);
-    memcpy(*out_buf, PyByteArray_AsString(py_value), out_buf_size);
+    py_out_buf_size = PyByteArray_Size(py_value);
+    if (py_out_buf_size > out_buf_size) {
+
+      /* Not enough space!
+      We will get called again right after resizing the buf.
+      Keep the references to our data for now. */
+      ((py_mutator_t *)py_mutator)->scratch_buf = py_value;
+      ((py_mutator_t *)py_mutator)->scratch_size = py_out_buf_size;
+      return py_out_buf_size;
+
+    }
+
+    memcpy(out_buf, PyByteArray_AsString(py_value), py_out_buf_size);
     Py_DECREF(py_value);
-    return out_buf_size;
+    return py_out_buf_size;
 
   } else {
 
diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index e8ef3049..a539b469 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -216,14 +216,49 @@ void write_to_testcase(afl_state_t *afl, void *mem, u32 len) {
 
   if (afl->mutator && afl->mutator->afl_custom_pre_save) {
 
-    u8 *   new_data;
-    size_t new_size = afl->mutator->afl_custom_pre_save(afl->mutator->data, mem,
-                                                        len, &new_data);
-    ck_write(fd, new_data, new_size, afl->fsrv.out_file);
-    free(new_data);
+    if (unlikely(afl->mutator->pre_save_size < len)) {
+
+      afl->mutator->pre_save_buf =
+          ck_realloc(afl->mutator->pre_save_buf, len * sizeof(u8));
+      afl->mutator->pre_save_size = len;
+
+    }
+
+    u8 buf_written = 0;
+    while (!buf_written) {
+
+      buf_written = 1;
+      size_t new_size = afl->mutator->afl_custom_pre_save(
+          afl->mutator->data, mem, len, afl->mutator->pre_save_buf,
+          afl->mutator->pre_save_size);
+
+      if (unlikely(new_size) == 0) {
+
+        /* custom_pre_save wants us to use the old buf */
+        ck_write(fd, mem, len, afl->fsrv.out_file);
+
+      } else if (unlikely(new_size) > afl->mutator->pre_save_size) {
+
+        /* The custom func needs more space.
+           Realloc and call again. */
+        afl->mutator->pre_save_buf =
+            ck_realloc(afl->mutator->pre_save_buf, new_size * sizeof(u8));
+        afl->mutator->pre_save_size = new_size;
+        buf_written = 0;
+        continue;
+
+      } else {
+
+        /* everything as planned. use the new data. */
+        ck_write(fd, afl->mutator->pre_save_buf, new_size, afl->fsrv.out_file);
+
+      }
+
+    }
 
   } else {
 
+    /* boring uncustom. */
     ck_write(fd, mem, len, afl->fsrv.out_file);
 
   }
@@ -505,8 +540,8 @@ void sync_fuzzers(afl_state_t *afl) {
     afl->stage_cur = 0;
     afl->stage_max = 0;
 
-    /* For every file queued by this fuzzer, parse ID and see if we have looked
-       at it before; exec a test case if not. */
+    /* For every file queued by this fuzzer, parse ID and see if we have
+       looked at it before; exec a test case if not. */
 
     while ((qd_ent = readdir(qd))) {
 
@@ -645,7 +680,8 @@ u8 trim_case(afl_state_t *afl, struct queue_entry *q, u8 *in_buf) {
 
       if (afl->stop_soon || fault == FAULT_ERROR) goto abort_trimming;
 
-      /* Note that we don't keep track of crashes or hangs here; maybe TODO? */
+      /* Note that we don't keep track of crashes or hangs here; maybe TODO?
+       */
 
       cksum = hash32(afl->fsrv.trace_bits, MAP_SIZE, HASH_CONST);