error handling for custom mutators

6 files changed, 54 insertions, 36 deletions

diff --git a/examples/custom_mutators/custom_mutator_helpers.h b/examples/custom_mutators/custom_mutator_helpers.h
index 0dc00d96..e5ce3569 100644
--- a/examples/custom_mutators/custom_mutator_helpers.h
+++ b/examples/custom_mutators/custom_mutator_helpers.h
@@ -277,37 +277,48 @@ static void surgical_havoc_mutate(u8 *out_buf, s32 begin, s32 end) {
 
 }
 
+
+/* This function calculates the next power of 2 greater or equal its argument.
+ @return The rounded up power of 2 (if no overflow) or 0 on overflow.
+*/
+static inline size_t next_pow2(size_t in) {
+  if (in == 0 || in > (size_t)-1) return 0; /* avoid undefined behaviour under-/overflow */
+  size_t out = in - 1;
+  out |= out >> 1;
+  out |= out >> 2;
+  out |= out >> 4;
+  out |= out >> 8;
+  out |= out >> 16;
+  return out + 1;
+}
+
 /* This function makes sure *size is > size_needed after call.
- It changes buf and size in-place, if needed.
  It will realloc *buf otherwise.
  *size will grow exponentially as per:
  https://blog.mozilla.org/nnethercote/2014/11/04/please-grow-your-buffers-exponentially/
- Will return NULL if size_needed is <1 or *size is negative or malloc Failed.
- @return For convenience, this function returns *buf. NULL on error.
+ Will return NULL and free *buf if size_needed is <1 or realloc failed.
+ @return For convenience, this function returns *buf.
  */
