fix token usage in normal and MOpt mode

2 files changed, 243 insertions, 123 deletions

diff --git a/src/afl-fuzz-one.c b/src/afl-fuzz-one.c
index b28ee80a..38eda769 100644
--- a/src/afl-fuzz-one.c
+++ b/src/afl-fuzz-one.c
@@ -1779,6 +1779,62 @@ skip_user_extras:
   afl->stage_finds[STAGE_EXTRAS_AO] += new_hit_cnt - orig_hit_cnt;
   afl->stage_cycles[STAGE_EXTRAS_AO] += afl->stage_max;
 
+  /* Insertion of auto extras. */
+
+  afl->stage_name = "auto extras (insert)";
+  afl->stage_short = "ext_AI";
+  afl->stage_cur = 0;
+  afl->stage_max = afl->a_extras_cnt * (len + 1);
+
+  orig_hit_cnt = new_hit_cnt;
+
+  ex_tmp = afl_realloc(AFL_BUF_PARAM(ex), len + MAX_DICT_FILE);
+  if (unlikely(!ex_tmp)) { PFATAL("alloc"); }
+
+  for (i = 0; i <= (u32)len; ++i) {
+
+    afl->stage_cur_byte = i;
+
+    for (j = 0; j < afl->a_extras_cnt; ++j) {
+
+      if (len + afl->a_extras[j].len > MAX_FILE) {
+
+        --afl->stage_max;
+        continue;
+
+      }
+
+      /* Insert token */
+      memcpy(ex_tmp + i, afl->a_extras[j].data, afl->a_extras[j].len);
+
+      /* Copy tail */
+      memcpy(ex_tmp + i + afl->a_extras[j].len, out_buf + i, len - i);
+
+#ifdef INTROSPECTION
+      snprintf(afl->mutation, sizeof(afl->mutation),
+               "%s AUTO_EXTRAS_insert-%u-%u", afl->queue_cur->fname, i, j);
+#endif
+
+      if (common_fuzz_stuff(afl, ex_tmp, len + afl->a_extras[j].len)) {
+
+        goto abandon_entry;
+
+      }
+
+      ++afl->stage_cur;
+
+    }
+
+    /* Copy head */
+    ex_tmp[i] = out_buf[i];
+
+  }
+
+  new_hit_cnt = afl->queued_items + afl->saved_crashes;
+
+  afl->stage_finds[STAGE_EXTRAS_AI] += new_hit_cnt - orig_hit_cnt;
+  afl->stage_cycles[STAGE_EXTRAS_AI] += afl->stage_max;
+
 skip_extras:
 
   /* If we made this to here without jumping to havoc_stage or abandon_entry,
@@ -4367,6 +4423,62 @@ skip_user_extras:
   afl->stage_finds[STAGE_EXTRAS_AO] += new_hit_cnt - orig_hit_cnt;
   afl->stage_cycles[STAGE_EXTRAS_AO] += afl->stage_max;
 
+  /* Insertion of auto extras. */
+
+  afl->stage_name = "auto extras (insert)";
+  afl->stage_short = "ext_AI";
+  afl->stage_cur = 0;
+  afl->stage_max = afl->a_extras_cnt * (len + 1);
+
+  orig_hit_cnt = new_hit_cnt;
+
+  ex_tmp = afl_realloc(AFL_BUF_PARAM(ex), len + MAX_DICT_FILE);
+  if (unlikely(!ex_tmp)) { PFATAL("alloc"); }
+
+  for (i = 0; i <= (u32)len; ++i) {
+
+    afl->stage_cur_byte = i;
+
+    for (j = 0; j < afl->a_extras_cnt; ++j) {
+
+      if (len + afl->a_extras[j].len > MAX_FILE) {
+
+        --afl->stage_max;
+        continue;
+
+      }
+
+      /* Insert token */
+      memcpy(ex_tmp + i, afl->a_extras[j].data, afl->a_extras[j].len);
+
+      /* Copy tail */
+      memcpy(ex_tmp + i + afl->a_extras[j].len, out_buf + i, len - i);
+
+#ifdef INTROSPECTION
+      snprintf(afl->mutation, sizeof(afl->mutation),
+               "%s MOPT_AUTO_EXTRAS_insert-%u-%u", afl->queue_cur->fname, i, j);
+#endif
+
+      if (common_fuzz_stuff(afl, ex_tmp, len + afl->a_extras[j].len)) {
+
+        goto abandon_entry;
+
+      }
+
+      ++afl->stage_cur;
+
+    }
+
+    /* Copy head */
+    ex_tmp[i] = out_buf[i];
+
+  }                                                  /* for i = 0; i <= len */
+
+  new_hit_cnt = afl->queued_items + afl->saved_crashes;
+
+  afl->stage_finds[STAGE_EXTRAS_AI] += new_hit_cnt - orig_hit_cnt;
+  afl->stage_cycles[STAGE_EXTRAS_AI] += afl->stage_max;
+
 skip_extras:
 
   /* If we made this to here without jumping to havoc_stage or abandon_entry,
@@ -4464,14 +4576,14 @@ pacemaker_fuzzing:
 
       havoc_queued = afl->queued_items;
 
-      u32 r_max;
+      u32 r_max, r;
 
       r_max = 15 + ((afl->extras_cnt + afl->a_extras_cnt) ? 2 : 0);
 
       if (unlikely(afl->expand_havoc && afl->ready_for_splicing_count > 1)) {
 
         /* add expensive havoc cases here, they are activated after a full
-           cycle without finds happened */
+           cycle without any finds happened */
 
         ++r_max;
 
