set max testcache entries automated if not specified by the user

7 files changed, 91 insertions, 18 deletions

diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index 220380b9..11feb9f7 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -366,7 +366,7 @@ typedef struct afl_env_vars {
   u8 *afl_tmpdir, *afl_custom_mutator_library, *afl_python_module, *afl_path,
       *afl_hang_tmout, *afl_forksrv_init_tmout, *afl_skip_crashes, *afl_preload,
       *afl_max_det_extras, *afl_statsd_host, *afl_statsd_port,
-      *afl_statsd_tags_flavor, *afl_testcache_size;
+      *afl_statsd_tags_flavor, *afl_testcache_size, *afl_testcache_entries;
 
 } afl_env_vars_t;
 
@@ -695,6 +695,9 @@ typedef struct afl_state {
   /* This is the user specified maximum size to use for the testcase cache */
   u64 q_testcase_max_cache_size;
 
+  /* This is the user specified maximum entries in the testcase cache */
+  u32 q_testcase_max_cache_entries;
+
   /* How much of the testcase cache is used so far */
   u64 q_testcase_cache_size;
 
@@ -712,7 +715,7 @@ typedef struct afl_state {
 
   /* Refs to each queue entry with cached testcase (for eviction, if cache_count
    * is too large) */
-  struct queue_entry *q_testcase_cache[TESTCASE_ENTRIES];
+  struct queue_entry **q_testcase_cache;
 
 } afl_state_t;
 
diff --git a/include/config.h b/include/config.h
index b4f3a775..491d8132 100644
--- a/include/config.h
+++ b/include/config.h
@@ -295,14 +295,12 @@
 
 #define RESEED_RNG 100000
 
-/* The maximum number of testcases to cache */
-
-#define TESTCASE_ENTRIES 16384
-
 /* The default maximum testcase cache size in MB, 0 = disable.
-   A value between 50 and 250 is a good default value. */
+   A value between 50 and 250 is a good default value. Note that the
+   number of entries will be auto assigned if not specified via the
+   AFL_TESTCACHE_ENTRIES env variable */
 
-#define TESTCASE_CACHE 0
+#define TESTCASE_CACHE_SIZE 50
 
 /* Maximum line length passed from GCC to 'as' and used for parsing
    configuration files: */
diff --git a/include/envs.h b/include/envs.h
index a1b3ad12..b753d5f8 100644
--- a/include/envs.h
+++ b/include/envs.h
@@ -140,6 +140,7 @@ static char *afl_environment_variables[] = {
     "AFL_STATSD_PORT",
     "AFL_STATSD_TAGS_FLAVOR",
     "AFL_TESTCACHE_SIZE",
+    "AFL_TESTCACHE_ENTRIES",
     "AFL_TMIN_EXACT",
     "AFL_TMPDIR",
     "AFL_TOKEN_FILE",
diff --git a/src/afl-fuzz-queue.c b/src/afl-fuzz-queue.c
index db387c33..baa80e61 100644
--- a/src/afl-fuzz-queue.c
+++ b/src/afl-fuzz-queue.c
@@ -978,9 +978,9 @@ inline u8 *queue_testcase_get(afl_state_t *afl, struct queue_entry *q) {
     /* Buf not cached, let's load it */
     u32 tid = afl->q_testcase_max_cache_count;
 
-    while (unlikely(afl->q_testcase_cache_size + len >=
-                        afl->q_testcase_max_cache_size ||
-                    afl->q_testcase_cache_count >= TESTCASE_ENTRIES - 1)) {
+    while (unlikely(
+        afl->q_testcase_cache_size + len >= afl->q_testcase_max_cache_size ||
+        afl->q_testcase_cache_count >= afl->q_testcase_max_cache_entries - 1)) {
 
       /* Cache full. We neet to evict one or more to map one.
          Get a random one which is not in use */
@@ -1009,7 +1009,7 @@ inline u8 *queue_testcase_get(afl_state_t *afl, struct queue_entry *q) {
 
     }
 
-    if (unlikely(tid >= TESTCASE_ENTRIES)) {
+    if (unlikely(tid >= afl->q_testcase_max_cache_entries)) {
 
       // uh we were full, so now we have to search from start
       tid = afl->q_testcase_smallest_free;
@@ -1062,7 +1062,8 @@ inline void queue_testcase_store_mem(afl_state_t *afl, struct queue_entry *q,
 
   if (unlikely(afl->q_testcase_cache_size + len >=
                    afl->q_testcase_max_cache_size ||
-               afl->q_testcase_cache_count >= TESTCASE_ENTRIES - 1)) {
+               afl->q_testcase_cache_count >=
+                   afl->q_testcase_max_cache_entries - 1)) {
 
     // no space? will be loaded regularly later.
     return;
diff --git a/src/afl-fuzz-state.c b/src/afl-fuzz-state.c
index 0824b77f..ae7d410b 100644
--- a/src/afl-fuzz-state.c
+++ b/src/afl-fuzz-state.c
@@ -103,7 +103,8 @@ void afl_state_init(afl_state_t *afl, uint32_t map_size) {
   afl->stats_avg_exec = -1;
   afl->skip_deterministic = 1;
   afl->use_splicing = 1;
-  afl->q_testcase_max_cache_size = TESTCASE_CACHE * 1024000;
+  afl->q_testcase_max_cache_size = TESTCASE_CACHE_SIZE * 1048576UL;
+  afl->q_testcase_max_cache_entries = 4096;
 
 #ifdef HAVE_AFFINITY
   afl->cpu_aff = -1;                    /* Selected CPU core                */
@@ -361,6 +362,13 @@ void read_afl_environment(afl_state_t *afl, char **envp) {
             afl->afl_env.afl_testcache_size =
                 (u8 *)get_afl_env(afl_environment_variables[i]);
 
+          } else if (!strncmp(env, "AFL_TESTCACHE_ENTRIES",
+
+                              afl_environment_variable_len)) {
+
+            afl->afl_env.afl_testcache_entries =
+                (u8 *)get_afl_env(afl_environment_variables[i]);
+
           } else if (!strncmp(env, "AFL_STATSD_HOST",
 
                               afl_environment_variable_len)) {
diff --git a/src/afl-fuzz-stats.c b/src/afl-fuzz-stats.c
index d213d054..bec90519 100644
--- a/src/afl-fuzz-stats.c
+++ b/src/afl-fuzz-stats.c
@@ -1028,7 +1028,7 @@ void show_init_stats(afl_state_t *afl) {
 
   }
 
-  SAYF("\n");
+  // SAYF("\n");
 
   if (avg_us > ((afl->fsrv.qemu_mode || afl->unicorn_mode) ? 50000 : 10000)) {
 
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index 7215ecec..637e1985 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -251,7 +251,7 @@ static int stricmp(char const *a, char const *b) {
 
 int main(int argc, char **argv_orig, char **envp) {
 
-  s32 opt, i, auto_sync = 0;
+  s32 opt, i, auto_sync = 0, user_set_cache = 0;
   u64 prev_queued = 0;
   u32 sync_interval_cnt = 0, seek_to = 0, show_help = 0, map_size = MAP_SIZE;
   u8 *extras_dir[4];
@@ -1015,6 +1015,22 @@ int main(int argc, char **argv_orig, char **envp) {
 
   }
 
+  if (afl->afl_env.afl_testcache_entries) {
+
+    afl->q_testcase_max_cache_entries =
+        (u32)atoi(afl->afl_env.afl_testcache_entries);
+
+    user_set_cache = 1;
+
+  }
+
+  if (!afl->afl_env.afl_testcache_size || !afl->afl_env.afl_testcache_entries) {
+
+    afl->afl_env.afl_testcache_entries = 0;
+    afl->afl_env.afl_testcache_size = 0;
+
+  }
+
   if (!afl->q_testcase_max_cache_size) {
 
     ACTF(
@@ -1347,6 +1363,52 @@ int main(int argc, char **argv_orig, char **envp) {
 
   perform_dry_run(afl);
 
+  if (!user_set_cache && afl->q_testcase_max_cache_size) {
+
+    /* The user defined not a fixed number of entries for the cache.
+       Hence we autodetect a good value. After the dry run inputs are
+       trimmed and we know the average and max size of the input seeds.
+       We use this information to set a fitting size to max entries
+       based on the cache size. */
+
+    struct queue_entry *q = afl->queue;
+    u64                 size = 0, count = 0, avg = 0, max = 0;
+
+    while (q) {
+
+      ++count;
+      size += q->len;
+      if (max < q->len) { max = q->len; }
+      q = q->next;
+
+    }
+
+    if (count) {
+
+      avg = size / count;
+      avg = ((avg + max) / 2) + 1;
+
+    }
+
+    if (avg < 10240) { avg = 10240; }
+
+    afl->q_testcase_max_cache_entries = afl->q_testcase_max_cache_size / avg;
+
+    if (afl->q_testcase_max_cache_entries > 32768)
+      afl->q_testcase_max_cache_entries = 32768;
+
+  }
+
+  if (afl->q_testcase_max_cache_entries) {
+
+    OKF("Setting %u maximum entries for the testcase cache",
+        afl->q_testcase_max_cache_entries);
+    afl->q_testcase_cache =
+        ck_alloc(afl->q_testcase_max_cache_entries * sizeof(size_t));
+    if (!afl->q_testcase_cache) { PFATAL("malloc failed for cache entries"); }
+
+  }
+
   cull_queue(afl);
 
   if (!afl->pending_not_fuzzed)
@@ -1366,8 +1428,7 @@ int main(int argc, char **argv_orig, char **envp) {
 
   if (!afl->not_on_tty) {
 
-    sleep(4);
-    afl->start_time += 4000;
+    sleep(1);
     if (afl->stop_soon) { goto stop_fuzzing; }
 
   }
@@ -1654,6 +1715,7 @@ stop_fuzzing:
   ck_free(afl->fsrv.target_path);
   ck_free(afl->fsrv.out_file);
   ck_free(afl->sync_id);
+  if (afl->q_testcase_cache) { ck_free(afl->q_testcase_cache); }
   afl_state_deinit(afl);
   free(afl);                                                 /* not tracked */