custom mutators might work again like this

2 files changed, 177 insertions, 143 deletions

diff --git a/src/afl-fuzz-mutators.c b/src/afl-fuzz-mutators.c
index c158aa6f..0ded4ba1 100644
--- a/src/afl-fuzz-mutators.c
+++ b/src/afl-fuzz-mutators.c
@@ -27,7 +27,7 @@
 
 void load_custom_mutator(afl_state_t *, const char *);
 #ifdef USE_PYTHON
-void load_custom_mutator_py(afl_state_t *, const char *);
+void load_custom_mutator_py(afl_state_t *, char *);
 #endif
 
 void setup_custom_mutator(afl_state_t *afl) {
@@ -59,10 +59,7 @@ void setup_custom_mutator(afl_state_t *afl) {
       FATAL(
           "MOpt and Python mutator are mutually exclusive. We accept pull "
           "requests that integrates MOpt with the optional mutators "
-          "(custom/radamsa/redquenn/...).");
-
-    if (init_py_module(afl, module_name))
-      FATAL("Failed to initialize Python module");
+          "(custom/radamsa/redqueen/...).");
 
     load_custom_mutator_py(afl, module_name);
 
@@ -79,18 +76,13 @@ void destroy_custom_mutator(afl_state_t *afl) {
 
   if (afl->mutator) {
 
+    afl->mutator->afl_custom_deinit(afl->mutator->data);
+
     if (afl->mutator->dh)
       dlclose(afl->mutator->dh);
-    else {
-
-      /* Python mutator */
-#ifdef USE_PYTHON
-      finalize_py_module(afl);
-#endif
-
-    }
 
     ck_free(afl->mutator);
+    afl->mutator = NULL;
 
   }
 
@@ -109,10 +101,13 @@ void load_custom_mutator(afl_state_t *afl, const char *fn) {
   afl->mutator->dh = dh;
 
   /* Mutator */
-  /* "afl_custom_init", optional for backward compatibility */
+  /* "afl_custom_init", required */
   afl->mutator->afl_custom_init = dlsym(dh, "afl_custom_init");
-  if (!afl->mutator->afl_custom_init)
-    WARNF("Symbol 'afl_custom_init' not found.");
+  if (!afl->mutator->afl_custom_init) FATAL("Symbol 'afl_custom_init' not found.");
+
+  /* "afl_custom_deinit", required */
+  afl->mutator->afl_custom_deinit = dlsym(dh, "afl_custom_deinit");
+  if (!afl->mutator->afl_custom_deinit) FATAL("Symbol 'afl_custom_deinit' not found.");
 
   /* "afl_custom_fuzz" or "afl_custom_mutator", required */
   afl->mutator->afl_custom_fuzz = dlsym(dh, "afl_custom_fuzz");
@@ -203,7 +198,7 @@ u8 trim_case_custom(afl_state_t *afl, struct queue_entry *q, u8 *in_buf) {
 
   /* Initialize trimming in the custom mutator */
   afl->stage_cur = 0;
-  afl->stage_max = afl->mutator->afl_custom_init_trim(afl, in_buf, q->len);
+  afl->stage_max = afl->mutator->afl_custom_init_trim(afl->mutator->data, in_buf, q->len);
 
   if (afl->not_on_tty && afl->debug)
     SAYF("[Custom Trimming] START: Max %d iterations, %u bytes", afl->stage_max,
@@ -309,54 +304,3 @@ abort_trimming:
   return fault;
 
 }
-
-#ifdef USE_PYTHON
-void load_custom_mutator_py(afl_state_t *afl, const char *module_name) {
-
-  PyObject **py_functions = afl->py_functions;
-
-  afl->mutator = ck_alloc(sizeof(struct custom_mutator));
-
-  afl->mutator->name = module_name;
-  ACTF("Loading Python mutator library from '%s'...", module_name);
-
-  if (py_functions[PY_FUNC_INIT]) afl->mutator->afl_custom_init = init_py;
-
-  /* "afl_custom_fuzz" should not be NULL, but the interface of Python mutator
-     is quite different from the custom mutator. */
-  afl->mutator->afl_custom_fuzz = fuzz_py;
-
-  if (py_functions[PY_FUNC_PRE_SAVE])
-    afl->mutator->afl_custom_pre_save = pre_save_py;
-
-  if (py_functions[PY_FUNC_INIT_TRIM])
-    afl->mutator->afl_custom_init_trim = init_trim_py;
-
-  if (py_functions[PY_FUNC_POST_TRIM])
-    afl->mutator->afl_custom_post_trim = post_trim_py;
-
-  if (py_functions[PY_FUNC_TRIM]) afl->mutator->afl_custom_trim = trim_py;
-
-  if (py_functions[PY_FUNC_HAVOC_MUTATION])
-    afl->mutator->afl_custom_havoc_mutation = havoc_mutation_py;
-
-  if (py_functions[PY_FUNC_HAVOC_MUTATION_PROBABILITY])
-    afl->mutator->afl_custom_havoc_mutation_probability =
-        havoc_mutation_probability_py;
-
-  if (py_functions[PY_FUNC_QUEUE_GET])
-    afl->mutator->afl_custom_queue_get = queue_get_py;
-
-  if (py_functions[PY_FUNC_QUEUE_NEW_ENTRY])
-    afl->mutator->afl_custom_queue_new_entry = queue_new_entry_py;
-
-  OKF("Python mutator '%s' installed successfully.", module_name);
-
-  /* Initialize the custom mutator */
-  if (afl->mutator->afl_custom_init)
-    afl->mutator->afl_custom_init(afl, rand_below(afl, 0xFFFFFFFF));
-
-}
-
-#endif
-
diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
index 595c1ed0..d3027d2b 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -28,9 +28,84 @@
 /* Python stuff */
 #ifdef USE_PYTHON
 
-int init_py_module(afl_state_t *afl, u8 *module_name) {
+static void *unsupported(afl_state_t *afl, unsigned int seed) {
+  FATAL("Python Mutator cannot be called twice yet");
+  return NULL;
+}
+
+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) {
+
+  size_t    mutated_size;
+  PyObject *py_args, *py_value;
+  py_args = PyTuple_New(3);
+
+  /* buf */
+  py_value = PyByteArray_FromStringAndSize(*buf, buf_size);
+  if (!py_value) {
+
+    Py_DECREF(py_args);
+    FATAL("Failed to convert arguments");
+
+  }
+
+  PyTuple_SetItem(py_args, 0, py_value);
+
+  /* add_buf */
+  py_value = PyByteArray_FromStringAndSize(add_buf, add_buf_size);
+  if (!py_value) {
+
+    Py_DECREF(py_args);
+    FATAL("Failed to convert arguments");
+
+  }
+
+  PyTuple_SetItem(py_args, 1, py_value);
+
+  /* max_size */
+#if PY_MAJOR_VERSION >= 3
+  py_value = PyLong_FromLong(max_size);
+#else
+  py_value = PyInt_FromLong(max_size);
+#endif
+  if (!py_value) {
+
+    Py_DECREF(py_args);
+    FATAL("Failed to convert arguments");
+
+  }
+
+  PyTuple_SetItem(py_args, 2, py_value);
+
+  py_value = PyObject_CallObject(((py_mutator_t *)py_mutator)->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);
+    Py_DECREF(py_value);
+    return mutated_size;
+
+  } else {
+
+    PyErr_Print();
+    FATAL("Call failed");
+
+  }
+
+}
 
-  if (!module_name) return 1;
+
+static py_mutator_t *init_py_module(afl_state_t *afl, u8 *module_name) {
+
+  if (!module_name) return NULL;
+
+  py_mutator_t *py = calloc(1, sizeof(py_mutator_t));
+  if (!py) PFATAL("Could not allocate memory for python mutator!");
 
   Py_Initialize();
 
@@ -40,17 +115,18 @@ int init_py_module(afl_state_t *afl, u8 *module_name) {
   PyObject *py_name = PyString_FromString(module_name);
 #endif
 
-  afl->py_module = PyImport_Import(py_name);
+  py->py_module = PyImport_Import(py_name);
   Py_DECREF(py_name);
 
-  PyObject * py_module = afl->py_module;
-  PyObject **py_functions = afl->py_functions;
+  PyObject * py_module = py->py_module;
+  PyObject **py_functions = py->py_functions;
 
-  if (afl->py_module != NULL) {
+  if (py_module != NULL) {
 
     u8 py_notrim = 0, py_idx;
-    py_functions[PY_FUNC_INIT] = PyObject_GetAttrString(afl->py_module, "init");
-    py_functions[PY_FUNC_FUZZ] = PyObject_GetAttrString(afl->py_module, "fuzz");
+    py_functions[PY_FUNC_INIT] = PyObject_GetAttrString(py_module, "init");
+    py_functions[PY_FUNC_DEINIT] = PyObject_GetAttrString(py_module, "deinit");
+    py_functions[PY_FUNC_FUZZ] = PyObject_GetAttrString(py_module, "fuzz");
     py_functions[PY_FUNC_PRE_SAVE] =
         PyObject_GetAttrString(py_module, "pre_save");
     py_functions[PY_FUNC_INIT_TRIM] =
@@ -96,7 +172,7 @@ int init_py_module(afl_state_t *afl, u8 *module_name) {
                   "Cannot find/call function with index %d in external "
                   "Python module.\n",
                   py_idx);
-          return 1;
+          return NULL;
 
         }
 
@@ -119,23 +195,27 @@ int init_py_module(afl_state_t *afl, u8 *module_name) {
 
     PyErr_Print();
     fprintf(stderr, "Failed to load \"%s\"\n", module_name);
-    return 1;
+    return NULL;
 
   }
 
-  return 0;
+  return py;
 
 }
 
-void finalize_py_module(afl_state_t *afl) {
+void finalize_py_module(void *py_mutator) {
+
+  py_mutator_t *py = (py_mutator_t *)py_mutator;
 
-  if (afl->py_module != NULL) {
+  if (py->py_module != NULL) {
+
+    deinit_py(py_mutator);
 
     u32 i;
     for (i = 0; i < PY_FUNC_COUNT; ++i)
-      Py_XDECREF(afl->py_functions[i]);
+      Py_XDECREF(py->py_functions[i]);
 
-    Py_DECREF(afl->py_module);
+    Py_DECREF(py->py_module);
 
   }
 
@@ -143,7 +223,7 @@ void finalize_py_module(afl_state_t *afl) {
 
 }
 
-void init_py(afl_state_t *afl, unsigned int seed) {
+static void init_py(afl_state_t *afl, py_mutator_t *py_mutator, unsigned int seed) {
 
   PyObject *py_args, *py_value;
 
@@ -158,14 +238,13 @@ void init_py(afl_state_t *afl, unsigned int seed) {
   if (!py_value) {
 
     Py_DECREF(py_args);
-    fprintf(stderr, "Cannot convert argument\n");
-    return;
+    FATAL("Cannot convert argument in python init.");
 
   }
 
   PyTuple_SetItem(py_args, 0, py_value);
 
-  py_value = PyObject_CallObject(afl->py_functions[PY_FUNC_INIT], py_args);
+  py_value = PyObject_CallObject(py_mutator->py_functions[PY_FUNC_INIT], py_args);
 
   Py_DECREF(py_args);
 
@@ -173,79 +252,90 @@ void init_py(afl_state_t *afl, unsigned int seed) {
 
     PyErr_Print();
     fprintf(stderr, "Call failed\n");
-    return;
+    FATAL("Custom py mutator INIT failed.");
 
   }
 
 }
 
-size_t fuzz_py(afl_state_t *afl, u8 **buf, size_t buf_size, u8 *add_buf,
-               size_t add_buf_size, size_t max_size) {
+void deinit_py(void *py_mutator) {
 
-  size_t    mutated_size;
   PyObject *py_args, *py_value;
-  py_args = PyTuple_New(3);
 
-  /* buf */
-  py_value = PyByteArray_FromStringAndSize(*buf, buf_size);
-  if (!py_value) {
+  py_args = PyTuple_New(0);
+  py_value = PyObject_CallObject(
+      ((py_mutator_t *)py_mutator)->py_functions[PY_FUNC_DEINIT], py_args);
+  Py_DECREF(py_args);
 
-    Py_DECREF(py_args);
-    FATAL("Failed to convert arguments");
+  if (py_value != NULL) {
+
+    Py_DECREF(py_value);
+
+  } else {
+
+    PyErr_Print();
+    FATAL("Call failed");
 
   }
 
-  PyTuple_SetItem(py_args, 0, py_value);
+}
 
-  /* add_buf */
-  py_value = PyByteArray_FromStringAndSize(add_buf, add_buf_size);
-  if (!py_value) {
+void load_custom_mutator_py(afl_state_t *afl, char *module_name) {
 
-    Py_DECREF(py_args);
-    FATAL("Failed to convert arguments");
+  afl->mutator = ck_alloc(sizeof(struct custom_mutator));
 
+  afl->mutator->name = module_name;
+  ACTF("Loading Python mutator library from '%s'...", module_name);
+
+  py_mutator_t *py_mutator;
+  py_mutator = init_py_module(afl, module_name);
+  if (!py_mutator) {
+    FATAL("Failed to load python mutator.");
   }
 
-  PyTuple_SetItem(py_args, 1, py_value);
+  PyObject **py_functions = py_mutator->py_functions;
 
-  /* max_size */
-#if PY_MAJOR_VERSION >= 3
-  py_value = PyLong_FromLong(max_size);
-#else
-  py_value = PyInt_FromLong(max_size);
-#endif
-  if (!py_value) {
+  if (py_functions[PY_FUNC_INIT]) afl->mutator->afl_custom_init = unsupported;
 
-    Py_DECREF(py_args);
-    FATAL("Failed to convert arguments");
+  if (py_functions[PY_FUNC_DEINIT]) afl->mutator->afl_custom_deinit = deinit_py;
 
-  }
+  /* "afl_custom_fuzz" should not be NULL, but the interface of Python mutator
+     is quite different from the custom mutator. */
+  afl->mutator->afl_custom_fuzz = fuzz_py;
 
-  PyTuple_SetItem(py_args, 2, py_value);
+  if (py_functions[PY_FUNC_PRE_SAVE])
+    afl->mutator->afl_custom_pre_save = pre_save_py;
 
-  py_value = PyObject_CallObject(afl->py_functions[PY_FUNC_FUZZ], py_args);
+  if (py_functions[PY_FUNC_INIT_TRIM])
+    afl->mutator->afl_custom_init_trim = init_trim_py;
 
-  Py_DECREF(py_args);
+  if (py_functions[PY_FUNC_POST_TRIM])
+    afl->mutator->afl_custom_post_trim = post_trim_py;
 
-  if (py_value != NULL) {
+  if (py_functions[PY_FUNC_TRIM]) afl->mutator->afl_custom_trim = trim_py;
 
-    mutated_size = PyByteArray_Size(py_value);
-    if (buf_size < mutated_size) *buf = ck_realloc(*buf, mutated_size);
+  if (py_functions[PY_FUNC_HAVOC_MUTATION])
+    afl->mutator->afl_custom_havoc_mutation = havoc_mutation_py;
 
-    memcpy(*buf, PyByteArray_AsString(py_value), mutated_size);
-    Py_DECREF(py_value);
-    return mutated_size;
+  if (py_functions[PY_FUNC_HAVOC_MUTATION_PROBABILITY])
+    afl->mutator->afl_custom_havoc_mutation_probability =
+        havoc_mutation_probability_py;
 
-  } else {
+  if (py_functions[PY_FUNC_QUEUE_GET])
+    afl->mutator->afl_custom_queue_get = queue_get_py;
 
-    PyErr_Print();
-    FATAL("Call failed");
+  if (py_functions[PY_FUNC_QUEUE_NEW_ENTRY])
+    afl->mutator->afl_custom_queue_new_entry = queue_new_entry_py;
 
-  }
+  OKF("Python mutator '%s' installed successfully.", module_name);
+
+  /* Initialize the custom mutator */
+  init_py(afl, py_mutator, rand_below(afl, 0xFFFFFFFF));
 
 }
 
-size_t pre_save_py(afl_state_t *afl, 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;
   PyObject *py_args, *py_value;
@@ -260,7 +350,7 @@ size_t pre_save_py(afl_state_t *afl, u8 *buf, size_t buf_size, u8 **out_buf) {
 
   PyTuple_SetItem(py_args, 0, py_value);
 
-  py_value = PyObject_CallObject(afl->py_functions[PY_FUNC_PRE_SAVE], py_args);
+  py_value = PyObject_CallObject(((py_mutator_t *)py_mutator)->py_functions[PY_FUNC_PRE_SAVE], py_args);
 
   Py_DECREF(py_args);
 
@@ -281,7 +371,7 @@ size_t pre_save_py(afl_state_t *afl, u8 *buf, size_t buf_size, u8 **out_buf) {
 
 }
 
-u32 init_trim_py(afl_state_t *afl, u8 *buf, size_t buf_size) {
+u32 init_trim_py(void *py_mutator, u8 *buf, size_t buf_size) {
 
   PyObject *py_args, *py_value;
 
@@ -296,7 +386,7 @@ u32 init_trim_py(afl_state_t *afl, u8 *buf, size_t buf_size) {
 
   PyTuple_SetItem(py_args, 0, py_value);
 
-  py_value = PyObject_CallObject(afl->py_functions[PY_FUNC_INIT_TRIM], py_args);
+  py_value = PyObject_CallObject(((py_mutator_t *)py_mutator)->py_functions[PY_FUNC_INIT_TRIM], py_args);
   Py_DECREF(py_args);
 
   if (py_value != NULL) {
@@ -318,7 +408,7 @@ u32 init_trim_py(afl_state_t *afl, u8 *buf, size_t buf_size) {
 
 }
 
-u32 post_trim_py(afl_state_t *afl, u8 success) {
+u32 post_trim_py(void *py_mutator, u8 success) {
 
   PyObject *py_args, *py_value;
 
@@ -334,7 +424,7 @@ u32 post_trim_py(afl_state_t *afl, u8 success) {
 
   PyTuple_SetItem(py_args, 0, py_value);
 
-  py_value = PyObject_CallObject(afl->py_functions[PY_FUNC_POST_TRIM], py_args);
+  py_value = PyObject_CallObject(((py_mutator_t *)py_mutator)->py_functions[PY_FUNC_POST_TRIM], py_args);
   Py_DECREF(py_args);
 
   if (py_value != NULL) {
@@ -356,12 +446,12 @@ u32 post_trim_py(afl_state_t *afl, u8 success) {
 
 }
 
-void trim_py(afl_state_t *afl, u8 **out_buf, size_t *out_buf_size) {
+void trim_py(void *py_mutator, u8 **out_buf, size_t *out_buf_size) {
 
   PyObject *py_args, *py_value;
 
   py_args = PyTuple_New(0);
-  py_value = PyObject_CallObject(afl->py_functions[PY_FUNC_TRIM], py_args);
+  py_value = PyObject_CallObject(((py_mutator_t *)py_mutator)->py_functions[PY_FUNC_TRIM], py_args);
   Py_DECREF(py_args);
 
   if (py_value != NULL) {
@@ -380,7 +470,7 @@ void trim_py(afl_state_t *afl, u8 **out_buf, size_t *out_buf_size) {
 
 }
 
-size_t havoc_mutation_py(afl_state_t *afl, u8 **buf, size_t buf_size,
+size_t havoc_mutation_py(void *py_mutator, u8 **buf, size_t buf_size,
                          size_t max_size) {
 
   size_t    mutated_size;
@@ -414,7 +504,7 @@ size_t havoc_mutation_py(afl_state_t *afl, u8 **buf, size_t buf_size,
   PyTuple_SetItem(py_args, 1, py_value);
 
   py_value =
-      PyObject_CallObject(afl->py_functions[PY_FUNC_HAVOC_MUTATION], py_args);
+      PyObject_CallObject(((py_mutator_t *)py_mutator)->py_functions[PY_FUNC_HAVOC_MUTATION], py_args);
 
   Py_DECREF(py_args);
 
@@ -437,13 +527,13 @@ size_t havoc_mutation_py(afl_state_t *afl, u8 **buf, size_t buf_size,
 
 }
 
-u8 havoc_mutation_probability_py(afl_state_t *afl) {
+u8 havoc_mutation_probability_py(void *py_mutator) {
 
   PyObject *py_args, *py_value;
 
   py_args = PyTuple_New(0);
   py_value = PyObject_CallObject(
-      afl->py_functions[PY_FUNC_HAVOC_MUTATION_PROBABILITY], py_args);
+      ((py_mutator_t *)py_mutator)->py_functions[PY_FUNC_HAVOC_MUTATION_PROBABILITY], py_args);
   Py_DECREF(py_args);
 
   if (py_value != NULL) {
@@ -461,7 +551,7 @@ u8 havoc_mutation_probability_py(afl_state_t *afl) {
 
 }
 
-u8 queue_get_py(afl_state_t *afl, const u8 *filename) {
+u8 queue_get_py(void *py_mutator, const u8 *filename) {
 
   PyObject *py_args, *py_value;
 
@@ -483,7 +573,7 @@ u8 queue_get_py(afl_state_t *afl, const u8 *filename) {
   PyTuple_SetItem(py_args, 0, py_value);
 
   // Call Python function
-  py_value = PyObject_CallObject(afl->py_functions[PY_FUNC_QUEUE_GET], py_args);
+  py_value = PyObject_CallObject(((py_mutator_t *)py_mutator)->py_functions[PY_FUNC_QUEUE_GET], py_args);
   Py_DECREF(py_args);
 
   if (py_value != NULL) {
@@ -509,7 +599,7 @@ u8 queue_get_py(afl_state_t *afl, const u8 *filename) {
 
 }
 
-void queue_new_entry_py(afl_state_t *afl, const u8 *filename_new_queue,
+void queue_new_entry_py(void *py_mutator, const u8 *filename_new_queue,
                         const u8 *filename_orig_queue) {
 
   PyObject *py_args, *py_value;
@@ -553,7 +643,7 @@ void queue_new_entry_py(afl_state_t *afl, const u8 *filename_new_queue,
 
   // Call
   py_value =
-      PyObject_CallObject(afl->py_functions[PY_FUNC_QUEUE_NEW_ENTRY], py_args);
+      PyObject_CallObject(((py_mutator_t *)py_mutator)->py_functions[PY_FUNC_QUEUE_NEW_ENTRY], py_args);
   Py_DECREF(py_args);
 
   if (py_value == NULL) {