Use buffer protocol to retrieve result from python post_process (#605)

Saves an extra copy, gives post processing functions more flexibility

3 files changed, 28 insertions, 12 deletions

diff --git a/docs/custom_mutators.md b/docs/custom_mutators.md
index 2516e511..53f783fe 100644
--- a/docs/custom_mutators.md
+++ b/docs/custom_mutators.md
@@ -130,6 +130,9 @@ def introspection():
     `post_process` function. This function is then transforming the data into the
     format expected by the API before executing the target.
 
+    This can return any python object that implements the buffer protocol and
+    supports PyBUF_SIMPLE. These include bytes, bytearray, etc.
+
 - `queue_new_entry` (optional):
 
     This methods is called after adding a new test case to the queue.
diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index 423230f1..933af65d 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -326,8 +326,7 @@ typedef struct py_mutator {
   u8 *   fuzz_buf;
   size_t fuzz_size;
 
-  u8 *   post_process_buf;
-  size_t post_process_size;
+  Py_buffer post_process_buf;
 
   u8 *   trim_buf;
   size_t trim_size;
diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
index 80532774..9ac4403b 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -134,6 +134,18 @@ static py_mutator_t *init_py_module(afl_state_t *afl, u8 *module_name) {
   PyObject * py_module = py->py_module;
   PyObject **py_functions = py->py_functions;
 
+  // initialize the post process buffer; ensures it's always valid
+  PyObject *unused_bytes = PyByteArray_FromStringAndSize("OHAI", 4);
+  if (!unused_bytes) { FATAL("allocation failed!"); }
+  if (PyObject_GetBuffer(unused_bytes, &py->post_process_buf, PyBUF_SIMPLE) ==
+      -1) {
+
+    FATAL("buffer initialization failed");
+
+  }
+
+  Py_DECREF(unused_bytes);
+
   if (py_module != NULL) {
 
     u8 py_notrim = 0, py_idx;
@@ -313,7 +325,6 @@ struct custom_mutator *load_custom_mutator_py(afl_state_t *afl,
   struct custom_mutator *mutator;
 
   mutator = ck_alloc(sizeof(struct custom_mutator));
-  mutator->post_process_buf = NULL;
 
   mutator->name = module_name;
   ACTF("Loading Python mutator library from '%s'...", module_name);
@@ -403,10 +414,13 @@ struct custom_mutator *load_custom_mutator_py(afl_state_t *afl,
 size_t post_process_py(void *py_mutator, u8 *buf, size_t buf_size,
                        u8 **out_buf) {
 
-  size_t        py_out_buf_size;
   PyObject *    py_args, *py_value;
   py_mutator_t *py = (py_mutator_t *)py_mutator;
 
+  // buffer returned previously must be released; initialized during init
+  // so we don't need to do comparisons
+  PyBuffer_Release(&py->post_process_buf);
+
   py_args = PyTuple_New(1);
   py_value = PyByteArray_FromStringAndSize(buf, buf_size);
   if (!py_value) {
@@ -426,20 +440,20 @@ size_t post_process_py(void *py_mutator, u8 *buf, size_t buf_size,
 
   if (py_value != NULL) {
 
-    py_out_buf_size = PyByteArray_Size(py_value);
+    if (PyObject_GetBuffer(py_value, &py->post_process_buf, PyBUF_SIMPLE) ==
+        -1) {
 
-    if (unlikely(!afl_realloc(BUF_PARAMS(post_process), py_out_buf_size))) {
-
-      PFATAL("alloc");
+      PyErr_Print();
+      FATAL(
+          "Python custom mutator: post_process call return value not a "
+          "bytes-like object");
 
     }
 
-    memcpy(py->post_process_buf, PyByteArray_AsString(py_value),
-           py_out_buf_size);
     Py_DECREF(py_value);
 
-    *out_buf = py->post_process_buf;
-    return py_out_buf_size;
+    *out_buf = (u8 *)py->post_process_buf.buf;
+    return py->post_process_buf.len;
 
   } else {