Add new APIs for the custom mutator

6 files changed, 242 insertions, 65 deletions

diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index 46bead3a..ae04554a 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -458,29 +458,83 @@ u8* (*post_handler)(u8* buf, u32* len);
 extern u8* cmplog_binary;
 extern s32 cmplog_child_pid, cmplog_forksrv_pid;
 
-/* hooks for the custom mutator function */
-/**
- * Perform custom mutations on a given input
- * @param data Input data to be mutated
- * @param size Size of input data
- * @param mutated_out Buffer to store the mutated input
- * @param max_size Maximum size of the mutated output. The mutation must not
- * produce data larger than max_size.
- * @param seed Seed used for the mutation. The mutation should produce the same
- * output given the same seed.
- * @return Size of the mutated output.
- */
+/* Custom mutators */
+
+struct custom_mutator {
+  const char* name;
+  void* dh;
+
+  /* hooks for the custom mutator function */
+
+  /**
+   * Initialize the custom mutator.
+   *
+   * (Optional)
+   */
+  u32 (*afl_custom_init)(void);
+
+  /**
+   * Perform custom mutations on a given input
+   *
+   * (Required)
+   *
+   * @param[in] data Input data to be mutated
+   * @param[in] size Size of input data
+   * @param[out] mutated_out Buffer to store the mutated input
+   * @param[in] max_size Maximum size of the mutated output. The mutation must not
+   *     produce data larger than max_size.
+   * @param[in] seed Seed used for the mutation. The mutation should produce the
+   *     same output given the same seed.
+   * @return Size of the mutated output.
+   */
+  size_t (*afl_custom_fuzz)(u8* data, size_t size, u8* mutated_out,
+                            size_t max_size, unsigned int seed);
+
+  /**
+   * A post-processing function to use right before AFL writes the test case to
+   * disk in order to execute the target.
+   *
+   * (Optional) If this functionality is not needed, simply don't define this
+   * function.
+   *
+   * @param[in] data Buffer containing the test case to be executed.
+   * @param[in] size Size of the test case.
+   * @param[out] new_data Buffer to store the test case after processing
+   * @return Size of data after processing.
+   */
+  size_t (*afl_custom_pre_save)(u8* data, size_t size, u8** new_data);
+
+  /**
+   * TODO: figure out what `trim` is
+   *
+   * (Optional)
+   */
+  u32 (*afl_custom_init_trim)(u8*, size_t);
+
+  /**
+   * TODO: figure out how `trim` works
+   *
+   * (Optional)
+   *
+   * @param[out] ret (TODO: finish here)
+   * @param[out] ret_len (TODO: finish here)
+   */
+  void (*afl_custom_trim)(u8** ret, size_t* ret_len);
+
+  /**
+   * A post-processing function for the last trim operation.
+   *
+   * (Optional)
+   *
+   * @param success Indicates if the last trim operation was successful.
+   */
+  u32 (*afl_custom_post_trim)(u8 success);
+};
+
+extern struct custom_mutator* mutator;
+
 size_t (*custom_mutator)(u8* data, size_t size, u8* mutated_out,
                          size_t max_size, unsigned int seed);
-/**
- * A post-processing function to use right before AFL writes the test case to
- * disk in order to execute the target. If this functionality is not needed,
- * Simply don't define this function.
- * @param data Buffer containing the test case to be executed.
- * @param size Size of the test case.
- * @param new_data Buffer to store the test case after processing
- * @return Size of data after processing.
- */
 size_t (*pre_save_handler)(u8* data, size_t size, u8** new_data);
 
 /* Interesting values, as per config.h */
