1 files changed, 38 insertions, 19 deletions

diff --git a/instrumentation/afl-compiler-rt.o.c b/instrumentation/afl-compiler-rt.o.c
index 9871d7f4..a88396d4 100644
--- a/instrumentation/afl-compiler-rt.o.c
+++ b/instrumentation/afl-compiler-rt.o.c
@@ -1539,12 +1539,16 @@ void __sanitizer_cov_trace_pc_guard_init(uint32_t *start, uint32_t *stop) {
   if (start == stop || *start) return;
 
   x = getenv("AFL_INST_RATIO");
-  if (x) { inst_ratio = (u32)atoi(x); }
+  if (x) {
 
-  if (!inst_ratio || inst_ratio > 100) {
+    inst_ratio = (u32)atoi(x);
 
-    fprintf(stderr, "[-] ERROR: Invalid AFL_INST_RATIO (must be 1-100).\n");
-    abort();
+    if (!inst_ratio || inst_ratio > 100) {
+
+      fprintf(stderr, "[-] ERROR: Invalid AFL_INST_RATIO (must be 1-100).\n");
+      abort();
+
+    }
 
   }
 
@@ -1568,10 +1572,16 @@ void __sanitizer_cov_trace_pc_guard_init(uint32_t *start, uint32_t *stop) {
 
       while (start < stop) {
 
-        if (likely(inst_ratio == 100) || R(100) < inst_ratio)
-          *start = offset;
-        else
-          *start = 0;  // write to map[0]
+        if (likely(inst_ratio == 100) || R(100) < inst_ratio) {
+
+          *(start++) = offset;
+
+        } else {
+
+          *(start++) = 0;  // write to map[0]
+
+        }
+
         if (unlikely(++offset >= __afl_final_loc)) { offset = 4; }
 
       }
@@ -1592,12 +1602,15 @@ void __sanitizer_cov_trace_pc_guard_init(uint32_t *start, uint32_t *stop) {
 
   while (start < stop) {
 
-    if (likely(inst_ratio == 100) || R(100) < inst_ratio)
-      *start = ++__afl_final_loc;
-    else
-      *start = 0;  // write to map[0]
+    if (likely(inst_ratio == 100) || R(100) < inst_ratio) {
+
+      *(start++) = ++__afl_final_loc;
+
+    } else {
+
+      *(start++) = 0;  // write to map[0]
 
-    start++;
+    }
 
   }
 
@@ -1609,17 +1622,23 @@ void __sanitizer_cov_trace_pc_guard_init(uint32_t *start, uint32_t *stop) {
 
   }
 
-  if (__afl_already_initialized_shm && __afl_final_loc > __afl_map_size) {
+  if (__afl_already_initialized_shm) {
 
-    if (__afl_debug) {
+    if (__afl_final_loc > __afl_map_size) {
+
+      if (__afl_debug) {
+
+        fprintf(stderr, "Reinit shm necessary (+%u)\n",
+                __afl_final_loc - __afl_map_size);
+
+      }
 
-      fprintf(stderr, "Reinit shm necessary (+%u)\n",
-              __afl_final_loc - __afl_map_size);
+      __afl_unmap_shm();
+      __afl_map_shm();
 
     }
 
-    __afl_unmap_shm();
-    __afl_map_shm();
+    __afl_map_size = __afl_final_loc + 1;
 
   }