Add two new hooks for the custom mutator

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

8 files changed, 313 insertions, 20 deletions

diff --git a/docs/custom_mutators.md b/docs/custom_mutators.md
index b31a2e4f..6e951008 100644
--- a/docs/custom_mutators.md
+++ b/docs/custom_mutators.md
@@ -28,14 +28,17 @@ performed with the custom mutator.
 C/C++:
 ```c
 void afl_custom_init(unsigned int seed);
-size_t afl_custom_fuzz(u8** buf, size_t buf_size, u8* add_buf,
+size_t afl_custom_fuzz(uint8_t** buf, size_t buf_size, uint8_t* add_buf,
                        size_t add_buf_size, size_t max_size);
-size_t afl_custom_pre_save(u8* buf, size_t buf_size, u8** out_buf);
-u32 afl_custom_init_trim(u8* buf, size_t buf_size);
-void afl_custom_trim(u8** out_buf, size_t* out_buf_size);
-u32 afl_custom_post_trim(u8 success);
+size_t afl_custom_pre_save(uint8_t* buf, size_t buf_size, uint8_t** out_buf);
+uint32_t afl_custom_init_trim(uint8_t* buf, size_t buf_size);
+void afl_custom_trim(uint8_t** out_buf, size_t* out_buf_size);
+uint32_t afl_custom_post_trim(uint8_t success);
 size_t afl_custom_havoc_mutation(uint8_t** buf, size_t buf_size, size_t max_size);
 uint8_t afl_custom_havoc_mutation_probability(void);
+uint8_t afl_custom_queue_get(const uint8_t* filename);
+void afl_custom_queue_new_entry(const uint8_t* filename_new_queue,
+                                const uint8_t* filename_orig_queue);
 ```
 
 Python:
@@ -63,6 +66,12 @@ def havoc_mutation(buf, max_size):
 
 def havoc_mutation_probability():
     return probability # int in [0, 100]
