Merge pull request #1915 from yangzao/dev

add custom mutator function for running script after target gets executed

5 files changed, 82 insertions, 25 deletions

diff --git a/src/afl-fuzz-mutators.c b/src/afl-fuzz-mutators.c
index 64dbe7c6..17fb9368 100644
--- a/src/afl-fuzz-mutators.c
+++ b/src/afl-fuzz-mutators.c
@@ -397,6 +397,18 @@ struct custom_mutator *load_custom_mutator(afl_state_t *afl, const char *fn) {
 
   }
 
+  /* "afl_custom_post_run", optional */
+  mutator->afl_custom_post_run = dlsym(dh, "afl_custom_post_run");
+  if (!mutator->afl_custom_post_run) {
+
+    ACTF("optional symbol 'afl_custom_post_run' not found.");
+
+  } else {
+
+    OKF("Found 'afl_custom_post_run'.");
+
+  }
+
   /* "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) {
diff --git a/src/afl-fuzz-one.c b/src/afl-fuzz-one.c
index 67dafda8..01e34b69 100644
--- a/src/afl-fuzz-one.c
+++ b/src/afl-fuzz-one.c
@@ -1894,6 +1894,7 @@ custom_mutator_stage:
   LIST_FOREACH(&afl->custom_mutator_list, struct custom_mutator, {
 
     if (el->afl_custom_fuzz) {
+
       havoc_queued = afl->queued_items;
 
       afl->current_custom_fuzz = el;
diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
index 7dad0770..4c7da774 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -249,6 +249,8 @@ static py_mutator_t *init_py_module(afl_state_t *afl, u8 *module_name) {
         PyObject_GetAttrString(py_module, "queue_get");
     py_functions[PY_FUNC_FUZZ_SEND] =
         PyObject_GetAttrString(py_module, "fuzz_send");
+    py_functions[PY_FUNC_POST_RUN] =
+        PyObject_GetAttrString(py_module, "post_run");
     py_functions[PY_FUNC_SPLICE_OPTOUT] =
         PyObject_GetAttrString(py_module, "splice_optout");
     if (py_functions[PY_FUNC_SPLICE_OPTOUT]) { afl->custom_splice_optout = 1; }
@@ -468,6 +470,12 @@ struct custom_mutator *load_custom_mutator_py(afl_state_t *afl,
 
   }
 
+  if (py_functions[PY_FUNC_POST_RUN]) {
+
+    mutator->afl_custom_post_run = post_run_py;
+
+  }
+
   if (py_functions[PY_FUNC_SPLICE_OPTOUT]) {
 
     mutator->afl_custom_splice_optout = splice_optout_py;
@@ -925,6 +933,28 @@ void fuzz_send_py(void *py_mutator, const u8 *buf, size_t buf_size) {
 
 }
 
+void post_run_py(void *py_mutator) {
+
+  PyObject *py_args, *py_value;
+
+  py_args = PyTuple_New(0);
+  py_value = PyObject_CallObject(
+      ((py_mutator_t *)py_mutator)->py_functions[PY_FUNC_POST_RUN], py_args);
+  Py_DECREF(py_args);
+
+  if (py_value != NULL) {
+
+    Py_DECREF(py_value);
+
+  } else {
+
+    PyErr_Print();
+    FATAL("Call failed");
+
+  }
+
+}
+
 u8 queue_new_entry_py(void *py_mutator, const u8 *filename_new_queue,
                       const u8 *filename_orig_queue) {
 
diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index ac4fb4a9..ae7969a6 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -60,6 +60,19 @@ fuzz_run_target(afl_state_t *afl, afl_forkserver_t *fsrv, u32 timeout) {
 
   fsrv_run_result_t res = afl_fsrv_run_target(fsrv, timeout, &afl->stop_soon);
 
+  /* If post_run() function is defined in custom mutator, the function will be
+     called each time after AFL++ executes the target program. */
+  
+  if (unlikely(afl->custom_mutators_count)) {
+
+    LIST_FOREACH(&afl->custom_mutator_list, struct custom_mutator, {
+
+      if (el->afl_custom_post_run) { el->afl_custom_post_run(el->data); }
+
+    });
+
+  }
+
 #ifdef PROFILING
   clock_gettime(CLOCK_REALTIME, &spec);
   time_spent_start = (spec.tv_sec * 1000000000) + spec.tv_nsec;
@@ -1110,4 +1123,3 @@ common_fuzz_stuff(afl_state_t *afl, u8 *out_buf, u32 len) {
   return 0;
 
 }
