changes

4 files changed, 31 insertions, 8 deletions

diff --git a/include/afl-fuzz.h b/include/afl-fuzz.h
index 37e2dc6c..5e4e5a19 100644
--- a/include/afl-fuzz.h
+++ b/include/afl-fuzz.h
@@ -896,6 +896,7 @@ void write_bitmap(afl_state_t *);
 u32  count_bits(afl_state_t *, u8 *);
 u32  count_bits_len(afl_state_t *, u8 *, u32);
 u32  count_bytes(afl_state_t *, u8 *);
+u32  count_bytes_len(afl_state_t *, u8 *, u32);
 u32  count_non_255_bytes(afl_state_t *, u8 *);
 #ifdef WORD_SIZE_64
 void simplify_trace(afl_state_t *, u64 *);
diff --git a/src/afl-fuzz-bitmap.c b/src/afl-fuzz-bitmap.c
index 9cb1b83f..8aaa4ae1 100644
--- a/src/afl-fuzz-bitmap.c
+++ b/src/afl-fuzz-bitmap.c
@@ -235,6 +235,29 @@ u32 count_bytes(afl_state_t *afl, u8 *mem) {
 
 }
 
+u32 count_bytes_len(afl_state_t *afl, u8 *mem, u32 len) {
+
+  u32 *ptr = (u32 *)mem;
+  u32  i = (len >> 2);
+  u32  ret = 0;
+
+  while (i--) {
+
+    u32 v = *(ptr++);
+
+    if (!v) { continue; }
+    if (v & 0x000000ff) { ++ret; }
+    if (v & 0x0000ff00) { ++ret; }
+    if (v & 0x00ff0000) { ++ret; }
+    if (v & 0xff000000) { ++ret; }
+
+  }
+
+  return ret;
+
+}
+
+
 /* Count the number of non-255 bytes set in the bitmap. Used strictly for the
    status screen, several calls per second or so. */
 
diff --git a/src/afl-fuzz-run.c b/src/afl-fuzz-run.c
index badc2239..b325f788 100644
--- a/src/afl-fuzz-run.c
+++ b/src/afl-fuzz-run.c
@@ -479,10 +479,11 @@ abort_calibration:
     if (afl_fsrv_run_target(&afl->taint_fsrv, use_tmout, &afl->stop_soon) ==
         0) {
 
-      u32 len = q->len / 8;
-      if (q->len % 8) len++;
-      u32 bits = count_bits_len(afl, afl->taint_fsrv.trace_bits, len);
-      if (afl->debug) fprintf(stderr, "Debug: tainted bytes: %u\n", bits);
+      u32 len = q->len;
+      if (len % 4)
+        len = len + 4 - (q->len % 4);
+      u32 bytes = count_bytes_len(afl, afl->taint_fsrv.trace_bits, len);
+      if (afl->debug) fprintf(stderr, "Debug: tainted bytes: %u\n", bytes);
 
     }
 
diff --git a/src/afl-fuzz.c b/src/afl-fuzz.c
index 684b123e..4a3d2e97 100644
--- a/src/afl-fuzz.c
+++ b/src/afl-fuzz.c
@@ -825,11 +825,9 @@ int main(int argc, char **argv_orig, char **envp) {
 
   }
 
-  if (afl->fsrv.taint_mode && afl->fsrv.map_size < (MAX_FILE / 8) + 1) {
+  if (afl->fsrv.taint_mode && afl->fsrv.map_size < MAX_FILE) {
 
-    afl->shm.map_size = (MAX_FILE / 8);
-    if (MAX_FILE % 8) afl->shm.map_size++;
-    afl->fsrv.map_size = afl->shm.map_size;
+    afl->fsrv.map_size = afl->shm.map_size = MAX_FILE;
 
   }