more custom mutator remodelling

7 files changed, 103 insertions, 48 deletions

diff --git a/src/afl-fuzz-cmplog.c b/src/afl-fuzz-cmplog.c
index 6c6f05ac..7c398507 100644
--- a/src/afl-fuzz-cmplog.c
+++ b/src/afl-fuzz-cmplog.c
@@ -527,8 +527,13 @@ u8 common_fuzz_cmplog_stuff(afl_state_t *afl, u8 *out_buf, u32 len) {
 
   if (afl->post_handler) {
 
-    out_buf = afl->post_handler(out_buf, &len);
-    if (!out_buf || !len) return 0;
+    u8 *post_buf = NULL;
+
+    size_t post_len =
+        afl->post_handler(afl->post_data, out_buf, len, &post_buf);
+    if (!post_buf || !post_len) return 0;
+    out_buf = post_buf;
+    len = post_len;
 
   }
 
diff --git a/src/afl-fuzz-globals.c b/src/afl-fuzz-globals.c
index 3e573dd2..1d99e1fa 100644
--- a/src/afl-fuzz-globals.c
+++ b/src/afl-fuzz-globals.c
@@ -349,6 +349,8 @@ void read_afl_environment(afl_state_t *afl, char **envp) {
 
 void afl_state_deinit(afl_state_t *afl) {
 
+  if (afl->post_deinit) afl->post_deinit(afl->post_data);
+
   free(afl->out_buf);
   free(afl->out_scratch_buf);
   free(afl->eff_buf);
diff --git a/src/afl-fuzz-init.c b/src/afl-fuzz-init.c
index a69c3b61..fe2be4d2 100644
--- a/src/afl-fuzz-init.c
+++ b/src/afl-fuzz-init.c
@@ -270,7 +270,9 @@ void setup_post(afl_state_t *afl) {
 
   void *dh;
   u8 *  fn = afl->afl_env.afl_post_library;
+  u8    tbuf[6];
   u32   tlen = 6;
+  strncpy(tbuf, "hello", tlen);
 
   if (!fn) return;
 
@@ -281,10 +283,20 @@ void setup_post(afl_state_t *afl) {
 
   afl->post_handler = dlsym(dh, "afl_postprocess");
   if (!afl->post_handler) FATAL("Symbol 'afl_postprocess' not found.");
+  afl->post_init = dlsym(dh, "afl_postprocess_init");
+  if (!afl->post_init) FATAL("Symbol 'afl_postprocess_init' not found.");
+  afl->post_deinit = dlsym(dh, "afl_postprocess_deinit");
+  if (!afl->post_deinit) FATAL("Symbol 'afl_postprocess_deinit' not found.");
 
   /* Do a quick test. It's better to segfault now than later =) */
 
-  afl->post_handler("hello", &tlen);
+  u8 *post_buf = NULL;
+  afl->post_data = afl->post_init(afl);
+  if (!afl->post_data) FATAL("Could not initialize post handler.");
+
+  size_t post_len = afl->post_handler(afl->post_data, tbuf, tlen, &post_buf);
+  if (!post_len || !post_buf)
+    SAYF("Empty return in test post handler for buf=\"hello\\0\".");
 
   OKF("Postprocessor installed successfully.");
 
diff --git a/src/afl-fuzz-mutators.c b/src/afl-fuzz-mutators.c
index 41bc69c8..0692ebb0 100644
--- a/src/afl-fuzz-mutators.c
+++ b/src/afl-fuzz-mutators.c
@@ -7,6 +7,7 @@
    Now maintained by  Marc Heuse <mh@mh-sec.de>,
                         Heiko Eißfeldt <heiko.eissfeldt@hexco.de> and
                         Andrea Fioraldi <andreafioraldi@gmail.com>
+                        Dominik Maier <mail@dmnk.co>
 
    Copyright 2016, 2017 Google Inc. All rights reserved.
    Copyright 2019-2020 AFLplusplus Project. All rights reserved.
@@ -192,7 +193,8 @@ void load_custom_mutator(afl_state_t *afl, const char *fn) {
 
   /* Initialize the custom mutator */
   if (afl->mutator->afl_custom_init)
-    afl->mutator->data = afl->mutator->afl_custom_init(afl, rand_below(afl, 0xFFFFFFFF));
+    afl->mutator->data =
+        afl->mutator->afl_custom_init(afl, rand_below(afl, 0xFFFFFFFF));
 
 }
 
@@ -218,17 +220,18 @@ u8 trim_case_custom(afl_state_t *afl, struct queue_entry *q, u8 *in_buf) {
 
   while (afl->stage_cur < afl->stage_max) {
 
+    u8 *retbuf = NULL;
+
     sprintf(afl->stage_name_buf, "ptrim %s",
             u_stringify_int(val_buf, trim_exec));
 
     u32 cksum;
 
-    u8 *   retbuf = NULL;
-    size_t retlen = 0;
-
-    afl->mutator->afl_custom_trim(afl->mutator->data, &retbuf, &retlen);
+    size_t retlen = afl->mutator->afl_custom_trim(afl->mutator->data, &retbuf);
 
-    if (retlen > orig_len)
+    if (unlikely(retlen < 0 || !retbuf))
+      FATAL("custom_trim failed (ret %zd)", retlen);
+    else if (unlikely(retlen > orig_len))
       FATAL(
           "Trimmed data returned by custom mutator is larger than original "
           "data");
@@ -238,12 +241,7 @@ u8 trim_case_custom(afl_state_t *afl, struct queue_entry *q, u8 *in_buf) {
     fault = run_target(afl, afl->fsrv.exec_tmout);
     ++afl->trim_execs;
 
-    if (afl->stop_soon || fault == FAULT_ERROR) {
-
-      ck_free(retbuf);
-      goto abort_trimming;
-
-    }
+    if (afl->stop_soon || fault == FAULT_ERROR) { goto abort_trimming; }
 
     cksum = hash32(afl->fsrv.trace_bits, MAP_SIZE, HASH_CONST);
 
@@ -281,8 +279,6 @@ u8 trim_case_custom(afl_state_t *afl, struct queue_entry *q, u8 *in_buf) {
 
     }
 
-    ck_free(retbuf);
-
     /* Since this can be slow, update the screen every now and then. */
 
     if (!(trim_exec++ % afl->stats_update_freq)) show_stats(afl);
diff --git a/src/afl-fuzz-one.c b/src/afl-fuzz-one.c
index 024b4665..8dfafb7b 100644
--- a/src/afl-fuzz-one.c
+++ b/src/afl-fuzz-one.c
@@ -1612,15 +1612,20 @@ custom_mutator_stage:
     ck_read(fd, new_buf, target->len, target->fname);
     close(fd);
 
-    // TODO: clean up this mess.
+    u8 *mutated_buf = NULL;
+
     size_t mutated_size = afl->mutator->afl_custom_fuzz(
-        afl->mutator->data, &out_buf, len, new_buf, target->len, max_seed_size);
+        afl->mutator->data, out_buf, len, &mutated_buf, new_buf, target->len,
+        max_seed_size);
+
+    if (unlikely(mutated_size < 0))
+      FATAL("custom_fuzz returned %zd", mutated_size);
 
     if (mutated_size > len) afl->out_size = mutated_size;
 
     if (mutated_size > 0) {
 
-      if (common_fuzz_stuff(afl, out_buf, (u32)mutated_size)) {
+      if (common_fuzz_stuff(afl, mutated_buf, (u32)mutated_size)) {
 
         goto abandon_entry;
 
@@ -1726,8 +1731,22 @@ havoc_stage:
 
       if (stacked_custom && rand_below(afl, 100) < stacked_custom_prob) {
 
-        temp_len = afl->mutator->afl_custom_havoc_mutation(
-            afl->mutator->data, &out_buf, temp_len, MAX_FILE);
+        u8 *   custom_havoc_buf = NULL;
+        size_t new_len = afl->mutator->afl_custom_havoc_mutation(
+            afl->mutator->data, out_buf, temp_len, &custom_havoc_buf, MAX_FILE);
+        if (unlikely(new_len < 0))
+          FATAL("Error in custom_havoc (return %zd)", new_len);
+        if (likely(new_len > 0 && custom_havoc_buf)) {
+
+          temp_len = new_len;
+          if (out_buf != custom_havoc_buf) {
+
+            ck_maybe_grow(BUF_PARAMS(out), temp_len);
+            memcpy(out_buf, custom_havoc_buf, temp_len);
+
+          }
+
+        }
 
       }
 
diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
index 6fbdb678..f9f71929 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -35,15 +35,22 @@ static void *unsupported(afl_state_t *afl, unsigned int seed) {
 
 }
 
-size_t fuzz_py(void *py_mutator, u8 **buf, size_t buf_size, u8 *add_buf,
-               size_t add_buf_size, size_t max_size) {
+/* sorry for this makro...
+it just filles in `&py_mutator->something_buf, &py_mutator->something_size`. */
+#define BUF_PARAMS(name)                              \
+  (void **)&((py_mutator_t *)py_mutator)->name##_buf, \
+      &((py_mutator_t *)py_mutator)->name##_size
+
+size_t fuzz_py(void *py_mutator, u8 *buf, size_t buf_size, u8 **out_buf,
+               u8 *add_buf, size_t add_buf_size, size_t max_size) {
 
   size_t    mutated_size;
   PyObject *py_args, *py_value;
   py_args = PyTuple_New(3);
+  py_mutator_t *py = (py_mutator_t *)py_mutator;
 
   /* buf */
-  py_value = PyByteArray_FromStringAndSize(*buf, buf_size);
+  py_value = PyByteArray_FromStringAndSize(buf, buf_size);
   if (!py_value) {
 
     Py_DECREF(py_args);
@@ -79,17 +86,17 @@ size_t fuzz_py(void *py_mutator, u8 **buf, size_t buf_size, u8 *add_buf,
 
   PyTuple_SetItem(py_args, 2, py_value);
 
-  py_value = PyObject_CallObject(
-      ((py_mutator_t *)py_mutator)->py_functions[PY_FUNC_FUZZ], py_args);
+  py_value = PyObject_CallObject(py->py_functions[PY_FUNC_FUZZ], py_args);
 
   Py_DECREF(py_args);
 
   if (py_value != NULL) {
 
     mutated_size = PyByteArray_Size(py_value);
-    if (buf_size < mutated_size) *buf = ck_realloc(*buf, mutated_size);
 
-    memcpy(*buf, PyByteArray_AsString(py_value), mutated_size);
+    *out_buf = ck_maybe_grow(BUF_PARAMS(fuzz), mutated_size);
+
+    memcpy(*out_buf, PyByteArray_AsString(py_value), mutated_size);
     Py_DECREF(py_value);
     return mutated_size;
 
@@ -364,14 +371,8 @@ size_t pre_save_py(void *py_mutator, u8 *buf, size_t buf_size, u8 **out_buf) {
 
     py_out_buf_size = PyByteArray_Size(py_value);
 
-    if (py_out_buf_size > py->pre_save_size) {
-
-      /* Not enough space!
-      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;
-
-    }
+    ck_maybe_grow((void **)&py->pre_save_buf, &py->pre_save_size,
+                  py_out_buf_size);
 
     memcpy(py->pre_save_buf, PyByteArray_AsString(py_value), py_out_buf_size);
     Py_DECREF(py_value);
@@ -465,9 +466,10 @@ u32 post_trim_py(void *py_mutator, u8 success) {
 
 }
 
-void trim_py(void *py_mutator, u8 **out_buf, size_t *out_buf_size) {
+size_t trim_py(void *py_mutator, u8 **out_buf) {
 
   PyObject *py_args, *py_value;
+  size_t    ret;
 
   py_args = PyTuple_New(0);
   py_value = PyObject_CallObject(
@@ -476,9 +478,9 @@ void trim_py(void *py_mutator, u8 **out_buf, size_t *out_buf_size) {
 
   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);
+    ret = PyByteArray_Size(py_value);
+    *out_buf = ck_maybe_grow(BUF_PARAMS(trim), ret);
+    memcpy(*out_buf, PyByteArray_AsString(py_value), ret);
     Py_DECREF(py_value);
 
   } else {
@@ -488,17 +490,19 @@ void trim_py(void *py_mutator, u8 **out_buf, size_t *out_buf_size) {
 
   }
 
+  return ret;
+
 }
 
-size_t havoc_mutation_py(void *py_mutator, u8 **buf, size_t buf_size,
-                         size_t max_size) {
+size_t havoc_mutation_py(void *py_mutator, u8 *buf, size_t buf_size,
+                         u8 **out_buf, size_t max_size) {
 
   size_t    mutated_size;
   PyObject *py_args, *py_value;
   py_args = PyTuple_New(2);
 
   /* buf */
-  py_value = PyByteArray_FromStringAndSize(*buf, buf_size);
+  py_value = PyByteArray_FromStringAndSize(buf, buf_size);
   if (!py_value) {
 
     Py_DECREF(py_args);
@@ -532,9 +536,19 @@ size_t havoc_mutation_py(void *py_mutator, u8 **buf, size_t buf_size,
   if (py_value != NULL) {
 
     mutated_size = PyByteArray_Size(py_value);
-    if (buf_size < mutated_size) *buf = ck_realloc(*buf, mutated_size);
+    if (mutated_size <= buf_size) {
+
+      /* We reuse the input buf here. */
+      *out_buf = buf;
+
+    } else {
 
-    memcpy(*buf, PyByteArray_AsString(py_value), mutated_size);
+      /* A new buf is needed... */
+      *out_buf = ck_maybe_grow(BUF_PARAMS(havoc), mutated_size);
+
+    }
+
+    memcpy(*out_buf, PyByteArray_AsString(py_value), mutated_size);
 
     Py_DECREF(py_value);
     return mutated_size;
@@ -680,5 +694,7 @@ void queue_new_entry_py(void *py_mutator, const u8 *filename_new_queue,
 
 }
 
+#undef BUF_PARAMS
+
 #endif                                                        /* USE_PYTHON */
 
diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index a43bfad2..9ba2c5f7 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -741,8 +741,13 @@ u8 common_fuzz_stuff(afl_state_t *afl, u8 *out_buf, u32 len) {
 
   if (afl->post_handler) {
 
-    out_buf = afl->post_handler(out_buf, &len);
-    if (!out_buf || !len) return 0;
+    u8 *post_buf = NULL;
+
+    size_t post_len =
+        afl->post_handler(afl->post_data, out_buf, len, &post_buf);
+    if (!post_buf || !post_len) return 0;
+    out_buf = post_buf;
+    len = post_len;
 
   }