Finish refactoring APIs for the custom mutator and Python module

- Remove AFL_PYTHON_ONLY (env) and python_only (variable)
- Unify fuzz API of the custom mutator and Python module
- Merge the custom mutator into the old python_stage, which is now renamed to custom_mutator_stage

6 files changed, 66 insertions, 191 deletions

diff --git a/src/afl-fuzz-globals.c b/src/afl-fuzz-globals.c
index a054499f..87418753 100644
--- a/src/afl-fuzz-globals.c
+++ b/src/afl-fuzz-globals.c
@@ -88,8 +88,7 @@ u8 cal_cycles = CAL_CYCLES,             /* Calibration cycles defaults      */
     no_unlink,                          /* do not unlink cur_input          */
     use_stdin = 1,                      /* use stdin for sending data       */
     be_quiet,                           /* is AFL_QUIET set?                */
-    custom_only,                        /* Custom mutator only mode         */
-    python_only;                        /* Python-only mode                 */
+    custom_only;                        /* Custom mutator only mode         */
 
 u32 stats_update_freq = 1;              /* Stats update frequency (execs)   */
 
diff --git a/src/afl-fuzz-mutators.c b/src/afl-fuzz-mutators.c
index 9365d487..26eaea59 100644
--- a/src/afl-fuzz-mutators.c
+++ b/src/afl-fuzz-mutators.c
@@ -27,6 +27,7 @@
 
 void setup_custom_mutator(void) {
 
+  /* Try mutator library first */
   u8* fn = getenv("AFL_CUSTOM_MUTATOR_LIBRARY");
 
   if (fn) {
@@ -41,6 +42,7 @@ void setup_custom_mutator(void) {
     return;
   }
 
+  /* Try Python module */
 #ifdef USE_PYTHON
   u8* module_name = getenv("AFL_PYTHON_MODULE");
 
@@ -286,7 +288,7 @@ void load_custom_mutator_py(const char* module_name) {
 
   /* "afl_custom_fuzz" should not be NULL, but the interface of Python mutator
      is quite different from the custom mutator. */
-  mutator->afl_custom_fuzz = NULL;
+  mutator->afl_custom_fuzz = fuzz_py;
 
   if (py_functions[PY_FUNC_PRE_SAVE])
     mutator->afl_custom_pre_save = pre_save_py;
diff --git a/src/afl-fuzz-one.c b/src/afl-fuzz-one.c
index 1e6dd45d..5d00e8df 100644
--- a/src/afl-fuzz-one.c
+++ b/src/afl-fuzz-one.c
@@ -482,56 +482,6 @@ u8 fuzz_one_original(char** argv) {
 
   if (use_radamsa > 1) goto radamsa_stage;
 
-  // custom_stage:	// not used - yet
-
-  if (mutator->afl_custom_fuzz) {
-
-    stage_short = "custom";
-    stage_name = "custom mutator";
-    stage_max = len << 3;
-    stage_val_type = STAGE_VAL_NONE;
-
-    const u32 max_seed_size = 4096 * 4096;
-    u8*       mutated_buf = ck_alloc(max_seed_size);
-
-    orig_hit_cnt = queued_paths + unique_crashes;
-
-    for (stage_cur = 0; stage_cur < stage_max; ++stage_cur) {
-
-      size_t orig_size = (size_t)len;
-      size_t mutated_size = mutator->afl_custom_fuzz(in_buf, orig_size,
-                                                     mutated_buf, max_seed_size,
-                                                     UR(UINT32_MAX));
-      if (mutated_size > 0) {
-
-        out_buf = ck_realloc(out_buf, mutated_size);
-        memcpy(out_buf, mutated_buf, mutated_size);
-        if (common_fuzz_stuff(argv, out_buf, (u32)mutated_size)) {
-
-          goto abandon_entry;
-
-        }
-
-      }
-
-    }
-
-    ck_free(mutated_buf);
-    new_hit_cnt = queued_paths + unique_crashes;
-
-    stage_finds[STAGE_CUSTOM_MUTATOR] += new_hit_cnt - orig_hit_cnt;
-    stage_cycles[STAGE_CUSTOM_MUTATOR] += stage_max;
-
-    if (custom_only) {
-
-      /* Skip other stages */
-      ret_val = 0;
-      goto abandon_entry;
-
-    }
-
-  }
-
   if (cmplog_mode) {
 
     if (input_to_state_stage(argv, in_buf, out_buf, len, queue_cur->exec_cksum))
@@ -551,11 +501,7 @@ u8 fuzz_one_original(char** argv) {
                          : havoc_max_mult * 100)) ||
       queue_cur->passed_det) {
 
-#ifdef USE_PYTHON
-    goto python_stage;
-#else
-    goto havoc_stage;
-#endif
+    goto custom_mutator_stage;
 
   }
 
@@ -564,11 +510,7 @@ u8 fuzz_one_original(char** argv) {
 
   if (master_max && (queue_cur->exec_cksum % master_max) != master_id - 1) {
 
-#ifdef USE_PYTHON
-    goto python_stage;
-#else
-    goto havoc_stage;
-#endif
+    goto custom_mutator_stage;
 
   }
 
@@ -1583,24 +1525,25 @@ skip_extras:
 
   if (!queue_cur->passed_det) mark_as_det_done(queue_cur);
 
-#ifdef USE_PYTHON
-python_stage:
-  /**********************************
-   * EXTERNAL MUTATORS (Python API) *
-   **********************************/
+custom_mutator_stage:
+  /*******************
+   * CUSTOM MUTATORS *
+   *******************/
 
-  if (!py_module) goto havoc_stage;
+  if (!mutator) goto havoc_stage;
+  if (!mutator->afl_custom_fuzz) goto havoc_stage;
 
-  stage_name = "python";
-  stage_short = "python";
+  stage_name = "custom mutator";
+  stage_short = "custom";
   stage_max = HAVOC_CYCLES * perf_score / havoc_div / 100;
+  stage_val_type = STAGE_VAL_NONE;
 
   if (stage_max < HAVOC_MIN) stage_max = HAVOC_MIN;
 
-  orig_hit_cnt = queued_paths + unique_crashes;
+  const u32 max_seed_size = 4096 * 4096;
+  u8*       mutated_buf = ck_alloc(max_seed_size);
 
-  char*  retbuf = NULL;
-  size_t retlen = 0;
+  orig_hit_cnt = queued_paths + unique_crashes;
 
   for (stage_cur = 0; stage_cur < stage_max; ++stage_cur) {
 
@@ -1647,26 +1590,24 @@ python_stage:
     ck_read(fd, new_buf, target->len, target->fname);
     close(fd);
 
-    fuzz_py_original(out_buf, len, new_buf, target->len, &retbuf, &retlen);
+    size_t mutated_size = mutator->afl_custom_fuzz(out_buf, len,
+                                                   new_buf, target->len,
+                                                   mutated_buf, max_seed_size);
 
     ck_free(new_buf);
 
-    if (retbuf) {
+    if (mutated_size > 0) {
 
-      if (!retlen) goto abandon_entry;
+      out_buf = ck_realloc(out_buf, mutated_size);
+      memcpy(out_buf, mutated_buf, mutated_size);
 
-      if (common_fuzz_stuff(argv, retbuf, retlen)) {
+      if (common_fuzz_stuff(argv, out_buf, (u32)mutated_size)) {
 
-        free(retbuf);
+        ck_free(mutated_buf);
         goto abandon_entry;
 
       }
 
-      /* Reset retbuf/retlen */
-      free(retbuf);
-      retbuf = NULL;
-      retlen = 0;
-
       /* If we're finding new stuff, let's run for a bit longer, limits
          permitting. */
 
@@ -1687,12 +1628,13 @@ python_stage:
 
   }
 
+  ck_free(mutated_buf);
   new_hit_cnt = queued_paths + unique_crashes;
 
-  stage_finds[STAGE_PYTHON] += new_hit_cnt - orig_hit_cnt;
-  stage_cycles[STAGE_PYTHON] += stage_max;
+  stage_finds[STAGE_CUSTOM_MUTATOR] += new_hit_cnt - orig_hit_cnt;
+  stage_cycles[STAGE_CUSTOM_MUTATOR] += stage_max;
 
-  if (python_only) {
+  if (custom_only) {
 
     /* Skip other stages */
     ret_val = 0;
@@ -1700,8 +1642,6 @@ python_stage:
 
   }
 
-#endif
-
   /****************
    * RANDOM HAVOC *
    ****************/
@@ -2270,11 +2210,10 @@ retry_splicing:
     out_buf = ck_alloc_nozero(len);
     memcpy(out_buf, in_buf, len);
 
-#ifdef USE_PYTHON
-    goto python_stage;
-#else
-    goto havoc_stage;
-#endif
+    goto custom_mutator_stage;
+    /* ???: While integrating Python module, the author decided to jump to
+       python stage, but the reason behind this is not clear.*/
+    // goto havoc_stage;
 
   }
 
diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
index c8caa4c1..c22e4402 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -159,67 +159,16 @@ void init_py(unsigned int seed) {
   }
 }
 
-void fuzz_py_original(char* buf, size_t buflen,
-                      char* add_buf, size_t add_buflen,
-                      char** ret, size_t* retlen) {
+size_t fuzz_py(u8* buf, size_t buf_size,
+               u8* add_buf, size_t add_buf_size,
+               u8* mutated_out, size_t max_size) {
 
-  if (py_module != NULL) {
-
-    PyObject *py_args, *py_value;
-    py_args = PyTuple_New(2);
-    py_value = PyByteArray_FromStringAndSize(buf, buflen);
-    if (!py_value) {
-
-      Py_DECREF(py_args);
-      fprintf(stderr, "Cannot convert argument\n");
-      return;
-
-    }
-
-    PyTuple_SetItem(py_args, 0, py_value);
-
-    py_value = PyByteArray_FromStringAndSize(add_buf, add_buflen);
-    if (!py_value) {
-
-      Py_DECREF(py_args);
-      fprintf(stderr, "Cannot convert argument\n");
-      return;
-
-    }
-
-    PyTuple_SetItem(py_args, 1, py_value);
-
-    py_value = PyObject_CallObject(py_functions[PY_FUNC_FUZZ], py_args);
-
-    Py_DECREF(py_args);
-
-    if (py_value != NULL) {
-
-      *retlen = PyByteArray_Size(py_value);
-      *ret = malloc(*retlen);
-      memcpy(*ret, PyByteArray_AsString(py_value), *retlen);
-      Py_DECREF(py_value);
-
-    } else {
-
-      PyErr_Print();
-      fprintf(stderr, "Call failed\n");
-      return;
-
-    }
-
-  }
-
-}
-
-size_t fuzz_py(u8* data, size_t size, u8* mutated_out, size_t max_size,
-               unsigned int seed) {
-
-  size_t out_size;
+  size_t mutated_size;
   PyObject *py_args, *py_value;
   py_args = PyTuple_New(3);
 
-  py_value = PyByteArray_FromStringAndSize(data, size);
+  /* buf */
+  py_value = PyByteArray_FromStringAndSize(buf, buf_size);
   if (!py_value) {
 
     Py_DECREF(py_args);
@@ -229,11 +178,8 @@ size_t fuzz_py(u8* data, size_t size, u8* mutated_out, size_t max_size,
 
   PyTuple_SetItem(py_args, 0, py_value);
 
-#if PY_MAJOR_VERSION >= 3
-  py_value = PyLong_FromLong(max_size);
-#else
-  py_value = PyInt_FromLong(max_size);
-#endif
+  /* add_buf */
+  py_value = PyByteArray_FromStringAndSize(add_buf, add_buf_size);
   if (!py_value) {
 
     Py_DECREF(py_args);
@@ -243,10 +189,11 @@ size_t fuzz_py(u8* data, size_t size, u8* mutated_out, size_t max_size,
 
   PyTuple_SetItem(py_args, 1, py_value);
 
+  /* max_size */
 #if PY_MAJOR_VERSION >= 3
-  py_value = PyLong_FromLong(seed);
+  py_value = PyLong_FromLong(max_size);
 #else
-  py_value = PyInt_FromLong(seed);
+  py_value = PyInt_FromLong(max_size);
 #endif
   if (!py_value) {
 
@@ -263,11 +210,10 @@ size_t fuzz_py(u8* data, size_t size, u8* mutated_out, size_t max_size,
 
   if (py_value != NULL) {
 
-    out_size = PyByteArray_Size(py_value);
-    memcpy(mutated_out, PyByteArray_AsString(py_value), out_size);
+    mutated_size = PyByteArray_Size(py_value);
+    memcpy(mutated_out, PyByteArray_AsString(py_value), mutated_size);
     Py_DECREF(py_value);
-
-    return out_size;
+    return mutated_size;
 
   } else {
 
@@ -278,12 +224,12 @@ size_t fuzz_py(u8* data, size_t size, u8* mutated_out, size_t max_size,
 
 }
 
-size_t pre_save_py(u8* data, size_t size, u8** new_data) {
+size_t pre_save_py(u8* buf, size_t buf_size, u8** out_buf) {
 
-  size_t new_size;
+  size_t out_buf_size;
   PyObject *py_args, *py_value;
   py_args = PyTuple_New(2);
-  py_value = PyByteArray_FromStringAndSize(data, size);
+  py_value = PyByteArray_FromStringAndSize(buf, buf_size);
   if (!py_value) {
 
     Py_DECREF(py_args);
@@ -299,11 +245,11 @@ size_t pre_save_py(u8* data, size_t size, u8** new_data) {
 
   if (py_value != NULL) {
 
-    new_size = PyByteArray_Size(py_value);
-    *new_data = malloc(new_size);
-    memcpy(*new_data, PyByteArray_AsString(py_value), new_size);
+    out_buf_size = PyByteArray_Size(py_value);
+    *out_buf = malloc(out_buf_size);
+    memcpy(*out_buf, PyByteArray_AsString(py_value), out_buf_size);
     Py_DECREF(py_value);
-    return new_size;
+    return out_buf_size;
 
   } else {
 
@@ -314,12 +260,12 @@ size_t pre_save_py(u8* data, size_t size, u8** new_data) {
 
 }
 
-u32 init_trim_py(u8* buf, size_t buflen) {
+u32 init_trim_py(u8* buf, size_t buf_size) {
 
   PyObject *py_args, *py_value;
 
   py_args = PyTuple_New(1);
-  py_value = PyByteArray_FromStringAndSize(buf, buflen);
+  py_value = PyByteArray_FromStringAndSize(buf, buf_size);
   if (!py_value) {
 
     Py_DECREF(py_args);
@@ -389,7 +335,7 @@ u32 post_trim_py(u8 success) {
 
 }
 
-void trim_py(u8** ret, size_t* retlen) {
+void trim_py(u8** out_buf, size_t* out_buf_size) {
 
   PyObject *py_args, *py_value;
 
@@ -399,9 +345,9 @@ void trim_py(u8** ret, size_t* retlen) {
 
   if (py_value != NULL) {
 
-    *retlen = PyByteArray_Size(py_value);
-    *ret = malloc(*retlen);
-    memcpy(*ret, PyByteArray_AsString(py_value), *retlen);
+    *out_buf_size = PyByteArray_Size(py_value);
+    *out_buf = malloc(*out_buf_size);
+    memcpy(*out_buf, PyByteArray_AsString(py_value), *out_buf_size);
     Py_DECREF(py_value);
 
   } else {
diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index a1d10387..12352355 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -309,11 +309,12 @@ void write_to_testcase(void* mem, u32 len) {
 
     lseek(fd, 0, SEEK_SET);
 
-  if (mutator->afl_custom_pre_save) {
+  if (mutator && mutator->afl_custom_pre_save) {
 
     u8*    new_data;
     size_t new_size = mutator->afl_custom_pre_save(mem, len, &new_data);
     ck_write(fd, new_data, new_size, out_file);
+    ck_free(new_data);
 
   } else {
 
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index d329a20e..a96ee1d0 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -158,7 +158,6 @@ static void usage(u8* argv0, int more_help) {
       "AFL_CUSTOM_MUTATOR_LIBRARY: lib with afl_custom_fuzz() to mutate inputs\n"
       "AFL_CUSTOM_MUTATOR_ONLY: avoid AFL++'s internal mutators\n"
       "AFL_PYTHON_MODULE: mutate and trim inputs with the specified Python module\n"
-      "AFL_PYTHON_ONLY: skip AFL++'s own mutators\n"
       "AFL_DEBUG: extra debugging output for Python mode trimming\n"
       "AFL_DISABLE_TRIM: disable the trimming of test cases\n"
       "AFL_NO_UI: switch status screen off\n"
@@ -658,11 +657,10 @@ int main(int argc, char** argv, char** envp) {
   OKF("afl-tmin fork server patch from github.com/nccgroup/TriforceAFL");
   OKF("MOpt Mutator from github.com/puppet-meteor/MOpt-AFL");
 
-  if (sync_id && force_deterministic &&
-      (getenv("AFL_CUSTOM_MUTATOR_ONLY") || getenv("AFL_PYTHON_ONLY")))
+  if (sync_id && force_deterministic && getenv("AFL_CUSTOM_MUTATOR_ONLY"))
     WARNF(
-        "Using -M master with the AFL_..._ONLY mutator options will result in "
-        "no deterministic mutations being done!");
+        "Using -M master with the AFL_CUSTOM_MUTATOR_ONLY mutator options will "
+        "result in no deterministic mutations being done!");
 
   check_environment_vars(envp);
 
@@ -832,16 +830,6 @@ int main(int argc, char** argv, char** envp) {
 
   if (get_afl_env("AFL_DEBUG")) debug = 1;
 
-  if (get_afl_env("AFL_PYTHON_ONLY")) {
-
-    /* This ensures we don't proceed to havoc/splice */
-    python_only = 1;
-
-    /* Ensure we also skip all deterministic steps */
-    skip_deterministic = 1;
-
-  }
-
   if (get_afl_env("AFL_CUSTOM_MUTATOR_ONLY")) {
 
     /* This ensures we don't proceed to havoc/splice */