-static inline void *maybe_grow(void **buf, size_t *size, size_t size_needed) {
-
-  /* Oops. found a bug? */
-  if (unlikely(size_needed < 1)) return NULL;
+static inline void *maybe_grow(void **buf, size_t *size,
+                                  size_t size_needed) {
 
   /* No need to realloc */
-  if (likely(*size >= size_needed)) return *buf;
-  if (unlikely(*size < 0)) return NULL;
-  /* No inital size was set */
-  if (*size == 0) *size = INITIAL_GROWTH_SIZE;
-  while (*size < size_needed) {
-
-    *size *= 2;
-    if ((*size) < 0) {
+  if (likely(size_needed && *size >= size_needed)) return *buf;
 
-      /* An overflow occurred. Fall back to size_needed */
-      *size = size_needed;
+  /* No initial size was set */
+  if (size_needed < INITIAL_GROWTH_SIZE) size_needed = INITIAL_GROWTH_SIZE;
 
-    }
+  /* grow exponentially */
+  size_t next_size = next_pow2(size_needed);
 
+  /* handle overflow */
+  if (!next_size) {
+    next_size = size_needed;
   }
 
-  *buf = realloc(*buf, *size);
+  /* alloc */
+  *buf = realloc(*buf, next_size);
+  *size = *buf ? next_size : 0;
 
   return *buf;
 
diff --git a/examples/custom_mutators/example.c b/examples/custom_mutators/example.c
index 9a62d1a7..488ece81 100644
--- a/examples/custom_mutators/example.c
+++ b/examples/custom_mutators/example.c
@@ -100,6 +100,7 @@ size_t afl_custom_fuzz(my_mutator_t *data, uint8_t *buf, size_t buf_size,
   u8 *mutated_out = maybe_grow(BUF_PARAMS(data, fuzz), mutated_size);
   if (!mutated_out) {
 
+    *out_buf = NULL;
     perror("custom mutator allocation (maybe_grow)");
     return 0;            /* afl-fuzz will very likely error out after this. */
 
@@ -189,16 +190,20 @@ size_t afl_custom_pre_save(my_mutator_t *data, uint8_t *buf, size_t buf_size,
  * @param data pointer returned in afl_custom_init for this fuzz case
  * @param buf Buffer containing the test case
  * @param buf_size Size of the test case
- * @return The amount of possible iteration steps to trim the input
+ * @return The amount of possible iteration steps to trim the input.
+ *        negative on error.
  */
-int afl_custom_init_trim(my_mutator_t *data, uint8_t *buf, size_t buf_size) {
+int32_t afl_custom_init_trim(my_mutator_t *data, uint8_t *buf, size_t buf_size) {
 
   // We simply trim once
   data->trimmming_steps = 1;
 
   data->cur_step = 0;
 
-  maybe_grow(BUF_PARAMS(data, trim), buf_size);
+  if (!maybe_grow(BUF_PARAMS(data, trim), buf_size)) {
+    perror("init_trim grow");
+    return -1;
+  }
   memcpy(data->trim_buf, buf, buf_size);
 
   data->trim_size_current = buf_size;
@@ -245,9 +250,9 @@ size_t afl_custom_trim(my_mutator_t *data, uint8_t **out_buf) {
  * @param[in] data pointer returned in afl_custom_init for this fuzz case
  * @param success Indicates if the last trim operation was successful.
  * @return The next trim iteration index (from 0 to the maximum amount of
- *     steps returned in init_trim)
+ *     steps returned in init_trim). negative ret on failure.
  */
-int afl_custom_post_trim(my_mutator_t *data, int success) {
+int32_t afl_custom_post_trim(my_mutator_t *data, int success) {
 
   if (success) {
 
diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index 79878cb6..a265c1a3 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -700,9 +700,10 @@ struct custom_mutator {
    * @param data pointer returned in afl_custom_init for this fuzz case
    * @param buf Buffer containing the test case
    * @param buf_size Size of the test case
-   * @return The amount of possible iteration steps to trim the input
+   * @return The amount of possible iteration steps to trim the input.
+   *        Negative on error.
    */
-  u32 (*afl_custom_init_trim)(void *data, u8 *buf, size_t buf_size);
+  s32 (*afl_custom_init_trim)(void *data, u8 *buf, size_t buf_size);
 
   /**
    * This method is called for each trimming operation. It doesn't have any
@@ -733,9 +734,9 @@ struct custom_mutator {
    * @param data pointer returned in afl_custom_init for this fuzz case
    * @param success Indicates if the last trim operation was successful.
    * @return The next trim iteration index (from 0 to the maximum amount of
-   *     steps returned in init_trim)
+   *     steps returned in init_trim). Negative on error.
    */
-  u32 (*afl_custom_post_trim)(void *data, u8 success);
+  s32 (*afl_custom_post_trim)(void *data, u8 success);
 
   /**
    * Perform a single custom mutation on a given input.
@@ -818,8 +819,8 @@ u8   trim_case_custom(afl_state_t *, struct queue_entry *q, u8 *in_buf);
 void finalize_py_module(void *);
 
 size_t pre_save_py(void *, u8 *, size_t, u8 **);
-u32    init_trim_py(void *, u8 *, size_t);
-u32    post_trim_py(void *, u8);
+s32    init_trim_py(void *, u8 *, size_t);
+s32    post_trim_py(void *, u8);
 size_t trim_py(void *, u8 **);
 size_t havoc_mutation_py(void *, u8 *, size_t, u8 **, size_t);
 u8     havoc_mutation_probability_py(void *);
diff --git a/include/alloc-inl.h b/include/alloc-inl.h
index b8c83db4..91564932 100644
--- a/include/alloc-inl.h
+++ b/include/alloc-inl.h
@@ -767,10 +767,10 @@ static inline void TRK_ck_free(void *ptr, const char *file, const char *func,
 
 #endif                                          /* _WANT_ORIGINAL_AFL_ALLOC */
 
-/* This function calculates the lowest power of 2 greater or equal its argument.
+/* This function calculates the next power of 2 greater or equal its argument.
  @return The rounded up power of 2 (if no overflow) or 0 on overflow.
 */
-static inline size_t powerOf2Ceil(size_t in) {
+static inline size_t next_pow2(size_t in) {
   if (in == 0 || in > (size_t)-1) return 0; /* avoid undefined behaviour under-/overflow */
   size_t out = in - 1;
   out |= out >> 1;
@@ -801,7 +801,7 @@ static inline void *ck_maybe_grow(void **buf, size_t *size,
   if (size_needed < INITIAL_GROWTH_SIZE) size_needed = INITIAL_GROWTH_SIZE;
 
   /* grow exponentially */
-  size_t next_size = powerOf2Ceil(size_needed);
+  size_t next_size = next_pow2(size_needed);
 
   /* handle overflow */
   if (!next_size) {
diff --git a/src/afl-fuzz-mutators.c b/src/afl-fuzz-mutators.c
index f14a57bb..90d7de40 100644
--- a/src/afl-fuzz-mutators.c
+++ b/src/afl-fuzz-mutators.c
@@ -213,7 +213,7 @@ u8 trim_case_custom(afl_state_t *afl, struct queue_entry *q, u8 *in_buf) {
   afl->stage_cur = 0;
   afl->stage_max =
       afl->mutator->afl_custom_init_trim(afl->mutator->data, in_buf, q->len);
-
+  if (unlikely(afl->stage_max) < 0) FATAL("custom_init_trim error ret: %d", afl->stage_max);
   if (afl->not_on_tty && afl->debug)
     SAYF("[Custom Trimming] START: Max %d iterations, %u bytes", afl->stage_max,
          q->len);
@@ -273,6 +273,7 @@ u8 trim_case_custom(afl_state_t *afl, struct queue_entry *q, u8 *in_buf) {
       /* Tell the custom mutator that the trimming was unsuccessful */
       afl->stage_cur =
           afl->mutator->afl_custom_post_trim(afl->mutator->data, 0);
+      if (unlikely(afl->stage_cur < 0)) FATAL("Error ret in custom_post_trim: %d", afl->stage_cur);
       if (afl->not_on_tty && afl->debug)
         SAYF("[Custom Trimming] FAILURE: %d/%d iterations", afl->stage_cur,
              afl->stage_max);
diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
index 76b5ca80..91e5b084 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -388,7 +388,7 @@ size_t pre_save_py(void *py_mutator, u8 *buf, size_t buf_size, u8 **out_buf) {
 
 }
 
-u32 init_trim_py(void *py_mutator, u8 *buf, size_t buf_size) {
+s32 init_trim_py(void *py_mutator, u8 *buf, size_t buf_size) {
 
   PyObject *py_args, *py_value;
 
@@ -426,7 +426,7 @@ u32 init_trim_py(void *py_mutator, u8 *buf, size_t buf_size) {
 
 }
 
-u32 post_trim_py(void *py_mutator, u8 success) {
+s32 post_trim_py(void *py_mutator, u8 success) {
 
   PyObject *py_args, *py_value;