1 files changed, 184 insertions, 23 deletions

diff --git a/src/afl-fuzz-mutators.c b/src/afl-fuzz-mutators.c
index 0e77a690..22ca1384 100644
--- a/src/afl-fuzz-mutators.c
+++ b/src/afl-fuzz-mutators.c
@@ -42,11 +42,22 @@ void setup_custom_mutator(void) {
   }
 
 #ifdef USE_PYTHON
-  if (init_py()) FATAL("Failed to initialize Python module");
+  u8* module_name = getenv("AFL_PYTHON_MODULE");
 
-  // u8* module_name = getenv("AFL_PYTHON_MODULE");
-  // if (py_module && module_name)
-  //   load_custom_mutator_py(module_name);
+  if (module_name) {
+
+    if (limit_time_sig)
+      FATAL(
+          "MOpt and Python mutator are mutually exclusive. We accept pull "
+          "requests that integrates MOpt with the optional mutators "
+          "(custom/radamsa/redquenn/...).");
+
+    if (init_py_module(module_name))
+      FATAL("Failed to initialize Python module");
+
+    load_custom_mutator_py(module_name);
+
+  }
 #else
   if (getenv("AFL_PYTHON_MODULE"))
     FATAL("Your AFL binary was built without Python support");
@@ -62,7 +73,7 @@ void destroy_custom_mutator(void) {
     else {
       /* Python mutator */
 #ifdef USE_PYTHON
-      finalize_py();
+      finalize_py_module();
 #endif
     }
 
@@ -80,11 +91,11 @@ void load_custom_mutator(const char* fn) {
   ACTF("Loading custom mutator library from '%s'...", fn);
 
   dh = dlopen(fn, RTLD_NOW);
-  if (!mutator->dh) FATAL("%s", dlerror());
+  if (!dh) FATAL("%s", dlerror());
   mutator->dh = dh;
 
   /* Mutator */
-  /* "afl_custom_init", optional */
+  /* "afl_custom_init", optional for backward compatibility */
   mutator->afl_custom_init = dlsym(dh, "afl_custom_init");
   if (!mutator->afl_custom_init)
     WARNF("Symbol 'afl_custom_init' not found.");
@@ -92,13 +103,14 @@ void load_custom_mutator(const char* fn) {
   /* "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) {
+    if (!mutator->afl_custom_fuzz)
       FATAL("Symbol 'afl_custom_mutator' not found.");
-    }
+
   }
 
   /* "afl_custom_pre_save", optional */
@@ -106,6 +118,7 @@ void load_custom_mutator(const char* fn) {
   if (!mutator->afl_custom_pre_save)
     WARNF("Symbol 'afl_custom_pre_save' not found.");
 
+  u8 notrim = 0;
   /* "afl_custom_init_trim", optional */
   mutator->afl_custom_init_trim = dlsym(dh, "afl_custom_init_trim");
   if (!mutator->afl_custom_init_trim)
@@ -121,29 +134,177 @@ void load_custom_mutator(const char* fn) {
   if (!mutator->afl_custom_post_trim)
     WARNF("Symbol 'afl_custom_post_trim' not found.");
 
+  if (notrim) {
+
+    mutator->afl_custom_init_trim = NULL;
+    mutator->afl_custom_trim = NULL;
+    mutator->afl_custom_post_trim = NULL;
+    WARNF(
+        "Custom mutator does not implement all three trim APIs, standard "
+        "trimming will be used.");
+
+  }
+
   OKF("Custom mutator '%s' installed successfully.", fn);
 
   /* Initialize the custom mutator */
   if (mutator->afl_custom_init)
-    mutator->afl_custom_init();
+    mutator->afl_custom_init(UR(0xFFFFFFFF));
+
+}
+
+u8 trim_case_custom(char** argv, struct queue_entry* q, u8* in_buf) {
+
+  static u8 tmp[64];
+  static u8 clean_trace[MAP_SIZE];
+
+  u8  needs_write = 0, fault = 0;
+  u32 trim_exec = 0;
+  u32 orig_len = q->len;
+
+  stage_name = tmp;
+  bytes_trim_in += q->len;
+
+  /* Initialize trimming in the custom mutator */
+  stage_cur = 0;
+  stage_max = mutator->afl_custom_init_trim(in_buf, q->len);
+
+  if (not_on_tty && debug)
+    SAYF("[Custom Trimming] START: Max %d iterations, %u bytes", stage_max,
+         q->len);
+
+  while (stage_cur < stage_max) {
+
+    sprintf(tmp, "ptrim %s", DI(trim_exec));
+
+    u32 cksum;
+
+    u8*    retbuf = NULL;
+    size_t retlen = 0;
+
+    mutator->afl_custom_trim(&retbuf, &retlen);
+
+    if (retlen > orig_len)
+      FATAL(
+          "Trimmed data returned by custom mutator is larger than original "
+          "data");
+
+    write_to_testcase(retbuf, retlen);
+
+    fault = run_target(argv, exec_tmout);
+    ++trim_execs;
+
+    if (stop_soon || fault == FAULT_ERROR) {
+
+      free(retbuf);
+      goto abort_trimming;
+
+    }
+
+    cksum = hash32(trace_bits, MAP_SIZE, HASH_CONST);
+
+    if (cksum == q->exec_cksum) {
+
+      q->len = retlen;
+      memcpy(in_buf, retbuf, retlen);
+
+      /* Let's save a clean trace, which will be needed by
+         update_bitmap_score once we're done with the trimming stuff. */
+
+      if (!needs_write) {
+
+        needs_write = 1;
+        memcpy(clean_trace, trace_bits, MAP_SIZE);
+
+      }
+
+      /* Tell the custom mutator that the trimming was successful */
+      stage_cur = mutator->afl_custom_post_trim(1);
+
+      if (not_on_tty && debug)
+        SAYF("[Custom Trimming] SUCCESS: %d/%d iterations (now at %u bytes)",
+             stage_cur, stage_max, q->len);
+
+    } else {
+
+      /* Tell the custom mutator that the trimming was unsuccessful */
+      stage_cur = mutator->afl_custom_post_trim(0);
+      if (not_on_tty && debug)
+        SAYF("[Custom Trimming] FAILURE: %d/%d iterations", stage_cur,
+             stage_max);
+
+    }
+
+    free(retbuf);
+
+    /* Since this can be slow, update the screen every now and then. */
+
+    if (!(trim_exec++ % stats_update_freq)) show_stats();
+
+  }
+
+  if (not_on_tty && debug)
+    SAYF("[Custom Trimming] DONE: %u bytes -> %u bytes", orig_len, q->len);
+
+  /* If we have made changes to in_buf, we also need to update the on-disk
+     version of the test case. */
+
+  if (needs_write) {
+
+    s32 fd;
+
+    unlink(q->fname);                                      /* ignore errors */
+
+    fd = open(q->fname, O_WRONLY | O_CREAT | O_EXCL, 0600);
+
+    if (fd < 0) PFATAL("Unable to create '%s'", q->fname);
+
+    ck_write(fd, in_buf, q->len, q->fname);
+    close(fd);
+
+    memcpy(trace_bits, clean_trace, MAP_SIZE);
+    update_bitmap_score(q);
+
+  }
+
+abort_trimming:
+
+  bytes_trim_out += q->len;
+  return fault;
 
 }
 
-// void load_custom_mutator_py(const char* module_name) {
+void load_custom_mutator_py(const char* module_name) {
 
-//   mutator = ck_alloc(sizeof(struct custom_mutator));
+  mutator = ck_alloc(sizeof(struct custom_mutator));
 
-//   mutator->name = module_name;
-//   ACTF("Loading Python mutator library from '%s'...", module_name);
+  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;
+  /* TODO: unify "init" and "fuzz" */
+  if (py_functions[PY_FUNC_INIT])
+    mutator->afl_custom_init = init_py;
 
-//   OKF("Python mutator '%s' installed successfully.", module_name);
+  /* "afl_custom_fuzz" should not be NULL, but the interface of Python mutator
+     is quite different from the custom mutator. */
+  mutator->afl_custom_fuzz = NULL;
 
-// }
+  if (py_functions[PY_FUNC_PRE_SAVE])
+    mutator->afl_custom_pre_save = pre_save_py;
+
+  if (py_functions[PY_FUNC_INIT_TRIM])
+    mutator->afl_custom_init_trim = init_trim_py;
+
+  if (py_functions[PY_FUNC_POST_TRIM])
+    mutator->afl_custom_post_trim = post_trim_py;
+
+  if (py_functions[PY_FUNC_TRIM])
+    mutator->afl_custom_trim = trim_py;
+
+  OKF("Python mutator '%s' installed successfully.", module_name);
+
+  /* Initialize the custom mutator */
+  if (mutator->afl_custom_init)
+    mutator->afl_custom_init(UR(0xFFFFFFFF));
+
+}