more api

7 files changed, 71 insertions, 112 deletions

diff --git a/docs/custom_mutators.md b/docs/custom_mutators.md
index ecbd50ba..2163b2d5 100644
--- a/docs/custom_mutators.md
+++ b/docs/custom_mutators.md
@@ -27,7 +27,7 @@ performed with the custom mutator.
 
 C/C++:
 ```c
-void afl_custom_init(unsigned int seed);
+void afl_custom_init(void *afl, 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_write_to_testcase(uint8_t* buf, size_t buf_size, uint8_t** out_buf);
@@ -39,6 +39,7 @@ uint8_t afl_custom_havoc_mutation_probability(void);
 uint8_t afl_custom_queue_get(const uint8_t* filename);
 void afl_custom_queue_new_entry(const uint8_t* filename_new_queue,
                                 const uint8_t* filename_orig_queue);
+void afl_custom_deinit(void *data);
 ```
 
 Python:
diff --git a/examples/custom_mutators/example.c b/examples/custom_mutators/example.c
index 3e708db8..04b72c77 100644
--- a/examples/custom_mutators/example.c
+++ b/examples/custom_mutators/example.c
@@ -15,6 +15,7 @@
 #include <stdio.h>
 
 #define DATA_SIZE (100)
+#define INITIAL_BUF_SIZE (16384)
 
 static const char *commands[] = {
 
@@ -28,6 +29,8 @@ typedef struct my_mutator {
 
   afl_t *afl;
   // any additional data here!
+  size_t pre_save_size;
+  u8 *   pre_save_buf;
 
 } my_mutator_t;
 
@@ -56,6 +59,16 @@ my_mutator_t *afl_custom_init(afl_t *afl, unsigned int seed) {
 
   data->afl = afl;
 
+  data->pre_save_buf = malloc(INITIAL_BUF_SIZE);
+  if (!data->pre_save_buf) {
+
+    free(data);
+    return NULL;
+
+  }
+
+  data->pre_save_size = INITIAL_BUF_SIZE;
+
   return data;
 
 }
@@ -125,12 +138,23 @@ size_t afl_custom_fuzz(my_mutator_t *data, uint8_t **buf, size_t buf_size,
 size_t afl_custom_pre_save(my_mutator_t *data, uint8_t *buf, size_t buf_size,
                            uint8_t *out_buf, 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;
+  if (data->pre_save_size < buf_size + 5) {
+
+    data->pre_save_buf = realloc(data->pre_save_buf, buf_size + 5);
+    if (!data->pre_save_buf) {
+
+      perror("custom mutator realloc");
+      free(data);
+      return -1;
+
+    }
+
+    data->pre_save_size = buf_size + 5;
+
+  }
 
-  memcpy(out_buf, buf, buf_size);
-  out_buf_size = buf_size;
+  memcpy(out_buf + 5, buf, buf_size);
+  out_buf_size = buf_size + 5;
   out_buf[0] = 'A';
   out_buf[1] = 'F';
   out_buf[2] = 'L';
@@ -322,6 +346,7 @@ void afl_custom_queue_new_entry(my_mutator_t * data,
  */
 void afl_custom_deinit(my_mutator_t *data) {
 
+  free(data->pre_save_buf);
   free(data);
 
 }
diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index 8bf66403..530a4b6a 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -30,9 +30,6 @@
 #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
@@ -296,8 +293,8 @@ typedef struct py_mutator {
   void *    afl_state;
   void *    py_data;
 
-  PyObject *scratch_buf;
-  size_t    scratch_size;
+  u8 *   pre_save_buf;
+  size_t pre_save_size;
 
 } py_mutator_t;
 
@@ -641,13 +638,13 @@ struct custom_mutator {
    * @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 of storing 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
+   * @param[out] out_buf Pointer to the buffer storing the test case after
+   *     processing. External library should allocate memory for out_buf.
+   *     It can chose to alter buf in-place, if the space is large enough.
+   * @return Size of the output buffer.
    */
   size_t (*afl_custom_pre_save)(void *data, u8 *buf, size_t buf_size,
-                                u8 *out_buf, size_t out_buf_size);
+                                u8 **out_buf);
 
   /**
    * This method is called at the start of each trimming operation and receives
@@ -783,7 +780,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);
+size_t pre_save_py(void *, u8 *, size_t, u8 **);
 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 fac97f8e..9d8610c0 100644
--- a/src/afl-fuzz-mutators.c
+++ b/src/afl-fuzz-mutators.c
@@ -135,18 +135,9 @@ 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 b053e8d5..6fbdb678 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -307,14 +307,8 @@ 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;
@@ -344,39 +338,18 @@ 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 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;
+size_t pre_save_py(void *py_mutator, u8 *buf, size_t buf_size, u8 **out_buf) {
 
-    memcpy(out_buf, PyByteArray_AsString(py_value), py_out_buf_size);
-    Py_DECREF(py_value);
-    return py_out_buf_size;
-
-  }
+  size_t        py_out_buf_size;
+  PyObject *    py_args, *py_value;
+  py_mutator_t *py = (py_mutator_t *)py_mutator;
 
   py_args = PyTuple_New(1);
   py_value = PyByteArray_FromStringAndSize(buf, buf_size);
   if (!py_value) {
 
     Py_DECREF(py_args);
-    FATAL("Failed to convert arguments");
+    FATAL("Failed to convert arguments in custom pre_save");
 
   }
 
@@ -390,25 +363,26 @@ size_t pre_save_py(void *py_mutator, u8 *buf, size_t buf_size, u8 *out_buf,
   if (py_value != NULL) {
 
     py_out_buf_size = PyByteArray_Size(py_value);
-    if (py_out_buf_size > out_buf_size) {
+
+    if (py_out_buf_size > py->pre_save_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;
+      Let's resize our buf */
+      py->pre_save_buf = ck_realloc(py->pre_save_buf, py_out_buf_size);
+      py->pre_save_size = py_out_buf_size;
 
     }
 
-    memcpy(out_buf, PyByteArray_AsString(py_value), py_out_buf_size);
+    memcpy(py->pre_save_buf, PyByteArray_AsString(py_value), py_out_buf_size);
     Py_DECREF(py_value);
+
+    *out_buf = py->pre_save_buf;
     return py_out_buf_size;
 
   } else {
 
     PyErr_Print();
-    FATAL("Call failed");
+    FATAL("Python custom mutator: pre_save call failed.");
 
   }
 
diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index a539b469..a43bfad2 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -214,47 +214,18 @@ void write_to_testcase(afl_state_t *afl, void *mem, u32 len) {
 
     lseek(fd, 0, SEEK_SET);
 
-  if (afl->mutator && afl->mutator->afl_custom_pre_save) {
+  if (unlikely(afl->mutator && afl->mutator->afl_custom_pre_save)) {
 
-    if (unlikely(afl->mutator->pre_save_size < len)) {
+    u8 *new_buf = NULL;
 
-      afl->mutator->pre_save_buf =
-          ck_realloc(afl->mutator->pre_save_buf, len * sizeof(u8));
-      afl->mutator->pre_save_size = len;
+    size_t new_size = afl->mutator->afl_custom_pre_save(afl->mutator->data, mem,
+                                                        len, &new_buf);
 
-    }
-
-    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);
+    if (unlikely(new_size <= 0 || !new_buf))
+      FATAL("Custom_pre_save failed (ret: %ld)", new_size);
 
-      }
-
-    }
+    /* everything as planned. use the new data. */
+    ck_write(fd, new_buf, new_size, afl->fsrv.out_file);
 
   } else {
 
diff --git a/src/afl-fuzz-stats.c b/src/afl-fuzz-stats.c
index 53c162a2..5536c201 100644
--- a/src/afl-fuzz-stats.c
+++ b/src/afl-fuzz-stats.c
@@ -361,9 +361,9 @@ void show_stats(afl_state_t *afl) {
 
   /* Lord, forgive me this. */
 
-  SAYF(SET_G1 bSTG bLT bH bSTOP cCYA
+  SAYF(SET_G1 bSTG bLT bH bSTOP                         cCYA
        " process timing " bSTG bH30 bH5 bH bHB bH bSTOP cCYA
-       " overall results " bSTG bH2 bH2 bRT "\n");
+       " overall results " bSTG bH2 bH2                 bRT "\n");
 
   if (afl->dumb_mode) {
 
@@ -445,9 +445,9 @@ void show_stats(afl_state_t *afl) {
                 "   uniq hangs : " cRST "%-6s" bSTG         bV "\n",
        time_tmp, tmp);
 
-  SAYF(bVR bH bSTOP            cCYA
+  SAYF(bVR bH bSTOP                                          cCYA
        " cycle progress " bSTG bH10 bH5 bH2 bH2 bHB bH bSTOP cCYA
-       " map coverage " bSTG bH bHT bH20 bH2 bVL "\n");
+       " map coverage " bSTG bH bHT bH20 bH2                 bVL "\n");
 
   /* This gets funny because we want to print several variable-length variables
      together, but then cram them into a fixed-width field - so we need to
@@ -476,9 +476,9 @@ void show_stats(afl_state_t *afl) {
 
   SAYF(bSTOP " count coverage : " cRST "%-21s" bSTG bV "\n", tmp);
 
-  SAYF(bVR bH bSTOP            cCYA
+  SAYF(bVR bH bSTOP                                         cCYA
        " stage progress " bSTG bH10 bH5 bH2 bH2 bX bH bSTOP cCYA
-       " findings in depth " bSTG bH10 bH5 bH2 bH2 bVL "\n");
+       " findings in depth " bSTG bH10 bH5 bH2 bH2          bVL "\n");
 
   sprintf(tmp, "%s (%0.02f%%)", u_stringify_int(IB(0), afl->queued_favored),
           ((double)afl->queued_favored) * 100 / afl->queued_paths);
@@ -552,7 +552,7 @@ void show_stats(afl_state_t *afl) {
 
   /* Aaaalmost there... hold on! */
 
-  SAYF(bVR bH cCYA                      bSTOP
+  SAYF(bVR bH cCYA                                                     bSTOP
        " fuzzing strategy yields " bSTG bH10 bHT bH10 bH5 bHB bH bSTOP cCYA
        " path geometry " bSTG bH5 bH2 bVL "\n");