@@ -4497,7 +4609,7 @@ pacemaker_fuzzing:
 
         for (i = 0; i < use_stacking; ++i) {
 
-          switch (select_algorithm(afl, r_max)) {
+          switch (r = (select_algorithm(afl, r_max))) {
 
             case 0:
               /* Flip a single bit somewhere. Spooky! */
@@ -4914,194 +5026,198 @@ pacemaker_fuzzing:
 
             }                                                    /* case 15 */
 
+            default: {
+
+            }
+
               /* Values 16 and 17 can be selected only if there are any extras
                  present in the dictionaries. */
 
-            case 16: {
+              r -= 16;
 
-              /* Overwrite bytes with an extra. */
+              if (r == 0 && (afl->extras_cnt || afl->a_extras_cnt)) {
 
-              if (!afl->extras_cnt ||
-                  (afl->a_extras_cnt && rand_below(afl, 2))) {
+                /* Overwrite bytes with an extra. */
 
-                /* No user-specified extras or odds in our favor. Let's use an
-                  auto-detected one. */
+                if (!afl->extras_cnt ||
+                    (afl->a_extras_cnt && rand_below(afl, 2))) {
 
-                u32 use_extra = rand_below(afl, afl->a_extras_cnt);
-                u32 extra_len = afl->a_extras[use_extra].len;
+                  /* No user-specified extras or odds in our favor. Let's use an
+                    auto-detected one. */
 
-                if (extra_len > (u32)temp_len) break;
+                  u32 use_extra = rand_below(afl, afl->a_extras_cnt);
+                  u32 extra_len = afl->a_extras[use_extra].len;
 
-                u32 insert_at = rand_below(afl, temp_len - extra_len + 1);
+                  if (extra_len > (u32)temp_len) break;
+
+                  u32 insert_at = rand_below(afl, temp_len - extra_len + 1);
 #ifdef INTROSPECTION
-                snprintf(afl->m_tmp, sizeof(afl->m_tmp),
-                         " AUTO_EXTRA_OVERWRITE-%u-%u", insert_at, extra_len);
-                strcat(afl->mutation, afl->m_tmp);
+                  snprintf(afl->m_tmp, sizeof(afl->m_tmp),
+                           " AUTO_EXTRA_OVERWRITE-%u-%u", insert_at, extra_len);
+                  strcat(afl->mutation, afl->m_tmp);
 #endif
-                memcpy(out_buf + insert_at, afl->a_extras[use_extra].data,
-                       extra_len);
+                  memcpy(out_buf + insert_at, afl->a_extras[use_extra].data,
+                         extra_len);
 
-              } else {
+                } else {
 
-                /* No auto extras or odds in our favor. Use the dictionary. */
+                  /* No auto extras or odds in our favor. Use the dictionary. */
 
-                u32 use_extra = rand_below(afl, afl->extras_cnt);
-                u32 extra_len = afl->extras[use_extra].len;
+                  u32 use_extra = rand_below(afl, afl->extras_cnt);
+                  u32 extra_len = afl->extras[use_extra].len;
 
-                if (extra_len > (u32)temp_len) break;
+                  if (extra_len > (u32)temp_len) break;
 
-                u32 insert_at = rand_below(afl, temp_len - extra_len + 1);
+                  u32 insert_at = rand_below(afl, temp_len - extra_len + 1);
 #ifdef INTROSPECTION
-                snprintf(afl->m_tmp, sizeof(afl->m_tmp),
-                         " EXTRA_OVERWRITE-%u-%u", insert_at, extra_len);
-                strcat(afl->mutation, afl->m_tmp);
+                  snprintf(afl->m_tmp, sizeof(afl->m_tmp),
+                           " EXTRA_OVERWRITE-%u-%u", insert_at, extra_len);
+                  strcat(afl->mutation, afl->m_tmp);
 #endif
-                memcpy(out_buf + insert_at, afl->extras[use_extra].data,
-                       extra_len);
+                  memcpy(out_buf + insert_at, afl->extras[use_extra].data,
+                         extra_len);
 
-              }
+                }
 
-              MOpt_globals.cycles_v2[STAGE_OverWriteExtra]++;
+                MOpt_globals.cycles_v2[STAGE_OverWriteExtra]++;
 
-              break;
+                break;
 
-            }
+              }
 
               /* Insert an extra. */
 
-            case 17: {
+              else if (r == 1 && (afl->extras_cnt || afl->a_extras_cnt)) {
 
-              u32 use_extra, extra_len,
-                  insert_at = rand_below(afl, temp_len + 1);
-              u8 *ptr;
+                u32 use_extra, extra_len,
+                    insert_at = rand_below(afl, temp_len + 1);
+                u8 *ptr;
 
-              /* Insert an extra. Do the same dice-rolling stuff as for the
-                previous case. */
+                /* Insert an extra. Do the same dice-rolling stuff as for the
+                  previous case. */
 
-              if (!afl->extras_cnt ||
-                  (afl->a_extras_cnt && rand_below(afl, 2))) {
+                if (!afl->extras_cnt ||
+                    (afl->a_extras_cnt && rand_below(afl, 2))) {
 
-                use_extra = rand_below(afl, afl->a_extras_cnt);
-                extra_len = afl->a_extras[use_extra].len;
-                ptr = afl->a_extras[use_extra].data;
+                  use_extra = rand_below(afl, afl->a_extras_cnt);
+                  extra_len = afl->a_extras[use_extra].len;
+                  ptr = afl->a_extras[use_extra].data;
 #ifdef INTROSPECTION
-                snprintf(afl->m_tmp, sizeof(afl->m_tmp),
-                         " AUTO_EXTRA_INSERT-%u-%u", insert_at, extra_len);
-                strcat(afl->mutation, afl->m_tmp);
+                  snprintf(afl->m_tmp, sizeof(afl->m_tmp),
+                           " AUTO_EXTRA_INSERT-%u-%u", insert_at, extra_len);
+                  strcat(afl->mutation, afl->m_tmp);
 #endif
 
-              } else {
+                } else {
 
-                use_extra = rand_below(afl, afl->extras_cnt);
-                extra_len = afl->extras[use_extra].len;
-                ptr = afl->extras[use_extra].data;
+                  use_extra = rand_below(afl, afl->extras_cnt);
+                  extra_len = afl->extras[use_extra].len;
+                  ptr = afl->extras[use_extra].data;
 #ifdef INTROSPECTION
-                snprintf(afl->m_tmp, sizeof(afl->m_tmp), " EXTRA_INSERT-%u-%u",
-                         insert_at, extra_len);
-                strcat(afl->mutation, afl->m_tmp);
+                  snprintf(afl->m_tmp, sizeof(afl->m_tmp),
+                           " EXTRA_INSERT-%u-%u", insert_at, extra_len);
+                  strcat(afl->mutation, afl->m_tmp);
 #endif
 
-              }
-
-              if (temp_len + extra_len >= MAX_FILE) break;
+                }
 
-              out_buf = afl_realloc(AFL_BUF_PARAM(out), temp_len + extra_len);
-              if (unlikely(!out_buf)) { PFATAL("alloc"); }
+                if (temp_len + extra_len >= MAX_FILE) break;
 
-              /* Tail */
-              memmove(out_buf + insert_at + extra_len, out_buf + insert_at,
-                      temp_len - insert_at);
+                out_buf = afl_realloc(AFL_BUF_PARAM(out), temp_len + extra_len);
+                if (unlikely(!out_buf)) { PFATAL("alloc"); }
 
-              /* Inserted part */
-              memcpy(out_buf + insert_at, ptr, extra_len);
+                /* Tail */
+                memmove(out_buf + insert_at + extra_len, out_buf + insert_at,
+                        temp_len - insert_at);
 
-              temp_len += extra_len;
-              MOpt_globals.cycles_v2[STAGE_InsertExtra]++;
-              break;
+                /* Inserted part */
+                memcpy(out_buf + insert_at, ptr, extra_len);
 
-            }
+                temp_len += extra_len;
+                MOpt_globals.cycles_v2[STAGE_InsertExtra]++;
+                break;
 
-            default: {
+              } else {
 
-              if (unlikely(afl->ready_for_splicing_count < 2)) break;
+                if (unlikely(afl->ready_for_splicing_count < 2)) break;
 
-              u32 tid;
-              do {
+                u32 tid;
+                do {
 
-                tid = rand_below(afl, afl->queued_items);
+                  tid = rand_below(afl, afl->queued_items);
 
-              } while (tid == afl->current_entry ||
+                } while (tid == afl->current_entry ||
 
-                       afl->queue_buf[tid]->len < 4);
+                         afl->queue_buf[tid]->len < 4);
 
-              /* Get the testcase for splicing. */
-              struct queue_entry *target = afl->queue_buf[tid];
-              u32                 new_len = target->len;
-              u8 *                new_buf = queue_testcase_get(afl, target);
+                /* Get the testcase for splicing. */
+                struct queue_entry *target = afl->queue_buf[tid];
+                u32                 new_len = target->len;
+                u8 *                new_buf = queue_testcase_get(afl, target);
 
-              if ((temp_len >= 2 && rand_below(afl, 2)) ||
-                  temp_len + HAVOC_BLK_XL >= MAX_FILE) {
+                if ((temp_len >= 2 && rand_below(afl, 2)) ||
+                    temp_len + HAVOC_BLK_XL >= MAX_FILE) {
 
-                /* overwrite mode */
+                  /* overwrite mode */
 
-                u32 copy_from, copy_to, copy_len;
+                  u32 copy_from, copy_to, copy_len;
 
-                copy_len = choose_block_len(afl, new_len - 1);
-                if (copy_len > temp_len) copy_len = temp_len;
+                  copy_len = choose_block_len(afl, new_len - 1);
+                  if (copy_len > temp_len) copy_len = temp_len;
 
-                copy_from = rand_below(afl, new_len - copy_len + 1);
-                copy_to = rand_below(afl, temp_len - copy_len + 1);
+                  copy_from = rand_below(afl, new_len - copy_len + 1);
+                  copy_to = rand_below(afl, temp_len - copy_len + 1);
 
 #ifdef INTROSPECTION
-                snprintf(afl->m_tmp, sizeof(afl->m_tmp),
-                         " SPLICE_OVERWRITE-%u-%u-%u-%s", copy_from, copy_to,
-                         copy_len, target->fname);
-                strcat(afl->mutation, afl->m_tmp);
+                  snprintf(afl->m_tmp, sizeof(afl->m_tmp),
+                           " SPLICE_OVERWRITE-%u-%u-%u-%s", copy_from, copy_to,
+                           copy_len, target->fname);
+                  strcat(afl->mutation, afl->m_tmp);
 #endif
-                memmove(out_buf + copy_to, new_buf + copy_from, copy_len);
+                  memmove(out_buf + copy_to, new_buf + copy_from, copy_len);
 
-              } else {
+                } else {
 
-                /* insert mode */
+                  /* insert mode */
 
-                u32 clone_from, clone_to, clone_len;
+                  u32 clone_from, clone_to, clone_len;
 
-                clone_len = choose_block_len(afl, new_len);
-                clone_from = rand_below(afl, new_len - clone_len + 1);
-                clone_to = rand_below(afl, temp_len + 1);
+                  clone_len = choose_block_len(afl, new_len);
+                  clone_from = rand_below(afl, new_len - clone_len + 1);
+                  clone_to = rand_below(afl, temp_len + 1);
 
-                u8 *temp_buf = afl_realloc(AFL_BUF_PARAM(out_scratch),
-                                           temp_len + clone_len + 1);
-                if (unlikely(!temp_buf)) { PFATAL("alloc"); }
+                  u8 *temp_buf = afl_realloc(AFL_BUF_PARAM(out_scratch),
+                                             temp_len + clone_len + 1);
+                  if (unlikely(!temp_buf)) { PFATAL("alloc"); }
 
 #ifdef INTROSPECTION
-                snprintf(afl->m_tmp, sizeof(afl->m_tmp),
-                         " SPLICE_INSERT-%u-%u-%u-%s", clone_from, clone_to,
-                         clone_len, target->fname);
-                strcat(afl->mutation, afl->m_tmp);
+                  snprintf(afl->m_tmp, sizeof(afl->m_tmp),
+                           " SPLICE_INSERT-%u-%u-%u-%s", clone_from, clone_to,
+                           clone_len, target->fname);
+                  strcat(afl->mutation, afl->m_tmp);
 #endif
-                /* Head */
+                  /* Head */
 
-                memcpy(temp_buf, out_buf, clone_to);
+                  memcpy(temp_buf, out_buf, clone_to);
 
-                /* Inserted part */
+                  /* Inserted part */
 
-                memcpy(temp_buf + clone_to, new_buf + clone_from, clone_len);
+                  memcpy(temp_buf + clone_to, new_buf + clone_from, clone_len);
 
-                /* Tail */
-                memcpy(temp_buf + clone_to + clone_len, out_buf + clone_to,
-                       temp_len - clone_to);
+                  /* Tail */
+                  memcpy(temp_buf + clone_to + clone_len, out_buf + clone_to,
+                         temp_len - clone_to);
 
-                out_buf = temp_buf;
-                afl_swap_bufs(AFL_BUF_PARAM(out), AFL_BUF_PARAM(out_scratch));
-                temp_len += clone_len;
+                  out_buf = temp_buf;
+                  afl_swap_bufs(AFL_BUF_PARAM(out), AFL_BUF_PARAM(out_scratch));
+                  temp_len += clone_len;
 
-              }
+                }
 
-              MOpt_globals.cycles_v2[STAGE_Splice]++;
-              break;
+                MOpt_globals.cycles_v2[STAGE_Splice]++;
+                break;
 
-            }  // end of default:
+              }  // end of default:
 
           }                                    /* switch select_algorithm() */
 
diff --git a/src/afl-fuzz-stats.c b/src/afl-fuzz-stats.c
index 9737c692..5b237748 100644
--- a/src/afl-fuzz-stats.c
+++ b/src/afl-fuzz-stats.c
@@ -1021,13 +1021,15 @@ void show_stats_normal(afl_state_t *afl) {
 
   if (unlikely(!afl->skip_deterministic)) {
 
-    sprintf(tmp, "%s/%s, %s/%s, %s/%s",
+    sprintf(tmp, "%s/%s, %s/%s, %s/%s, %s/%s",
             u_stringify_int(IB(0), afl->stage_finds[STAGE_EXTRAS_UO]),
             u_stringify_int(IB(1), afl->stage_cycles[STAGE_EXTRAS_UO]),
             u_stringify_int(IB(2), afl->stage_finds[STAGE_EXTRAS_UI]),
             u_stringify_int(IB(3), afl->stage_cycles[STAGE_EXTRAS_UI]),
             u_stringify_int(IB(4), afl->stage_finds[STAGE_EXTRAS_AO]),
-            u_stringify_int(IB(5), afl->stage_cycles[STAGE_EXTRAS_AO]));
+            u_stringify_int(IB(5), afl->stage_cycles[STAGE_EXTRAS_AO]),
+            u_stringify_int(IB(6), afl->stage_finds[STAGE_EXTRAS_AI]),
+            u_stringify_int(IB(7), afl->stage_cycles[STAGE_EXTRAS_AI]));
 
   } else if (unlikely(!afl->extras_cnt || afl->custom_only)) {
 
@@ -1839,13 +1841,15 @@ void show_stats_pizza(afl_state_t *afl) {
 
   if (unlikely(!afl->skip_deterministic)) {
 
-    sprintf(tmp, "%s/%s, %s/%s, %s/%s",
+    sprintf(tmp, "%s/%s, %s/%s, %s/%s, %s/%s",
             u_stringify_int(IB(0), afl->stage_finds[STAGE_EXTRAS_UO]),
             u_stringify_int(IB(1), afl->stage_cycles[STAGE_EXTRAS_UO]),
             u_stringify_int(IB(2), afl->stage_finds[STAGE_EXTRAS_UI]),
             u_stringify_int(IB(3), afl->stage_cycles[STAGE_EXTRAS_UI]),
             u_stringify_int(IB(4), afl->stage_finds[STAGE_EXTRAS_AO]),
-            u_stringify_int(IB(5), afl->stage_cycles[STAGE_EXTRAS_AO]));
+            u_stringify_int(IB(5), afl->stage_cycles[STAGE_EXTRAS_AO]),
+            u_stringify_int(IB(6), afl->stage_finds[STAGE_EXTRAS_AI]),
+            u_stringify_int(IB(7), afl->stage_cycles[STAGE_EXTRAS_AI]));
 
   } else if (unlikely(!afl->extras_cnt || afl->custom_only)) {