tiny changes in custom mut api

6 files changed, 39 insertions, 25 deletions

diff --git a/examples/custom_mutators/example.c b/examples/custom_mutators/example.c
index 54fc9d47..9a62d1a7 100644
--- a/examples/custom_mutators/example.c
+++ b/examples/custom_mutators/example.c
@@ -29,16 +29,15 @@ typedef struct my_mutator {
   afl_t *afl;
 
   // any additional data here!
-  uint8_t *trim_buf;
-  size_t   trim_buf_size;
-  int      trimmming_steps;
-  int      cur_step;
+  size_t trim_size_current;
+  int    trimmming_steps;
+  int    cur_step;
 
   // Reused buffers:
   BUF_VAR(u8, fuzz);
   BUF_VAR(u8, data);
   BUF_VAR(u8, havoc);
-  BUF_VAR(u8, trim_out);
+  BUF_VAR(u8, trim);
   BUF_VAR(u8, pre_save);
 
 } my_mutator_t;
@@ -52,7 +51,7 @@ typedef struct my_mutator {
  * in the same way.
  * @return Pointer to the data object this custom mutator instance should use.
  *         There may be multiple instances of this mutator in one afl-fuzz run!
- *         Returns NULL on error.
+ *         Return NULL on error.
  */
 my_mutator_t *afl_custom_init(afl_t *afl, unsigned int seed) {
 
@@ -80,11 +79,13 @@ my_mutator_t *afl_custom_init(afl_t *afl, unsigned int seed) {
  * @param[in] data pointer returned in afl_custom_init for this fuzz case
  * @param[in] buf Pointer to input data to be mutated
  * @param[in] buf_size Size of input data
+ * @param[out] out_buf the buffer we will work on. we can reuse *buf. NULL on
+ * error.
  * @param[in] add_buf Buffer containing the additional test case
  * @param[in] add_buf_size Size of the additional test case
  * @param[in] max_size Maximum size of the mutated output. The mutation must not
  *     produce data larger than max_size.
- * @return Size of the mutated output. Negative return will abort fuzzing.
+ * @return Size of the mutated output.
  */
 size_t afl_custom_fuzz(my_mutator_t *data, uint8_t *buf, size_t buf_size,
                        u8 **out_buf, uint8_t *add_buf,
@@ -100,7 +101,7 @@ size_t afl_custom_fuzz(my_mutator_t *data, uint8_t *buf, size_t buf_size,
   if (!mutated_out) {
 
     perror("custom mutator allocation (maybe_grow)");
-    return -1;           /* afl-fuzz will very likely error out after this. */
+    return 0;            /* afl-fuzz will very likely error out after this. */
 
   }
 
@@ -135,7 +136,7 @@ size_t afl_custom_fuzz(my_mutator_t *data, uint8_t *buf, size_t buf_size,
  *     processing. External library should allocate memory for out_buf.
  *     The buf pointer may be reused (up to the given buf_size);
  * @return Size of the output buffer after processing or the needed amount.
- *     A return smaller 1 indicates an error.
+ *     A return of 0 indicates an error.
  */
 size_t afl_custom_pre_save(my_mutator_t *data, uint8_t *buf, size_t buf_size,
                            uint8_t **out_buf) {
@@ -146,7 +147,8 @@ size_t afl_custom_pre_save(my_mutator_t *data, uint8_t *buf, size_t buf_size,
     if (!data->pre_save_buf) {
 
       perror("custom mutator realloc failed.");
-      return -1;
+      *out_buf = NULL;
+      return 0;
 
     }
 
@@ -195,8 +197,11 @@ int afl_custom_init_trim(my_mutator_t *data, uint8_t *buf, size_t buf_size) {
   data->trimmming_steps = 1;
 
   data->cur_step = 0;
-  data->trim_buf = buf;
-  data->trim_buf_size = buf_size;
+
+  maybe_grow(BUF_PARAMS(data, trim), buf_size);
+  memcpy(data->trim_buf, buf, buf_size);
+
+  data->trim_size_current = buf_size;
 
   return data->trimmming_steps;
 
@@ -218,15 +223,15 @@ int afl_custom_init_trim(my_mutator_t *data, uint8_t *buf, size_t buf_size) {
  *     External library should allocate memory for out_buf.
  *     AFL++ will not release the memory after saving the test case.
  *     Keep a ref in *data.
+ *     *out_buf = NULL is treated as error.
  * @return Pointer to the size of the trimmed test case
  */
 size_t afl_custom_trim(my_mutator_t *data, uint8_t **out_buf) {
 
-  size_t ret = data->trim_buf_size - 1;
+  *out_buf = data->trim_buf;
 
-  *out_buf = maybe_grow(BUF_PARAMS(data, trim_out), ret);
   // Remove the last byte of the trimming input
-  memcpy(*out_buf, data->trim_buf, ret);
+  return data->trim_size_current - 1;
 
 }
 
@@ -266,7 +271,7 @@ int afl_custom_post_trim(my_mutator_t *data, int success) {
  *     output
  * @param[in] buf_size Size of input data
  * @param[out] out_buf The output buffer. buf can be reused, if the content
- * fits.
+ * fits. *out_buf = NULL is treated as error.
  * @param[in] max_size Maximum size of the mutated output. The mutation must
  *     not produce data larger than max_size.
  * @return Size of the mutated output.
@@ -277,6 +282,13 @@ size_t afl_custom_havoc_mutation(my_mutator_t *data, u8 *buf, size_t buf_size,
   if (buf_size == 0) {
 
     *out_buf = maybe_grow(BUF_PARAMS(data, havoc), 1);
+    if (!*out_buf) {
+
+      perror("custom havoc: maybe_grow");
+      return 0;
+
+    }
+
     **out_buf = rand() % 256;
     buf_size = 1;
 
@@ -354,7 +366,7 @@ void afl_custom_deinit(my_mutator_t *data) {
   free(data->havoc_buf);
   free(data->data_buf);
   free(data->fuzz_buf);
-  free(data->trim_out_buf);
+  free(data->trim_buf);
   free(data);
 
 }
diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index 5f9891bc..79878cb6 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -649,14 +649,16 @@ struct custom_mutator {
    * (Optional for now. Required in the future)
    *
    * @param data pointer returned in afl_custom_init for this fuzz case
-   * @param[inout] buf Pointer to the input data to be mutated and the mutated
+   * @param[in] buf Pointer to the input data to be mutated and the mutated
    *     output
    * @param[in] buf_size Size of the input/output data
+   * @param[out] out_buf the new buffer. We may reuse *buf if large enough.
+   *             *out_buf = NULL is treated as FATAL.
    * @param[in] add_buf Buffer containing the additional test case
    * @param[in] add_buf_size Size of the additional test case
    * @param[in] max_size Maximum size of the mutated output. The mutation must
    * not produce data larger than max_size.
-   * @return Size of the mutated output. Negative on error will abort exeuction.
+   * @return Size of the mutated output.
    */
   size_t (*afl_custom_fuzz)(void *data, u8 *buf, size_t buf_size, u8 **out_buf,
                             u8 *add_buf, size_t add_buf_size, size_t max_size);
diff --git a/src/afl-fuzz-mutators.c b/src/afl-fuzz-mutators.c
index 0692ebb0..f14a57bb 100644
--- a/src/afl-fuzz-mutators.c
+++ b/src/afl-fuzz-mutators.c
@@ -229,7 +229,7 @@ u8 trim_case_custom(afl_state_t *afl, struct queue_entry *q, u8 *in_buf) {
 
     size_t retlen = afl->mutator->afl_custom_trim(afl->mutator->data, &retbuf);
 
-    if (unlikely(retlen < 0 || !retbuf))
+    if (unlikely(!retbuf))
       FATAL("custom_trim failed (ret %zd)", retlen);
     else if (unlikely(retlen > orig_len))
       FATAL(
diff --git a/src/afl-fuzz-one.c b/src/afl-fuzz-one.c
index 8dfafb7b..b1bbad0a 100644
--- a/src/afl-fuzz-one.c
+++ b/src/afl-fuzz-one.c
@@ -1618,8 +1618,8 @@ custom_mutator_stage:
         afl->mutator->data, out_buf, len, &mutated_buf, new_buf, target->len,
         max_seed_size);
 
-    if (unlikely(mutated_size < 0))
-      FATAL("custom_fuzz returned %zd", mutated_size);
+    if (unlikely(!mutated_buf))
+      FATAL("Error in custom_fuzz. Size returned: %zd", mutated_size);
 
     if (mutated_size > len) afl->out_size = mutated_size;
 
@@ -1734,7 +1734,7 @@ havoc_stage:
         u8 *   custom_havoc_buf = NULL;
         size_t new_len = afl->mutator->afl_custom_havoc_mutation(
             afl->mutator->data, out_buf, temp_len, &custom_havoc_buf, MAX_FILE);
-        if (unlikely(new_len < 0))
+        if (unlikely(!custom_havoc_buf))
           FATAL("Error in custom_havoc (return %zd)", new_len);
         if (likely(new_len > 0 && custom_havoc_buf)) {
 
diff --git a/src/afl-fuzz-python.c b/src/afl-fuzz-python.c
index 6f8982c0..76b5ca80 100644
--- a/src/afl-fuzz-python.c
+++ b/src/afl-fuzz-python.c
@@ -103,7 +103,7 @@ size_t fuzz_py(void *py_mutator, u8 *buf, size_t buf_size, u8 **out_buf,
   } else {
 
     PyErr_Print();
-    FATAL("Call failed");
+    FATAL("python custom fuzz: call failed");
 
   }
 
diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index 315539d1..90cb2ed5 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -221,7 +221,7 @@ void write_to_testcase(afl_state_t *afl, void *mem, u32 len) {
     size_t new_size = afl->mutator->afl_custom_pre_save(afl->mutator->data, mem,
                                                         len, &new_buf);
 
-    if (unlikely(new_size <= 0 || !new_buf))
+    if (unlikely(!new_buf))
       FATAL("Custom_pre_save failed (ret: %lu)", (long unsigned)new_size);
 
     /* everything as planned. use the new data. */