fix attempt at post_process implementation

5 files changed, 30 insertions, 45 deletions

diff --git a/docs/Changelog.md b/docs/Changelog.md
index 501300b1..9436fc9f 100644
--- a/docs/Changelog.md
+++ b/docs/Changelog.md
@@ -13,6 +13,7 @@
       everyone who was affected!
     - allow pizza mode to be disabled when AFL_PIZZA_MODE is set to -1
     - option `-p mmopt` now also selects new queue items more often
+    - fix bug in post_process custom mutator implementation
     - print name of custom mutator in UI
   - afl-cc:
     - add CFI sanitizer variant to gcc targets
diff --git a/docs/custom_mutators.md b/docs/custom_mutators.md
index 82131c92..a1de479e 100644
--- a/docs/custom_mutators.md
+++ b/docs/custom_mutators.md
@@ -118,7 +118,7 @@ def deinit():  # optional for Python
 
 ### Custom Mutation
 
-- `init`:
+- `init` (optional in Python):
 
     This method is called when AFL++ starts up and is used to seed RNG and set
     up buffers and state.
@@ -184,6 +184,11 @@ def deinit():  # optional for Python
     to the target, e.g. if it is too short, too corrupted, etc. If so,
     return a NULL buffer and zero length (or a 0 length string in Python).
 
+    NOTE: Do not make any random changes to the data in this function!
+
+    PERFORMANCE for C/C++: If possible make the changes in-place (so modify
+    the `*data` directly, and return it as `*outbuf = data`.
+
 - `fuzz_send` (optional):
 
     This method can be used if you want to send data to the target yourself,
@@ -202,7 +207,7 @@ def deinit():  # optional for Python
     discovered if compiled with INTROSPECTION. The custom mutator can then
     return a string (const char *) that reports the exact mutations used.
 
-- `deinit`:
+- `deinit` (optional in Python):
 
     The last method to be called, deinitializing the state.
 
diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index 5fd393dd..8b6502b4 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -885,14 +885,19 @@ struct custom_mutator {
    * A post-processing function to use right before AFL writes the test case to
    * disk in order to execute the target.
    *
-   * (Optional) If this functionality is not needed, simply don't define this
+   * NOTE: Do not do any random changes to the data in this function!
+   *
+   * PERFORMANCE: If you can modify the data in-place you will have a better
+   *              performance. Modify *data and set `*out_buf = data`.
+   *
+   * (Optional) If this functionality is not needed, simply do not define this
    * function.
    *
    * @param[in] data pointer returned in afl_custom_init by this custom mutator
    * @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 storing the test case after
-   *     processing. External library should allocate memory for out_buf.
+   *     processing. The 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.
    */
diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
index 673e5a6c..7dad0770 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -219,11 +219,14 @@ static py_mutator_t *init_py_module(afl_state_t *afl, u8 *module_name) {
 
   if (py_module != NULL) {
 
-    u8 py_notrim = 0, py_idx;
-    /* init, required */
+    u8 py_notrim = 0;
     py_functions[PY_FUNC_INIT] = PyObject_GetAttrString(py_module, "init");
-    if (!py_functions[PY_FUNC_INIT])
-      FATAL("init function not found in python module");
+    if (!py_functions[PY_FUNC_INIT]) {
+
+      WARNF("init function not found in python module");
+
+    }
+
     py_functions[PY_FUNC_FUZZ] = PyObject_GetAttrString(py_module, "fuzz");
     if (!py_functions[PY_FUNC_FUZZ])
       py_functions[PY_FUNC_FUZZ] = PyObject_GetAttrString(py_module, "mutate");
@@ -231,12 +234,6 @@ static py_mutator_t *init_py_module(afl_state_t *afl, u8 *module_name) {
         PyObject_GetAttrString(py_module, "describe");
     py_functions[PY_FUNC_FUZZ_COUNT] =
         PyObject_GetAttrString(py_module, "fuzz_count");
-    if (!py_functions[PY_FUNC_FUZZ]) {
-
-      WARNF("fuzz function not found in python module");
-
-    }
-
     py_functions[PY_FUNC_POST_PROCESS] =
         PyObject_GetAttrString(py_module, "post_process");
     py_functions[PY_FUNC_INIT_TRIM] =
@@ -263,36 +260,6 @@ static py_mutator_t *init_py_module(afl_state_t *afl, u8 *module_name) {
     if (!py_functions[PY_FUNC_DEINIT])
       WARNF("deinit function not found in python module");
 
-    for (py_idx = 0; py_idx < PY_FUNC_COUNT; ++py_idx) {
-
-      if (!py_functions[py_idx] || !PyCallable_Check(py_functions[py_idx])) {
-
-        if (py_idx >= PY_FUNC_INIT_TRIM && py_idx <= PY_FUNC_TRIM) {
-
-          // Implementing the trim API is optional for now
-          if (PyErr_Occurred()) { PyErr_Print(); }
-          py_notrim = 1;
-
-        } else if (py_idx >= PY_OPTIONAL) {
-
-          // Only _init and _deinit are not optional currently
-
-          if (PyErr_Occurred()) { PyErr_Print(); }
-
-        } else {
-
-          fprintf(stderr,
-                  "Cannot find/call function with index %d in external "
-                  "Python module.\n",
-                  py_idx);
-          return NULL;
-
-        }
-
-      }
-
-    }
-
     if (py_notrim) {
 
       py_functions[PY_FUNC_INIT_TRIM] = NULL;
@@ -345,6 +312,8 @@ static void init_py(afl_state_t *afl, py_mutator_t *py_mutator,
 
   (void)afl;
 
+  if (py_mutator->py_functions[PY_FUNC_INIT] == NULL) { return; }
+
   PyObject *py_args, *py_value;
 
   /* Provide the init function a seed for the Python RNG */
diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index f5425011..26e8549d 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -133,7 +133,12 @@ write_to_testcase(afl_state_t *afl, void **mem, u32 len, u32 fix) {
 
     }
 
-    if (new_mem != *mem) { *mem = new_mem; }
+    if (new_mem != *mem && new_mem != NULL && new_size > 0) {
+
+      *mem = afl_realloc((void **)mem, new_size);
+      memmove(*mem, new_mem, new_size);
+
+    }
 
     if (unlikely(afl->custom_mutators_count)) {