edited custom mutator pre_save api

1 files changed, 48 insertions, 7 deletions

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 {