retake from mem if possible

2 files changed, 40 insertions, 7 deletions

diff --git a/src/afl-fuzz-queue.c b/src/afl-fuzz-queue.c
index 095a391f..92b722f6 100644
--- a/src/afl-fuzz-queue.c
+++ b/src/afl-fuzz-queue.c
@@ -870,8 +870,10 @@ u32 calculate_score(afl_state_t *afl, struct queue_entry *q) {
 
 }
 
-void queue_testcase_retake(afl_state_t *afl, struct queue_entry *q,
-                           u32 old_len) {
+/* after a custom trim we need to reload the testcase from disk */
+
+inline void queue_testcase_retake(afl_state_t *afl, struct queue_entry *q,
+                                  u32 old_len) {
 
   if (likely(q->testcase_buf)) {
 
@@ -879,9 +881,9 @@ void queue_testcase_retake(afl_state_t *afl, struct queue_entry *q,
 
     if (len != old_len) {
 
-      afl->q_testcase_cache_size =
-          afl->q_testcase_cache_size + q->len - old_len;
+      afl->q_testcase_cache_size = afl->q_testcase_cache_size + len - old_len;
       q->testcase_buf = realloc(q->testcase_buf, len);
+
       if (unlikely(!q->testcase_buf)) {
 
         PFATAL("Unable to malloc '%s' with len %d", q->fname, len);
@@ -901,8 +903,35 @@ void queue_testcase_retake(afl_state_t *afl, struct queue_entry *q,
 
 }
 
+/* after a normal trim we need to replace the testcase with the new data */
+
+inline void queue_testcase_retake_mem(afl_state_t *afl, struct queue_entry *q,
+                                      u8 *in, u32 len, u32 old_len) {
+
+  if (likely(q->testcase_buf)) {
+
+    if (len != old_len) {
+
+      afl->q_testcase_cache_size = afl->q_testcase_cache_size + len - old_len;
+      q->testcase_buf = realloc(q->testcase_buf, len);
+
+      if (unlikely(!q->testcase_buf)) {
+
+        PFATAL("Unable to malloc '%s' with len %d", q->fname, len);
+
+      }
+
+    }
+
+    memcpy(q->testcase_buf, in, len);
+
+  }
+
+}
+
 /* Returns the testcase buf from the file behind this queue entry.
   Increases the refcount. */
+
 inline u8 *queue_testcase_get(afl_state_t *afl, struct queue_entry *q) {
 
   u32 len = q->len;
@@ -913,7 +942,7 @@ inline u8 *queue_testcase_get(afl_state_t *afl, struct queue_entry *q) {
 
     u8 *buf;
 
-    if (q == afl->queue_cur) {
+    if (unlikely(q == afl->queue_cur)) {
 
       buf = afl_realloc((void **)&afl->testcase_buf, len);
 
diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index ab870319..dfd3abfb 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -711,7 +711,11 @@ u8 trim_case(afl_state_t *afl, struct queue_entry *q, u8 *in_buf) {
 
     });
 
-    if (orig_len != q->len) { queue_testcase_retake(afl, q, orig_len); }
+    if (orig_len != q->len || custom_trimmed) {
+
+      queue_testcase_retake(afl, q, orig_len);
+
+    }
 
     if (custom_trimmed) return trimmed_case;
 
@@ -846,7 +850,7 @@ u8 trim_case(afl_state_t *afl, struct queue_entry *q, u8 *in_buf) {
 
     close(fd);
 
-    if (orig_len != q->len) queue_testcase_retake(afl, q, orig_len);
+    queue_testcase_retake_mem(afl, q, in_buf, q->len, orig_len);
 
     memcpy(afl->fsrv.trace_bits, afl->clean_trace, afl->fsrv.map_size);
     update_bitmap_score(afl, q);