error handling for custom mutators

2 files changed, 40 insertions, 24 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) {