-
diff --git a/src/afl-fuzz-stats.c b/src/afl-fuzz-stats.c
index f212a4b8..deb28b7a 100644
--- a/src/afl-fuzz-stats.c
+++ b/src/afl-fuzz-stats.c
@@ -138,7 +138,7 @@ void load_stats_file(afl_state_t *afl) {
 
   FILE *f;
   u8    buf[MAX_LINE];
-  u8 *  lptr;
+  u8   *lptr;
   u8    fn[PATH_MAX];
   u32   lineno = 0;
   snprintf(fn, PATH_MAX, "%s/fuzzer_stats", afl->out_dir);
@@ -421,7 +421,7 @@ void write_stats_file(afl_state_t *afl, u32 t_bytes, double bitmap_cvg,
 void write_queue_stats(afl_state_t *afl) {
 
   FILE *f;
-  u8 *  fn = alloc_printf("%s/queue_data", afl->out_dir);
+  u8   *fn = alloc_printf("%s/queue_data", afl->out_dir);
   if ((f = fopen(fn, "w")) != NULL) {
 
     u32 id;
@@ -857,8 +857,9 @@ void show_stats_normal(afl_state_t *afl) {
   /* Since `total_crashes` does not get reloaded from disk on restart,
     it indicates if we found crashes this round already -> paint red.
     If it's 0, but `saved_crashes` is set from a past run, paint in yellow. */
-  char *crash_color =
-      afl->total_crashes ? cLRD : afl->saved_crashes ? cYEL : cRST;
+  char *crash_color = afl->total_crashes   ? cLRD
+                      : afl->saved_crashes ? cYEL
+                                           : cRST;
 
   /* Lord, forgive me this. */
 
@@ -881,26 +882,26 @@ void show_stats_normal(afl_state_t *afl) {
 
     } else
 
-        /* Subsequent cycles, but we're still making finds. */
-        if (afl->cycles_wo_finds < 25 || min_wo_finds < 30) {
+      /* Subsequent cycles, but we're still making finds. */
+      if (afl->cycles_wo_finds < 25 || min_wo_finds < 30) {
 
-      strcpy(tmp, cYEL);
+        strcpy(tmp, cYEL);
 
-    } else
+      } else
 
         /* No finds for a long time and no test cases to try. */
         if (afl->cycles_wo_finds > 100 && !afl->pending_not_fuzzed &&
             min_wo_finds > 120) {
 
-      strcpy(tmp, cLGN);
+          strcpy(tmp, cLGN);
 
-      /* Default: cautiously OK to stop? */
+          /* Default: cautiously OK to stop? */
 
-    } else {
+        } else {
 
-      strcpy(tmp, cLBL);
+          strcpy(tmp, cLBL);
 
-    }
+        }
 
   }
 
@@ -1666,8 +1667,9 @@ void show_stats_pizza(afl_state_t *afl) {
   /* Since `total_crashes` does not get reloaded from disk on restart,
     it indicates if we found crashes this round already -> paint red.
     If it's 0, but `saved_crashes` is set from a past run, paint in yellow. */
-  char *crash_color =
-      afl->total_crashes ? cLRD : afl->saved_crashes ? cYEL : cRST;
+  char *crash_color = afl->total_crashes   ? cLRD
+                      : afl->saved_crashes ? cYEL
+                                           : cRST;
 
   /* Lord, forgive me this. */
 
@@ -1690,26 +1692,26 @@ void show_stats_pizza(afl_state_t *afl) {
 
     } else
 
-        /* Subsequent cycles, but we're still making finds. */
-        if (afl->cycles_wo_finds < 25 || min_wo_finds < 30) {
+      /* Subsequent cycles, but we're still making finds. */
+      if (afl->cycles_wo_finds < 25 || min_wo_finds < 30) {
 
-      strcpy(tmp, cYEL);
+        strcpy(tmp, cYEL);
 
-    } else
+      } else
 
         /* No finds for a long time and no test cases to try. */
         if (afl->cycles_wo_finds > 100 && !afl->pending_not_fuzzed &&
             min_wo_finds > 120) {
 
-      strcpy(tmp, cLGN);
+          strcpy(tmp, cLGN);
 
-      /* Default: cautiously OK to stop? */
+          /* Default: cautiously OK to stop? */
 
-    } else {
+        } else {
 
-      strcpy(tmp, cLBL);
+          strcpy(tmp, cLBL);
 
-    }
+        }
 
   }