Merge pull request #1915 from yangzao/dev

add custom mutator function for running script after target gets executed

custom_mutators/examples/custom_post_run.c

@@ -0,0 +1,53 @@
+//
+// This is an example on how to use afl_custom_post_run
+// It executes custom code each time after AFL++ executes the target
+//
+// cc -O3 -fPIC -shared -g -o custom_post_run.so -I../../include custom_post_run.c
+// cd ../..
+// afl-cc -o test-instr test-instr.c
+// AFL_CUSTOM_MUTATOR_LIBRARY=custom_mutators/examples/custom_post_run.so \
+//   afl-fuzz -i in -o out -- ./test-instr -f /tmp/foo
+//
+
+
+#include "afl-fuzz.h"
+
+#include <stdint.h>
+#include <stdlib.h>
+#include <string.h>
+#include <stdio.h>
+
+typedef struct my_mutator {
+
+  afl_state_t *afl;
+
+} my_mutator_t;
+
+my_mutator_t *afl_custom_init(afl_state_t *afl, unsigned int seed) {
+
+  my_mutator_t *data = calloc(1, sizeof(my_mutator_t));
+  if (!data) {
+
+    perror("afl_custom_init alloc");
+    return NULL;
+
+  }
+
+  data->afl = afl;
+
+  return data;
+
+}
+
+void afl_custom_post_run(my_mutator_t *data) {
+
+  printf("hello from afl_custom_post_run\n");
+  return;
+}
+
+
+void afl_custom_deinit(my_mutator_t *data) {
+
+  free(data);
+
+}

custom_mutators/examples/example.py

@@ -133,6 +133,11 @@ def fuzz(buf, add_buf, max_size):
 #     @return: The buffer containing the test case after
 #     '''
 #     return buf
+# def post_run():
+#     '''
+#     Called after each time the execution of the target program by AFL++
+#     '''
+#     pass
 #
 # def havoc_mutation(buf, max_size):
 #     '''

include/afl-fuzz.h

@@ -345,6 +345,7 @@ enum {
   /* 13 */ PY_FUNC_DESCRIBE,
   /* 14 */ PY_FUNC_FUZZ_SEND,
   /* 15 */ PY_FUNC_SPLICE_OPTOUT,
+  /* 16 */ PY_FUNC_POST_RUN,
   PY_FUNC_COUNT
 
 };
@@ -1020,6 +1021,16 @@ struct custom_mutator {
    */
   void (*afl_custom_fuzz_send)(void *data, const u8 *buf, size_t buf_size);
 
+  /**
+   * This method can be used if you want to run some code or scripts each time
+   * AFL++ executes the target with afl-fuzz.
+   *
+   * (Optional)
+   *
+   * @param data pointer returned in afl_custom_init by this custom mutator
+   */
+  void (*afl_custom_post_run)(void *data);
+
   /**
    * Allow for additional analysis (e.g. calling a different tool that does a
    * different kind of coverage and saves this for the custom mutator).
@@ -1075,6 +1086,7 @@ void                   finalize_py_module(void *);
 
 u32         fuzz_count_py(void *, const u8 *, size_t);
 void        fuzz_send_py(void *, const u8 *, size_t);
+void        post_run_py(void *);
 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-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) {

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;

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) {
 

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;
 
 }
-

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 =
+  char *crash_color = afl->total_crashes   ? cLRD
-      afl->total_crashes ? cLRD : afl->saved_crashes ? cYEL : cRST;
+                      : afl->saved_crashes ? cYEL
+                                           : cRST;
 
   /* Lord, forgive me this. */
 
@@ -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 =
+  char *crash_color = afl->total_crashes   ? cLRD
-      afl->total_crashes ? cLRD : afl->saved_crashes ? cYEL : cRST;
+                      : afl->saved_crashes ? cYEL
+                                           : cRST;
 
   /* Lord, forgive me this. */