added afl_custom_fuzz_count

docs/Changelog.md

@@ -18,6 +18,9 @@ sending a mail to <afl-users+subscribe@googlegroups.com>.
       dict entries without recompiling.
     - AFL_FORKSRV_INIT_TMOUT env variable added to control the time to wait for
       the forkserver to come up without the need to increase the overall timeout.
+  - custom mutators:
+    - added afl_custom_fuzz_count/fuzz_count function to allow specifying the 
+      number of fuzz attempts for custom_fuzz
   - llvm_mode:
     - Ported SanCov to LTO, and made it the default for LTO. better
       instrumentation locations

docs/custom_mutators.md

@@ -32,6 +32,7 @@ performed with the custom mutator.
 C/C++:
 ```c
 void *afl_custom_init(afl_t *afl, unsigned int seed);
+uint32_t afl_custom_fuzz_count(void *data, const u8 *buf, size_t buf_size);
 size_t afl_custom_fuzz(void *data, uint8_t *buf, size_t buf_size, u8 **out_buf, uint8_t *add_buf, size_t add_buf_size, size_t max_size);
 size_t afl_custom_post_process(void *data, uint8_t *buf, size_t buf_size, uint8_t **out_buf);
 int32_t afl_custom_init_trim(void *data, uint8_t *buf, size_t buf_size);
@@ -49,6 +50,9 @@ Python:
 def init(seed):
     pass
 
+def fuzz_count(buf, add_buf, max_size):
+    return cnt
+
 def fuzz(buf, add_buf, max_size):
     return mutated_out
 
@@ -88,6 +92,11 @@ def queue_new_entry(filename_new_queue, filename_orig_queue):
     This method determines whether the custom fuzzer should fuzz the current
     queue entry or not
 
+- `fuzz_count` (optional):
+
+    This method can be used to instruct afl-fuzz how often to perform a fuzz
+    attempt on this input data.
+
 - `fuzz` (optional):
 
     This method performs custom mutations on a given input. It also accepts an

include/afl-fuzz.h

@@ -288,6 +288,7 @@ enum {
 enum {
 
   /* 00 */ PY_FUNC_INIT,
+  /* 01 */ PY_FUNC_FUZZ_COUNT,
   /* 01 */ PY_FUNC_FUZZ,
   /* 02 */ PY_FUNC_POST_PROCESS,
   /* 03 */ PY_FUNC_INIT_TRIM,
@@ -679,6 +680,24 @@ struct custom_mutator {
    */
   void *(*afl_custom_init)(afl_state_t *afl, unsigned int seed);
 
+  /**
+   * This method is called just before fuzzing a queue entry with the custom
+   * mutator, and receives the initial buffer. It should return the number of
+   * fuzzes to perform.
+   *
+   * A value of 0 means no fuzzing of this queue entry.
+   *
+   * The function is now allowed to change the data.
+   *
+   * (Optional)
+   *
+   * @param data pointer returned in afl_custom_init for this fuzz case
+   * @param buf Buffer containing the test case
+   * @param buf_size Size of the test case
+   * @return The amount of fuzzes to perform on this queue entry, 0 = skip
+   */
+  u32 (*afl_custom_fuzz_count)(void *data, const u8 *buf, size_t buf_size);
+
   /**
    * Perform custom mutations on a given input
    *
@@ -867,6 +886,7 @@ u8   trim_case_custom(afl_state_t *, struct queue_entry *q, u8 *in_buf,
 struct custom_mutator *load_custom_mutator_py(afl_state_t *, char *);
 void                   finalize_py_module(void *);
 
+u32    fuzz_count_py(void *, const u8 *, size_t);
 size_t post_process_py(void *, u8 *, size_t, u8 **);
 s32    init_trim_py(void *, u8 *, size_t);
 s32    post_trim_py(void *, u8);

src/afl-forkserver.c

@@ -634,7 +634,7 @@ void afl_fsrv_start(afl_forkserver_t *fsrv, char **argv,
 
         if (fsrv->add_extra_func == NULL || fsrv->afl_ptr == NULL) {
 
-          // this is not afl-fuzz - we deny and return
+          // this is not afl-fuzz - or it is cmplog - we deny and return
           if (fsrv->use_shmem_fuzz) {
 
             status = (FS_OPT_ENABLED | FS_OPT_SHDMEM_FUZZ);

src/afl-fuzz-mutators.c

@@ -166,6 +166,11 @@ struct custom_mutator *load_custom_mutator(afl_state_t *afl, const char *fn) {
 
   }
 
+  /* "afl_custom_fuzz_count", optional */
+  mutator->afl_custom_fuzz_count = dlsym(dh, "afl_custom_fuzz_count");
+  if (!mutator->afl_custom_fuzz_count)
+    ACTF("optional symbol 'afl_custom_fuzz_count' not found.");
+
   /* "afl_custom_deinit", optional for backward compatibility */
   mutator->afl_custom_deinit = dlsym(dh, "afl_custom_deinit");
   if (!mutator->afl_custom_deinit)

src/afl-fuzz-one.c

@@ -1672,7 +1672,7 @@ custom_mutator_stage:
 
   if (afl->stage_max < HAVOC_MIN) { afl->stage_max = HAVOC_MIN; }
 
-  const u32 max_seed_size = MAX_FILE;
+  const u32 max_seed_size = MAX_FILE, saved_max = afl->stage_max;
 
   orig_hit_cnt = afl->queued_paths + afl->unique_crashes;
 
@@ -1680,10 +1680,17 @@ custom_mutator_stage:
 
     if (el->afl_custom_fuzz) {
 
+      if (el->afl_custom_fuzz_count)
+        afl->stage_max = el->afl_custom_fuzz_count(el->data, out_buf, len);
+      else
+        afl->stage_max = saved_max;
+
       has_custom_fuzz = true;
 
       afl->stage_short = el->name_short;
 
+      if (afl->stage_max) {
+
         for (afl->stage_cur = 0; afl->stage_cur < afl->stage_max;
              ++afl->stage_cur) {
 
@@ -1729,7 +1736,11 @@ custom_mutator_stage:
 
           /* Read the additional testcase into a new buffer. */
           fd = open(target->fname, O_RDONLY);
-        if (unlikely(fd < 0)) { PFATAL("Unable to open '%s'", target->fname); }
+          if (unlikely(fd < 0)) {
+
+            PFATAL("Unable to open '%s'", target->fname);
+
+          }
 
           new_buf = afl_realloc(AFL_BUF_PARAM(out_scratch), target->len);
           if (unlikely(!new_buf)) { PFATAL("alloc"); }
@@ -1774,8 +1785,10 @@ custom_mutator_stage:
 
           }
 
-        /* `(afl->)out_buf` may have been changed by the call to custom_fuzz */
-        /* TODO: Only do this when `mutated_buf` == `out_buf`? Branch vs Memcpy.
+          /* `(afl->)out_buf` may have been changed by the call to custom_fuzz
+           */
+          /* TODO: Only do this when `mutated_buf` == `out_buf`? Branch vs
+           * Memcpy.
            */
           memcpy(out_buf, in_buf, len);
 
@@ -1783,6 +1796,8 @@ custom_mutator_stage:
 
       }
 
+    }
+
   });
 
   if (!has_custom_fuzz) goto havoc_stage;

src/afl-fuzz-python.c

@@ -347,6 +347,12 @@ struct custom_mutator *load_custom_mutator_py(afl_state_t *afl,
 
   }
 
+  if (py_functions[PY_FUNC_FUZZ_COUNT]) {
+
+    mutator->afl_custom_fuzz_count = fuzz_count_py;
+
+  }
+
   if (py_functions[PY_FUNC_POST_TRIM]) {
 
     mutator->afl_custom_post_trim = post_trim_py;
@@ -477,6 +483,44 @@ s32 init_trim_py(void *py_mutator, u8 *buf, size_t buf_size) {
 
 }
 
+u32 fuzz_count_py(void *py_mutator, const u8 *buf, size_t buf_size) {
+
+  PyObject *py_args, *py_value;
+
+  py_args = PyTuple_New(1);
+  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);
+
+  py_value = PyObject_CallObject(
+      ((py_mutator_t *)py_mutator)->py_functions[PY_FUNC_FUZZ_COUNT], py_args);
+  Py_DECREF(py_args);
+
+  if (py_value != NULL) {
+
+  #if PY_MAJOR_VERSION >= 3
+    u32 retcnt = (u32)PyLong_AsLong(py_value);
+  #else
+    u32 retcnt = PyInt_AsLong(py_value);
+  #endif
+    Py_DECREF(py_value);
+    return retcnt;
+
+  } else {
+
+    PyErr_Print();
+    FATAL("Call failed");
+
+  }
+
+}
+
 s32 post_trim_py(void *py_mutator, u8 success) {
 
   PyObject *py_args, *py_value;