Merge pull request #361 from rish9101/pre_save_format

Add post library API as custom mutator and rename pre_save

6 files changed, 42 insertions, 66 deletions

diff --git a/src/afl-fuzz-cmplog.c b/src/afl-fuzz-cmplog.c
index 9452fa90..4be6a2c8 100644
--- a/src/afl-fuzz-cmplog.c
+++ b/src/afl-fuzz-cmplog.c
@@ -51,17 +51,6 @@ u8 common_fuzz_cmplog_stuff(afl_state_t *afl, u8 *out_buf, u32 len) {
 
   u8 fault;
 
-  if (afl->post_handler) {
-
-    u8 *post_buf = NULL;
-
-    size_t post_len =
-        afl->post_handler(afl->post_data, out_buf, len, &post_buf);
-    if (!post_buf || !post_len) { return 0; }
-    out_buf = post_buf;
-    len = post_len;
-
-  }
 
   write_to_testcase(afl, out_buf, len);
 
diff --git a/src/afl-fuzz-init.c b/src/afl-fuzz-init.c
index 78b92368..aea22f48 100644
--- a/src/afl-fuzz-init.c
+++ b/src/afl-fuzz-init.c
@@ -297,12 +297,16 @@ void setup_post(afl_state_t *afl) {
   dh = dlopen(fn, RTLD_NOW);
   if (!dh) { FATAL("%s", dlerror()); }
 
-  afl->post_handler = dlsym(dh, "afl_postprocess");
-  if (!afl->post_handler) { FATAL("Symbol 'afl_postprocess' not found."); }
-  afl->post_init = dlsym(dh, "afl_postprocess_init");
-  if (!afl->post_init) { FATAL("Symbol 'afl_postprocess_init' not found."); }
-  afl->post_deinit = dlsym(dh, "afl_postprocess_deinit");
-  if (!afl->post_deinit) {
+  struct custom_mutator * mutator;
+  mutator = ck_alloc(sizeof(struct custom_mutator));
+  memset(mutator, 0, sizeof(struct custom_mutator));
+
+  mutator->afl_custom_post_process = dlsym(dh, "afl_postprocess");
+  if (!mutator->afl_custom_post_process) { FATAL("Symbol 'afl_postprocess' not found."); }
+  mutator->afl_custom_init = dlsym(dh, "afl_postprocess_init");
+  if (!mutator->afl_custom_init) { FATAL("Symbol 'afl_postprocess_init' not found."); }
+  mutator->afl_custom_deinit = dlsym(dh, "afl_postprocess_deinit");
+  if (!mutator->afl_custom_post_process) {
 
     FATAL("Symbol 'afl_postprocess_deinit' not found.");
 
@@ -310,16 +314,10 @@ void setup_post(afl_state_t *afl) {
 
   /* Do a quick test. It's better to segfault now than later =) */
 
-  u8 *post_buf = NULL;
-  afl->post_data = afl->post_init(afl);
-  if (!afl->post_data) { FATAL("Could not initialize post handler."); }
-
-  size_t post_len = afl->post_handler(afl->post_data, tbuf, tlen, &post_buf);
-  if (!post_len || !post_buf) {
+  mutator->data = mutator->afl_custom_init(afl, rand_below(afl, 0xFFFFFFFF));
+  if (!mutator->data) { FATAL("Could not initialize post handler."); }
 
-    SAYF("Empty return in test post handler for buf=\"hello\\0\".");
-
-  }
+  afl->post_library_mutator = mutator;
 
   OKF("Postprocessor installed successfully.");
 
diff --git a/src/afl-fuzz-mutators.c b/src/afl-fuzz-mutators.c
index 2c2efc94..acc1b12d 100644
--- a/src/afl-fuzz-mutators.c
+++ b/src/afl-fuzz-mutators.c
@@ -108,6 +108,8 @@ void setup_custom_mutators(afl_state_t *afl) {
 
 #endif
 
+  if (afl->post_library_mutator) list_append(&afl->custom_mutator_list, afl->post_library_mutator);
+
 }
 
 void destroy_custom_mutators(afl_state_t *afl) {
@@ -120,11 +122,11 @@ void destroy_custom_mutators(afl_state_t *afl) {
       if (el->afl_custom_deinit) el->afl_custom_deinit(el->data);
       if (el->dh) dlclose(el->dh);
 
-      if (el->pre_save_buf) {
+      if (el->post_process_buf) {
 
-        ck_free(el->pre_save_buf);
-        el->pre_save_buf = NULL;
-        el->pre_save_size = 0;
+        ck_free(el->post_process_buf);
+        el->post_process_buf = NULL;
+        el->post_process_size = 0;
 
       }
 
@@ -170,10 +172,10 @@ struct custom_mutator *load_custom_mutator(afl_state_t *afl, const char *fn) {
   mutator->afl_custom_deinit = dlsym(dh, "afl_custom_deinit");
   if (!mutator->afl_custom_deinit) FATAL("Symbol 'afl_custom_init' not found.");
 
-  /* "afl_custom_pre_save", optional */
-  mutator->afl_custom_pre_save = dlsym(dh, "afl_custom_pre_save");
-  if (!mutator->afl_custom_pre_save)
-    ACTF("optional symbol 'afl_custom_pre_save' not found.");
+  /* "afl_custom_post_process", optional */
+  mutator->afl_custom_post_process = dlsym(dh, "afl_custom_post_process");
+  if (!mutator->afl_custom_post_process)
+    ACTF("optional symbol 'afl_custom_post_process' not found.");
 
   u8 notrim = 0;
   /* "afl_custom_init_trim", optional */
diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
index 460d6683..2569cdaf 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -142,8 +142,8 @@ static py_mutator_t *init_py_module(afl_state_t *afl, u8 *module_name) {
       py_functions[PY_FUNC_FUZZ] = PyObject_GetAttrString(py_module, "mutate");
     if (!py_functions[PY_FUNC_FUZZ])
       WARNF("fuzz function not found in python module");
-    py_functions[PY_FUNC_PRE_SAVE] =
-        PyObject_GetAttrString(py_module, "pre_save");
+    py_functions[PY_FUNC_POST_PROCESS] =
+        PyObject_GetAttrString(py_module, "post_process");
     py_functions[PY_FUNC_INIT_TRIM] =
         PyObject_GetAttrString(py_module, "init_trim");
     py_functions[PY_FUNC_POST_TRIM] =
@@ -165,9 +165,9 @@ static py_mutator_t *init_py_module(afl_state_t *afl, u8 *module_name) {
 
       if (!py_functions[py_idx] || !PyCallable_Check(py_functions[py_idx])) {
 
-        if (py_idx == PY_FUNC_PRE_SAVE) {
+        if (py_idx == PY_FUNC_POST_PROCESS) {
 
-          // Implenting the pre_save API is optional for now
+          // Implenting the post_process API is optional for now
           if (PyErr_Occurred()) { PyErr_Print(); }
 
         } else if (py_idx >= PY_FUNC_INIT_TRIM && py_idx <= PY_FUNC_TRIM) {
@@ -309,8 +309,8 @@ struct custom_mutator *load_custom_mutator_py(afl_state_t *afl,
   struct custom_mutator *mutator;
 
   mutator = ck_alloc(sizeof(struct custom_mutator));
-  mutator->pre_save_buf = NULL;
-  mutator->pre_save_size = 0;
+  mutator->post_process_buf = NULL;
+  mutator->post_process_size = 0;
 
   mutator->name = module_name;
   ACTF("Loading Python mutator library from '%s'...", module_name);
@@ -330,9 +330,9 @@ struct custom_mutator *load_custom_mutator_py(afl_state_t *afl,
      is quite different from the custom mutator. */
   mutator->afl_custom_fuzz = fuzz_py;
 
-  if (py_functions[PY_FUNC_PRE_SAVE]) {
+  if (py_functions[PY_FUNC_POST_PROCESS]) {
 
-    mutator->afl_custom_pre_save = pre_save_py;
+    mutator->afl_custom_post_process = post_process_py;
 
   }
 
@@ -384,7 +384,7 @@ struct custom_mutator *load_custom_mutator_py(afl_state_t *afl,
 
 }
 
-size_t pre_save_py(void *py_mutator, u8 *buf, size_t buf_size, u8 **out_buf) {
+size_t post_process_py(void *py_mutator, u8 *buf, size_t buf_size, u8 **out_buf) {
 
   size_t        py_out_buf_size;
   PyObject *    py_args, *py_value;
@@ -395,14 +395,14 @@ size_t pre_save_py(void *py_mutator, u8 *buf, size_t buf_size, u8 **out_buf) {
   if (!py_value) {
 
     Py_DECREF(py_args);
-    FATAL("Failed to convert arguments in custom pre_save");
+    FATAL("Failed to convert arguments in custom post_process");
 
   }
 
   PyTuple_SetItem(py_args, 0, py_value);
 
   py_value = PyObject_CallObject(
-      ((py_mutator_t *)py_mutator)->py_functions[PY_FUNC_PRE_SAVE], py_args);
+      ((py_mutator_t *)py_mutator)->py_functions[PY_FUNC_POST_PROCESS], py_args);
 
   Py_DECREF(py_args);
 
@@ -410,18 +410,18 @@ size_t pre_save_py(void *py_mutator, u8 *buf, size_t buf_size, u8 **out_buf) {
 
     py_out_buf_size = PyByteArray_Size(py_value);
 
-    ck_maybe_grow(BUF_PARAMS(pre_save), py_out_buf_size);
+    ck_maybe_grow(BUF_PARAMS(post_process), py_out_buf_size);
 
-    memcpy(py->pre_save_buf, PyByteArray_AsString(py_value), py_out_buf_size);
+    memcpy(py->post_process_buf, PyByteArray_AsString(py_value), py_out_buf_size);
     Py_DECREF(py_value);
 
-    *out_buf = py->pre_save_buf;
+    *out_buf = py->post_process_buf;
     return py_out_buf_size;
 
   } else {
 
     PyErr_Print();
-    FATAL("Python custom mutator: pre_save call failed.");
+    FATAL("Python custom mutator: post_process call failed.");
 
   }
 
diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index 468b5fc6..d9adb792 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -97,10 +97,10 @@ void write_to_testcase(afl_state_t *afl, void *mem, u32 len) {
 
     LIST_FOREACH(&afl->custom_mutator_list, struct custom_mutator, {
 
-      if (el->afl_custom_pre_save) {
+      if (el->afl_custom_post_process) {
 
         new_size =
-            el->afl_custom_pre_save(el->data, new_mem, new_size, &new_buf);
+            el->afl_custom_post_process(el->data, new_mem, new_size, &new_buf);
 
       }
 
@@ -110,7 +110,7 @@ void write_to_testcase(afl_state_t *afl, void *mem, u32 len) {
 
     if (unlikely(!new_buf && (new_size <= 0))) {
 
-      FATAL("Custom_pre_save failed (ret: %lu)", (long unsigned)new_size);
+      FATAL("Custom_post_process failed (ret: %lu)", (long unsigned)new_size);
 
     } else if (likely(new_buf)) {
 
@@ -119,7 +119,7 @@ void write_to_testcase(afl_state_t *afl, void *mem, u32 len) {
 
     } else {
 
-      /* custom mutators do not has a custom_pre_save function */
+      /* custom mutators do not has a custom_post_process function */
       afl_fsrv_write_to_testcase(&afl->fsrv, mem, len);
 
     }
@@ -690,18 +690,6 @@ u8 common_fuzz_stuff(afl_state_t *afl, u8 *out_buf, u32 len) {
 
   u8 fault;
 
-  if (afl->post_handler) {
-
-    u8 *post_buf = NULL;
-
-    size_t post_len =
-        afl->post_handler(afl->post_data, out_buf, len, &post_buf);
-    if (!post_buf || !post_len) { return 0; }
-    out_buf = post_buf;
-    len = post_len;
-
-  }
-
   write_to_testcase(afl, out_buf, len);
 
   fault = fuzz_run_target(afl, &afl->fsrv, afl->fsrv.exec_tmout);
diff --git a/src/afl-fuzz-state.c b/src/afl-fuzz-state.c
index 4f5389e3..316f9bd9 100644
--- a/src/afl-fuzz-state.c
+++ b/src/afl-fuzz-state.c
@@ -371,7 +371,6 @@ void read_afl_environment(afl_state_t *afl, char **envp) {
 
 void afl_state_deinit(afl_state_t *afl) {
 
-  if (afl->post_deinit) { afl->post_deinit(afl->post_data); }
   if (afl->in_place_resume) { ck_free(afl->in_dir); }
   if (afl->sync_id) { ck_free(afl->out_dir); }
   if (afl->pass_stats) { ck_free(afl->pass_stats); }