added afl_custom_fuzz_count

GNUmakefile

@@ -37,18 +37,18 @@ MANPAGES=$(foreach p, $(PROGS) $(SH_PROGS), $(p).8) afl-as.8
 ASAN_OPTIONS=detect_leaks=0
 
 ifeq "$(findstring android, $(shell $(CC) --version 2>/dev/null))" ""
-ifeq "$(shell echo 'int main() {return 0; }' | $(CC) $(CFLAGS) -Werror -x c - -flto=full -o .test 2>/dev/null && echo 1 || echo 0 ; rm -f .test )" "1"
+ ifeq "$(shell echo 'int main() {return 0; }' | $(CC) $(CFLAGS) -Werror -x c - -flto=full -o .test 2>/dev/null && echo 1 || echo 0 ; rm -f .test )" "1"
 	CFLAGS_FLTO ?= -flto=full
-else
- ifeq "$(shell echo 'int main() {return 0; }' | $(CC) $(CFLAGS) -Werror -x c - -flto=thin -o .test 2>/dev/null && echo 1 || echo 0 ; rm -f .test )" "1"
-	CFLAGS_FLTO ?= -flto=thin
  else
-  ifeq "$(shell echo 'int main() {return 0; }' | $(CC) $(CFLAGS) -Werror -x c - -flto -o .test 2>/dev/null && echo 1 || echo 0 ; rm -f .test )" "1"
+  ifeq "$(shell echo 'int main() {return 0; }' | $(CC) $(CFLAGS) -Werror -x c - -flto=thin -o .test 2>/dev/null && echo 1 || echo 0 ; rm -f .test )" "1"
+	CFLAGS_FLTO ?= -flto=thin
+  else
+   ifeq "$(shell echo 'int main() {return 0; }' | $(CC) $(CFLAGS) -Werror -x c - -flto -o .test 2>/dev/null && echo 1 || echo 0 ; rm -f .test )" "1"
 	CFLAGS_FLTO ?= -flto
+   endif
   endif
  endif
 endif
-endif
 
 ifeq "$(shell echo 'int main() {return 0; }' | $(CC) -fno-move-loop-invariants -fdisable-tree-cunrolli -x c - -o .test 2>/dev/null && echo 1 || echo 0 ; rm -f .test )" "1"
 	SPECIAL_PERFORMANCE += -fno-move-loop-invariants -fdisable-tree-cunrolli

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-extras.c

@@ -251,7 +251,7 @@ static void extras_check_and_sort(afl_state_t *afl, u32 min_len, u32 max_len,
   if (afl->extras_cnt > afl->max_det_extras) {
 
     OKF("More than %d tokens - will use them probabilistically.",
-          afl->max_det_extras);
+        afl->max_det_extras);
 
   }
 
@@ -406,7 +406,7 @@ void add_extra(afl_state_t *afl, u8 *mem, u32 len) {
   if (afl->extras_cnt == afl->max_det_extras + 1) {
 
     OKF("More than %d tokens - will use them probabilistically.",
-          afl->max_det_extras);
+        afl->max_det_extras);
 
   }
 

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,104 +1680,119 @@ 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;
 
-      for (afl->stage_cur = 0; afl->stage_cur < afl->stage_max;
+      if (afl->stage_max) {
-           ++afl->stage_cur) {
 
-        struct queue_entry *target;
+        for (afl->stage_cur = 0; afl->stage_cur < afl->stage_max;
-        u32                 tid;
+             ++afl->stage_cur) {
-        u8 *                new_buf;
 
-      retry_external_pick:
+          struct queue_entry *target;
-        /* Pick a random other queue entry for passing to external API */
+          u32                 tid;
+          u8 *                new_buf;
 
-        do {
+        retry_external_pick:
+          /* Pick a random other queue entry for passing to external API */
 
-          tid = rand_below(afl, afl->queued_paths);
+          do {
 
-        } while (tid == afl->current_entry && afl->queued_paths > 1);
+            tid = rand_below(afl, afl->queued_paths);
 
-        target = afl->queue;
+          } while (tid == afl->current_entry && afl->queued_paths > 1);
 
-        while (tid >= 100) {
+          target = afl->queue;
 
-          target = target->next_100;
+          while (tid >= 100) {
-          tid -= 100;
 
-        }
+            target = target->next_100;
-
+            tid -= 100;
-        while (tid--) {
-
-          target = target->next;
-
-        }
-
-        /* Make sure that the target has a reasonable length. */
-
-        while (target && (target->len < 2 || target == afl->queue_cur) &&
-               afl->queued_paths > 3) {
-
-          target = target->next;
-          ++afl->splicing_with;
-
-        }
-
-        if (!target) { goto retry_external_pick; }
-
-        /* 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); }
-
-        new_buf = afl_realloc(AFL_BUF_PARAM(out_scratch), target->len);
-        if (unlikely(!new_buf)) { PFATAL("alloc"); }
-        ck_read(fd, new_buf, target->len, target->fname);
-        close(fd);
-
-        u8 *mutated_buf = NULL;
-
-        size_t mutated_size =
-            el->afl_custom_fuzz(el->data, out_buf, len, &mutated_buf, new_buf,
-                                target->len, max_seed_size);
-
-        if (unlikely(!mutated_buf)) {
-
-          FATAL("Error in custom_fuzz. Size returned: %zd", mutated_size);
-
-        }
-
-        if (mutated_size > 0) {
-
-          if (common_fuzz_stuff(afl, mutated_buf, (u32)mutated_size)) {
-
-            goto abandon_entry;
 
           }
 
-          /* If we're finding new stuff, let's run for a bit longer, limits
+          while (tid--) {
-            permitting. */
 
-          if (afl->queued_paths != havoc_queued) {
+            target = target->next;
 
-            if (perf_score <= afl->havoc_max_mult * 100) {
+          }
 
-              afl->stage_max *= 2;
+          /* Make sure that the target has a reasonable length. */
-              perf_score *= 2;
+
+          while (target && (target->len < 2 || target == afl->queue_cur) &&
+                 afl->queued_paths > 3) {
+
+            target = target->next;
+            ++afl->splicing_with;
+
+          }
+
+          if (!target) { goto retry_external_pick; }
+
+          /* 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);
+
+          }
+
+          new_buf = afl_realloc(AFL_BUF_PARAM(out_scratch), target->len);
+          if (unlikely(!new_buf)) { PFATAL("alloc"); }
+          ck_read(fd, new_buf, target->len, target->fname);
+          close(fd);
+
+          u8 *mutated_buf = NULL;
+
+          size_t mutated_size =
+              el->afl_custom_fuzz(el->data, out_buf, len, &mutated_buf, new_buf,
+                                  target->len, max_seed_size);
+
+          if (unlikely(!mutated_buf)) {
+
+            FATAL("Error in custom_fuzz. Size returned: %zd", mutated_size);
+
+          }
+
+          if (mutated_size > 0) {
+
+            if (common_fuzz_stuff(afl, mutated_buf, (u32)mutated_size)) {
+
+              goto abandon_entry;
 
             }
 
-            havoc_queued = afl->queued_paths;
+            /* If we're finding new stuff, let's run for a bit longer, limits
+              permitting. */
+
+            if (afl->queued_paths != havoc_queued) {
+
+              if (perf_score <= afl->havoc_max_mult * 100) {
+
+                afl->stage_max *= 2;
+                perf_score *= 2;
+
+              }
+
+              havoc_queued = afl->queued_paths;
+
+            }
 
           }
 
-        }
+          /* `(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);
 
-        /* `(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);
 
       }
 

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;