@@ -524,9 +578,10 @@ enum {
 
   /* 00 */ PY_FUNC_INIT,
   /* 01 */ PY_FUNC_FUZZ,
-  /* 02 */ PY_FUNC_INIT_TRIM,
-  /* 03 */ PY_FUNC_POST_TRIM,
-  /* 04 */ PY_FUNC_TRIM,
+  /* 02 */ PY_FUNC_PRE_SAVE,
+  /* 03 */ PY_FUNC_INIT_TRIM,
+  /* 04 */ PY_FUNC_POST_TRIM,
+  /* 05 */ PY_FUNC_TRIM,
   PY_FUNC_COUNT
 
 };
@@ -537,11 +592,18 @@ extern PyObject* py_functions[PY_FUNC_COUNT];
 
 /**** Prototypes ****/
 
+/* Custom mutators */
+void setup_custom_mutator(void);
+void destroy_custom_mutator(void);
+void load_custom_mutator(const char*);
+void load_custom_mutator_py(const char*);
+
 /* Python */
 #ifdef USE_PYTHON
 int  init_py();
 void finalize_py();
 void fuzz_py(char*, size_t, char*, size_t, char**, size_t*);
+size_t pre_save_py(u8* data, size_t size, u8** new_data);
 u32  init_trim_py(char*, size_t);
 u32  post_trim_py(char);
 void trim_py(char**, size_t*);
@@ -628,7 +690,6 @@ u8   fuzz_one(char**);
 void bind_to_free_cpu(void);
 #endif
 void   setup_post(void);
-void   setup_custom_mutator(void);
 void   read_testcases(void);
 void   perform_dry_run(char**);
 void   pivot_inputs(void);
diff --git a/src/afl-fuzz-globals.c b/src/afl-fuzz-globals.c
index fc71d29d..8577f875 100644
--- a/src/afl-fuzz-globals.c
+++ b/src/afl-fuzz-globals.c
@@ -255,6 +255,9 @@ u8 *(*post_handler)(u8 *buf, u32 *len);
 u8 *cmplog_binary;
 s32 cmplog_child_pid, cmplog_forksrv_pid;
 
+/* Custom mutator */
+struct custom_mutator* mutator;
+
 /* hooks for the custom mutator function */
 size_t (*custom_mutator)(u8 *data, size_t size, u8 *mutated_out,
                          size_t max_size, unsigned int seed);
diff --git a/src/afl-fuzz-init.c b/src/afl-fuzz-init.c
index 632cdf6b..9ae95b7d 100644
--- a/src/afl-fuzz-init.c
+++ b/src/afl-fuzz-init.c
@@ -296,34 +296,6 @@ void setup_post(void) {
 
 }
 