+
+def queue_get(filename):
+    return True
+
+def queue_new_entry(filename_new_queue, filename_orig_queue):
+    pass
 ```
 
 ### Custom Mutation
@@ -71,21 +80,37 @@ def havoc_mutation_probability():
 
     This method is called when AFL++ starts up and is used to seed RNG.
 
+- `queue_get` (optional):
+
+    This method determines whether the fuzzer should fuzz the current queue
+    entry or not
+
 - `fuzz` (required):
 
     This method performs custom mutations on a given input. It also accepts an
     additional test case.
 
+- `havoc_mutation` and `havoc_mutation_probability` (optional):
+
+    `havoc_mutation` performs a single custom mutation on a given input. This
+    mutation is stacked with the other mutations in havoc. The other method,
+    `havoc_mutation_probability`, returns the probability that `havoc_mutation`
+    is called in havoc. By default, it is 6%.
+
 - `pre_save` (optional):
 
     For some cases, the format of the mutated data returned from the custom
     mutator is not suitable to directly execute the target with this input.
     For example, when using libprotobuf-mutator, the data returned is in a
     protobuf format which corresponds to a given grammar. In order to execute
-    the target, the protobuf data must be converted to the plain-text format expected by the target. In such scenarios, the user can define the
+    the target, the protobuf data must be converted to the plain-text format
+    expected by the target. In such scenarios, the user can define the
     `pre_save` function. This function is then transforms the data into the
     format expected by the API before executing the target.
 
+- `queue_new_entry` (optional):
+
+    This methods is called after adding a new test case to the queue. 
 
 ### Trimming Support
 
diff --git a/examples/custom_mutators/example.c b/examples/custom_mutators/example.c
index d34b3045..178d39b3 100644
--- a/examples/custom_mutators/example.c
+++ b/examples/custom_mutators/example.c
@@ -57,7 +57,7 @@ size_t afl_custom_fuzz(uint8_t **buf, size_t buf_size,
   // Mutate the payload of the packet
   for (int i = 3; i < mutated_size; i++) {
 
-    mutated_out[i] = (buf[i] + rand() % 10) & 0xff;
+    mutated_out[i] = (mutated_out[i] + rand() % 10) & 0xff;
 
   }
 
@@ -93,10 +93,10 @@ size_t afl_custom_pre_save(uint8_t *buf, size_t buf_size, uint8_t **out_buf) {
 
 }
 
-uint8_t *trim_buf;
-size_t trim_buf_size;
-int trimmming_steps;
-int cur_step;
+static uint8_t *trim_buf;
+static size_t trim_buf_size;
+static int trimmming_steps;
+static int cur_step;
 
 /**
  * This method is called at the start of each trimming operation and receives
@@ -186,9 +186,11 @@ int afl_custom_post_trim(int success) {
  *
  * (Optional)
  *
- * @param[in] buf Pointer to the input data to be mutated
+ * @param[inout] buf Pointer to the input data to be mutated and the mutated
+ *     output
  * @param[in] buf_size Size of input data
- * @param[in] max_size Maximum size of the mutated output. The mutation must not produce data larger than max_size.
+ * @param[in] max_size Maximum size of the mutated output. The mutation must
+ *     not produce data larger than max_size.
  * @return Size of the mutated output.
  */
 size_t afl_custom_havoc_mutation(uint8_t** buf, size_t buf_size, size_t max_size) {
@@ -221,3 +223,35 @@ uint8_t afl_custom_havoc_mutation_probability(void) {
   return 5; // 5 %
 
 }
+
+/**
+ * Determine whether the fuzzer should fuzz the queue entry or not.
+ *
+ * (Optional)
+ *
+ * @param filename File name of the test case in the queue entry
+ * @return Return True(1) if the fuzzer will fuzz the queue entry, and
+ *     False(0) otherwise.
+ */
+uint8_t afl_custom_queue_get(const uint8_t* filename) {
+
+  return 1;
+
+}
+
+/**
+ * Allow for additional analysis (e.g. calling a different tool that does a 
+ * different kind of coverage and saves this for the custom mutator).
+ *
+ * (Optional)
+ *
+ * @param filename_new_queue File name of the new queue entry
+ * @param filename_orig_queue File name of the original queue entry
+ */
+void afl_custom_queue_new_entry(const uint8_t* filename_new_queue,
+                                const uint8_t* filename_orig_queue) {
+
+  /* Additional analysis on the original or new test case */
+
+}
+
diff --git a/examples/custom_mutators/example.py b/examples/custom_mutators/example.py
index a68f2ee5..6bacfa05 100644
--- a/examples/custom_mutators/example.py
+++ b/examples/custom_mutators/example.py
@@ -120,3 +120,55 @@ def fuzz(buf, add_buf, max_size):
 #     '''
 #     return buf
 #
+# def havoc_mutation(buf, max_size):
+#     '''
+#     Perform a single custom mutation on a given input.
+#
+#     @type buf: bytearray
+#     @param buf: The buffer that should be mutated.
+#
+#     @type max_size: int
+#     @param max_size: Maximum size of the mutated output. The mutation must not
+#         produce data larger than max_size.
+#
+#     @rtype: bytearray
+#     @return: A new bytearray containing the mutated data
+#     '''
+#     return mutated_buf
+#
+# def havoc_mutation_probability():
+#     '''
+#     Called for each `havoc_mutation`. Return the probability (in percentage)
+#     that `havoc_mutation` is called in havoc. Be default it is 6%.
+#
+#     @rtype: int
+#     @return: The probability (0-100)
+#     '''
+#     return prob
+#
+# def queue_get(filename):
+#     '''
+#     Called at the beginning of each fuzz iteration to determine whether the
+#     test case should be fuzzed
+#
+#     @type filename: str
+#     @param filename: File name of the test case in the current queue entry
+#
+#     @rtype: bool
+#     @return: Return True if the custom mutator decides to fuzz the test case,
+#         and False otherwise
+#     '''
+#     return True
+#
+# def queue_new_entry(filename_new_queue, filename_orig_queue):
+#     '''
+#     Called after adding a new test case to the queue
+#
+#     @type filename_new_queue: str 
+#     @param filename_new_queue: File name of the new queue entry
+#
+#     @type filename_orig_queue: str
+#     @param filename_orig_queue: File name of the original queue entry
+#     '''
+#     pass
+
diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index 57639411..8c0e7ca9 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -480,8 +480,9 @@ struct custom_mutator {
    *
    * (Optional for now. Required in the future)
    *
-   * @param[in] buf Pointer to input data to be mutated
-   * @param[in] buf_size Size of input data
+   * @param[inout] buf Pointer to the input data to be mutated and the mutated
+   *     output
+   * @param[in] buf_size Size of the input/output data
    * @param[in] add_buf Buffer containing the additional test case
    * @param[in] add_buf_size Size of the additional test case
    * @param[in] max_size Maximum size of the mutated output. The mutation must not
@@ -566,9 +567,11 @@ struct custom_mutator {
    *
    * (Optional)
    *
-   * @param[in] buf Pointer to the input data to be mutated
+   * @param[inout] buf Pointer to the input data to be mutated and the mutated
+   *     output
    * @param[in] buf_size Size of input data
-   * @param[in] max_size Maximum size of the mutated output. The mutation must not produce data larger than max_size.
+   * @param[in] max_size Maximum size of the mutated output. The mutation must
+   *     not produce data larger than max_size.
    * @return Size of the mutated output.
    */
   size_t (*afl_custom_havoc_mutation)(u8** buf, size_t buf_size, size_t max_size);
@@ -582,7 +585,30 @@ struct custom_mutator {
    * @return The probability (0-100).
    */
   u8 (*afl_custom_havoc_mutation_probability)(void);
-  
+
+  /**
+   * Determine whether the fuzzer should fuzz the current queue entry or not.
+   *
+   * (Optional)
+   *
+   * @param filename File name of the test case in the queue entry
+   * @return Return True(1) if the fuzzer will fuzz the queue entry, and
+   *     False(0) otherwise.
+   */
+  u8 (*afl_custom_queue_get)(const u8* filename);
+
+  /**
+   * Allow for additional analysis (e.g. calling a different tool that does a 
+   * different kind of coverage and saves this for the custom mutator).
+   *
+   * (Optional)
+   *
+   * @param filename_new_queue File name of the new queue entry
+   * @param filename_orig_queue File name of the original queue entry. This
+   *     argument can be NULL while initializing the fuzzer
+   */
+  void (*afl_custom_queue_new_entry)(const u8* filename_new_queue,
+                                     const u8* filename_orig_queue);
 };
 
 extern struct custom_mutator* mutator;
@@ -634,6 +660,8 @@ enum {
   /* 05 */ PY_FUNC_TRIM,
   /* 06 */ PY_FUNC_HAVOC_MUTATION,
   /* 07 */ PY_FUNC_HAVOC_MUTATION_PROBABILITY,
+  /* 08 */ PY_FUNC_QUEUE_GET,
+  /* 09 */ PY_FUNC_QUEUE_NEW_ENTRY,
   PY_FUNC_COUNT
 
 };
@@ -663,6 +691,8 @@ u32    post_trim_py(u8);
 void   trim_py(u8**, size_t*);
 size_t havoc_mutation_py(u8**, size_t, size_t);
 u8     havoc_mutation_probability_py(void);
+u8     queue_get_py(const u8*);
+void   queue_new_entry_py(const u8*, const u8*);
 
 #endif
 
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. */