Add two new hooks for the custom mutator

- `afl_custom_queue_get` and `afl_custom_queue_new_entry`
- Update the corresponding document and examples

4 files changed, 154 insertions, 2 deletions

diff --git a/src/afl-fuzz-mutators.c b/src/afl-fuzz-mutators.c
index 76ce2c96..28d21636 100644
--- a/src/afl-fuzz-mutators.c
+++ b/src/afl-fuzz-mutators.c
@@ -162,6 +162,16 @@ void load_custom_mutator(const char* fn) {
   if (!mutator->afl_custom_havoc_mutation_probability)
     WARNF("Symbol 'afl_custom_havoc_mutation_probability' not found.");
 
+  /* "afl_custom_queue_get", optional */
+  mutator->afl_custom_queue_get = dlsym(dh, "afl_custom_queue_get");
+  if (!mutator->afl_custom_queue_get)
+    WARNF("Symbol 'afl_custom_queue_get' not found.");
+
+  /* "afl_custom_queue_new_entry", optional */
+  mutator->afl_custom_queue_new_entry = dlsym(dh, "afl_custom_queue_new_entry");
+  if (!mutator->afl_custom_queue_new_entry)
+    WARNF("Symbol 'afl_custom_queue_new_entry' not found");
+
   OKF("Custom mutator '%s' installed successfully.", fn);
 
   /* Initialize the custom mutator */
@@ -324,6 +334,12 @@ void load_custom_mutator_py(const char* module_name) {
   if (py_functions[PY_FUNC_HAVOC_MUTATION_PROBABILITY])
     mutator->afl_custom_havoc_mutation_probability = havoc_mutation_probability_py;
 
+  if (py_functions[PY_FUNC_QUEUE_GET])
+    mutator->afl_custom_queue_get = queue_get_py;
+
+  if (py_functions[PY_FUNC_QUEUE_NEW_ENTRY])
+    mutator->afl_custom_queue_new_entry = queue_new_entry_py;
+
   OKF("Python mutator '%s' installed successfully.", module_name);
 
   /* Initialize the custom mutator */
diff --git a/src/afl-fuzz-one.c b/src/afl-fuzz-one.c
index f12f4a67..1817bd03 100644
--- a/src/afl-fuzz-one.c
+++ b/src/afl-fuzz-one.c
@@ -355,6 +355,15 @@ u8 fuzz_one_original(char** argv) {
 
 #else
 
+  if (mutator && mutator->afl_custom_queue_get) {
+
+    /* The custom mutator will decide to skip this test case or not. */
+
+    if (!mutator->afl_custom_queue_get(queue_cur->fname))
+      return 1;
+
+  }
+
   if (pending_favored) {
 
     /* If we have any favored, non-fuzzed new arrivals in the queue,
diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
index 32f9f6ab..8ceb6957 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -55,8 +55,14 @@ int init_py_module(u8* module_name) {
     py_functions[PY_FUNC_POST_TRIM] =
         PyObject_GetAttrString(py_module, "post_trim");
     py_functions[PY_FUNC_TRIM] = PyObject_GetAttrString(py_module, "trim");
-    py_functions[PY_FUNC_HAVOC_MUTATION] = PyObject_GetAttrString(py_module, "havoc_mutation");
-    py_functions[PY_FUNC_HAVOC_MUTATION_PROBABILITY] = PyObject_GetAttrString(py_module, "havoc_mutation_probability");
+    py_functions[PY_FUNC_HAVOC_MUTATION] =
+        PyObject_GetAttrString(py_module, "havoc_mutation");
+    py_functions[PY_FUNC_HAVOC_MUTATION_PROBABILITY] =
+        PyObject_GetAttrString(py_module, "havoc_mutation_probability");
+    py_functions[PY_FUNC_QUEUE_GET] =
+        PyObject_GetAttrString(py_module, "queue_get");
+    py_functions[PY_FUNC_QUEUE_NEW_ENTRY] =
+        PyObject_GetAttrString(py_module, "queue_new_entry");
 
     for (py_idx = 0; py_idx < PY_FUNC_COUNT; ++py_idx) {
 
@@ -73,6 +79,12 @@ int init_py_module(u8* module_name) {
           if (PyErr_Occurred()) PyErr_Print();
           py_notrim = 1;
 
+        } else if ((py_idx >= PY_FUNC_HAVOC_MUTATION) &&
+                   (py_idx <= PY_FUNC_QUEUE_NEW_ENTRY)) {
+
+          // Implenting the havoc and queue API is optional for now
+          if (PyErr_Occurred()) PyErr_Print();
+
         } else {
 
           if (PyErr_Occurred()) PyErr_Print();
@@ -442,5 +454,109 @@ u8 havoc_mutation_probability_py(void) {
 
 }
 
+u8 queue_get_py(const u8* filename) {
+
+  PyObject *py_args, *py_value;
+
+  py_args = PyTuple_New(1);
+
+  // File name
+#if PY_MAJOR_VERSION >= 3
+  py_value = PyUnicode_FromString(filename);
+#else
+  py_value = PyString_FromString(filename);
+#endif
+  if (!py_value) {
+
+    Py_DECREF(py_args);
+    FATAL("Failed to convert arguments");
+
+  }
+
+  PyTuple_SetItem(py_args, 0, py_value);
+
+  // Call Python function
+  py_value = PyObject_CallObject(py_functions[PY_FUNC_QUEUE_GET], py_args);
+  Py_DECREF(py_args);
+
+  if (py_value != NULL) {
+
+    int ret = PyObject_IsTrue(py_value);
+    Py_DECREF(py_value);
+    
+    if (ret == -1) {
+
+      PyErr_Print();
+      FATAL("Failed to convert return value");
+
+    }
+
+    return (u8) ret & 0xFF;
+
+  } else {
+    
+    PyErr_Print();
+    FATAL("Call failed");
+
+  }
+
+}
+
+void queue_new_entry_py(const u8* filename_new_queue,
+                        const u8* filename_orig_queue) {
+
+  PyObject *py_args, *py_value;
+
+  py_args = PyTuple_New(2);
+
+  // New queue
+#if PY_MAJOR_VERSION >= 3
+  py_value = PyUnicode_FromString(filename_new_queue);
+#else
+  py_value = PyString_FromString(filename_new_queue);
+#endif
+  if (!py_value) {
+  
+    Py_DECREF(py_args);
+    FATAL("Failed to convert arguments");
+
+  }
+
+  PyTuple_SetItem(py_args, 0, py_value);
+
+  // Orig queue
+  py_value = Py_None;
+  if (filename_orig_queue) {
+
+#if PY_MAJOR_VERSION >= 3
+    py_value = PyUnicode_FromString(filename_orig_queue);
+#else
+    py_value = PyString_FromString(filename_orig_queue);
+#endif
+    if (!py_value) {
+ 
+      Py_DECREF(py_args);
+      FATAL("Failed to convert arguments");
+
+    }
+
+  }
+
+  PyTuple_SetItem(py_args, 1, py_value);
+
+  // Call
+  py_value = PyObject_CallObject(py_functions[PY_FUNC_QUEUE_NEW_ENTRY],
+                                 py_args);
+  Py_DECREF(py_args);
+
+  if (py_value == NULL) {
+
+    PyErr_Print();
+    FATAL("Call failed");
+
+  }
+
+}
+
 #endif                                                        /* USE_PYTHON */
 
diff --git a/src/afl-fuzz-queue.c b/src/afl-fuzz-queue.c
index 0880de75..ad9dad13 100644
--- a/src/afl-fuzz-queue.c
+++ b/src/afl-fuzz-queue.c
@@ -139,6 +139,17 @@ void add_to_queue(u8* fname, u32 len, u8 passed_det) {
 
   last_path_time = get_cur_time();
 
+  if (mutator && mutator->afl_custom_queue_new_entry) {
+
+    u8* fname_orig = NULL;
+
+    /* At the initialization stage, queue_cur is NULL */
+    if (queue_cur) fname_orig = queue_cur->fname;
+
+    mutator->afl_custom_queue_new_entry(fname, fname_orig);
+
+  }
+
 }
 
 /* Destroy the entire queue. */