-void setup_custom_mutator(void) {
-
-  void* dh;
-  u8*   fn = getenv("AFL_CUSTOM_MUTATOR_LIBRARY");
-
-  if (!fn) return;
-
-  if (limit_time_sig)
-    FATAL(
-        "MOpt and custom mutator are mutually exclusive. We accept pull "
-        "requests that integrates MOpt with the optional mutators "
-        "(custom/radamsa/redquenn/...).");
-
-  ACTF("Loading custom mutator library from '%s'...", fn);
-
-  dh = dlopen(fn, RTLD_NOW);
-  if (!dh) FATAL("%s", dlerror());
-
-  custom_mutator = dlsym(dh, "afl_custom_mutator");
-  if (!custom_mutator) FATAL("Symbol 'afl_custom_mutator' not found.");
-
-  pre_save_handler = dlsym(dh, "afl_pre_save_handler");
-  //  if (!pre_save_handler) WARNF("Symbol 'afl_pre_save_handler' not found.");
-
-  OKF("Custom mutator installed successfully.");
-
-}
-
 /* Shuffle an array of pointers. Might be slightly biased. */
 
 static void shuffle_ptrs(void** ptrs, u32 cnt) {
diff --git a/src/afl-fuzz-mutators.c b/src/afl-fuzz-mutators.c
new file mode 100644
index 00000000..0e77a690
--- /dev/null
+++ b/src/afl-fuzz-mutators.c
@@ -0,0 +1,149 @@
+/*
+   american fuzzy lop++ - custom mutators related routines
+   -------------------------------------------------------
+
+   Originally written by Shengtuo Hu
+
+   Now maintained by  Marc Heuse <mh@mh-sec.de>,
+                        Heiko Eißfeldt <heiko.eissfeldt@hexco.de> and
+                        Andrea Fioraldi <andreafioraldi@gmail.com>
+
+   Copyright 2016, 2017 Google Inc. All rights reserved.
+   Copyright 2019-2020 AFLplusplus Project. All rights reserved.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at:
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+   This is the real deal: the program takes an instrumented binary and
+   attempts a variety of basic fuzzing tricks, paying close attention to
+   how they affect the execution path.
+
+ */
+
+#include "afl-fuzz.h"
+
+void setup_custom_mutator(void) {
+
+  u8* fn = getenv("AFL_CUSTOM_MUTATOR_LIBRARY");
+
+  if (fn) {
+    if (limit_time_sig)
+      FATAL(
+          "MOpt and custom mutator are mutually exclusive. We accept pull "
+          "requests that integrates MOpt with the optional mutators "
+          "(custom/radamsa/redquenn/...).");
+
+    load_custom_mutator(fn);
+
+    return;
+  }
+
+#ifdef USE_PYTHON
+  if (init_py()) FATAL("Failed to initialize Python module");
+
+  // u8* module_name = getenv("AFL_PYTHON_MODULE");
+  // if (py_module && module_name)
+  //   load_custom_mutator_py(module_name);
+#else
+  if (getenv("AFL_PYTHON_MODULE"))
+    FATAL("Your AFL binary was built without Python support");
+#endif
+
+}
+
+void destroy_custom_mutator(void) {
+
+  if (mutator) {
+    if (mutator->dh)
+      dlclose(mutator->dh);
+    else {
+      /* Python mutator */
+#ifdef USE_PYTHON
+      finalize_py();
+#endif
+    }
+
+    ck_free(mutator);
+  }
+
+}
+
+void load_custom_mutator(const char* fn) {
+
+  void* dh;
+  mutator = ck_alloc(sizeof(struct custom_mutator));
+
+  mutator->name = fn;
+  ACTF("Loading custom mutator library from '%s'...", fn);
+
+  dh = dlopen(fn, RTLD_NOW);
+  if (!mutator->dh) FATAL("%s", dlerror());
+  mutator->dh = dh;
+
+  /* Mutator */
+  /* "afl_custom_init", optional */
+  mutator->afl_custom_init = dlsym(dh, "afl_custom_init");
+  if (!mutator->afl_custom_init)
+    WARNF("Symbol 'afl_custom_init' not found.");
+
+  /* "afl_custom_fuzz" or "afl_custom_mutator", required */
+  mutator->afl_custom_fuzz = dlsym(dh, "afl_custom_fuzz");
+  if (!mutator->afl_custom_fuzz) {
+    /* Try "afl_custom_mutator" for backward compatibility */
+    WARNF("Symbol 'afl_custom_fuzz' not found. Try 'afl_custom_mutator'.");
+
+    mutator->afl_custom_fuzz = dlsym(dh, "afl_custom_mutator");
+    if (!mutator->afl_custom_fuzz) {
+      FATAL("Symbol 'afl_custom_mutator' not found.");
+    }
+  }
+
+  /* "afl_custom_pre_save", optional */
+  mutator->afl_custom_pre_save = dlsym(dh, "afl_custom_pre_save");
+  if (!mutator->afl_custom_pre_save)
+    WARNF("Symbol 'afl_custom_pre_save' not found.");
+
+  /* "afl_custom_init_trim", optional */
+  mutator->afl_custom_init_trim = dlsym(dh, "afl_custom_init_trim");
+  if (!mutator->afl_custom_init_trim)
+    WARNF("Symbol 'afl_custom_init_trim' not found.");
+
+  /* "afl_custom_trim", optional */
+  mutator->afl_custom_trim = dlsym(dh, "afl_custom_trim");
+  if (!mutator->afl_custom_trim)
+    WARNF("Symbol 'afl_custom_trim' not found.");
+
+  /* "afl_custom_post_trim", optional */
+  mutator->afl_custom_post_trim = dlsym(dh, "afl_custom_post_trim");
+  if (!mutator->afl_custom_post_trim)
+    WARNF("Symbol 'afl_custom_post_trim' not found.");
+
+  OKF("Custom mutator '%s' installed successfully.", fn);
+
+  /* Initialize the custom mutator */
+  if (mutator->afl_custom_init)
+    mutator->afl_custom_init();
+
+}
+
+// void load_custom_mutator_py(const char* module_name) {
+
+//   mutator = ck_alloc(sizeof(struct custom_mutator));
+
+//   mutator->name = module_name;
+//   ACTF("Loading Python mutator library from '%s'...", module_name);
+
+//   /* Initialize of the Python mutator has been invoked in "init_py()" */
+//   mutator->afl_custom_init = NULL;
+//   mutator->afl_custom_fuzz = fuzz_py;
+//   mutator->afl_custom_pre_save = pre_save_py;
+//   mutator->afl_custom_init_trim = init_trim_py;
+//   mutator->afl_custom_trim = trim_py;
+//   mutator->afl_custom_post_trim = post_trim_py;
+
+//   OKF("Python mutator '%s' installed successfully.", module_name);
+
+// }
diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
index 42286527..be86ebba 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -55,6 +55,8 @@ int init_py() {
       u8 py_notrim = 0, py_idx;
       py_functions[PY_FUNC_INIT] = PyObject_GetAttrString(py_module, "init");
       py_functions[PY_FUNC_FUZZ] = PyObject_GetAttrString(py_module, "fuzz");
+      py_functions[PY_FUNC_PRE_SAVE] =
+          PyObject_GetAttrString(py_module, "pre_save");
       py_functions[PY_FUNC_INIT_TRIM] =
           PyObject_GetAttrString(py_module, "init_trim");
       py_functions[PY_FUNC_POST_TRIM] =
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index c7f8ccad..035f74dc 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -155,7 +155,7 @@ static void usage(u8* argv0, int more_help) {
       "LD_BIND_LAZY: do not set LD_BIND_NOW env var for target\n"
       "AFL_BENCH_JUST_ONE: run the target just once\n"
       "AFL_DUMB_FORKSRV: use fork server without feedback from target\n"
-      "AFL_CUSTOM_MUTATOR_LIBRARY: lib with afl_custom_mutator() to mutate inputs\n"
+      "AFL_CUSTOM_MUTATOR_LIBRARY: lib with afl_custom_fuzz() to mutate inputs\n"
       "AFL_CUSTOM_MUTATOR_ONLY: avoid AFL++'s internal mutators\n"
       "AFL_PYTHON_MODULE: mutate and trim inputs with the specified Python module\n"
       "AFL_PYTHON_ONLY: skip AFL++'s own mutators\n"
@@ -864,13 +864,6 @@ int main(int argc, char** argv, char** envp) {
 
   setup_dirs_fds();
 
-#ifdef USE_PYTHON
-  if (init_py()) FATAL("Failed to initialize Python module");
-#else
-  if (getenv("AFL_PYTHON_MODULE"))
-    FATAL("Your AFL binary was built without Python support");
-#endif
-
   setup_cmdline_file(argv + optind);
 
   read_testcases();
@@ -1147,13 +1140,10 @@ stop_fuzzing:
   destroy_extras();
   ck_free(target_path);
   ck_free(sync_id);
+  destroy_custom_mutator();
 
   alloc_report();
 
-#ifdef USE_PYTHON
-  finalize_py();
-#endif
-
   OKF("We're done here. Have a nice day!\n");
 
   exit(0);