edited custom mutator pre_save api

1 files changed, 20 insertions, 1 deletions

diff --git a/src/afl-fuzz-mutators.c b/src/afl-fuzz-mutators.c
index 0b0c3394..fac97f8e 100644
--- a/src/afl-fuzz-mutators.c
+++ b/src/afl-fuzz-mutators.c
@@ -80,6 +80,14 @@ void destroy_custom_mutator(afl_state_t *afl) {
 
     if (afl->mutator->dh) dlclose(afl->mutator->dh);
 
+    if (afl->mutator->pre_save_buf) {
+
+      ck_free(afl->mutator->pre_save_buf);
+      afl->mutator->pre_save_buf = NULL;
+      afl->mutator->pre_save_size = 0;
+
+    }
+
     ck_free(afl->mutator);
     afl->mutator = NULL;
 
@@ -91,6 +99,8 @@ void load_custom_mutator(afl_state_t *afl, const char *fn) {
 
   void *dh;
   afl->mutator = ck_alloc(sizeof(struct custom_mutator));
+  afl->mutator->pre_save_buf = NULL;
+  afl->mutator->pre_save_size = 0;
 
   afl->mutator->name = fn;
   ACTF("Loading custom mutator library from '%s'...", fn);
@@ -125,9 +135,18 @@ void load_custom_mutator(afl_state_t *afl, const char *fn) {
 
   /* "afl_custom_pre_save", optional */
   afl->mutator->afl_custom_pre_save = dlsym(dh, "afl_custom_pre_save");
-  if (!afl->mutator->afl_custom_pre_save)
+  if (!afl->mutator->afl_custom_pre_save) {
+
     WARNF("Symbol 'afl_custom_pre_save' not found.");
 
+  } else {
+
+    /* if we have a pre_save hook, prealloc some memory. */
+    afl->mutator->pre_save_buf = ck_alloc(PRE_SAVE_BUF_INIT_SIZE * sizeof(u8));
+    afl->mutator->pre_save_size = PRE_SAVE_BUF_INIT_SIZE;
+
+  }
+
   u8 notrim = 0;
   /* "afl_custom_init_trim", optional */
   afl->mutator->afl_custom_init_trim = dlsym(dh, "afl_custom_init_trim");