notzero for afl-gcc

2 files changed, 12 insertions, 0 deletions

diff --git a/afl-as.h b/afl-as.h
index ebd57109..2c84f9f3 100644
--- a/afl-as.h
+++ b/afl-as.h
@@ -189,6 +189,7 @@ static const u8* main_payload_32 =
   "  orb  $1, (%edx, %edi, 1)\n"
 #else
   "  incb (%edx, %edi, 1)\n"
+  "  adcb $0, (%edx, %edi, 1)\n"
 #endif /* ^SKIP_COUNTS */
   "\n"
   "__afl_return:\n"
@@ -417,6 +418,7 @@ static const u8* main_payload_64 =
   "  orb  $1, (%rdx, %rcx, 1)\n"
 #else
   "  incb (%rdx, %rcx, 1)\n"
+  "  adcb $0, (%rdx, %rcx, 1)\n"
 #endif /* ^SKIP_COUNTS */
   "\n"
   "__afl_return:\n"
diff --git a/llvm_mode/afl-llvm-pass.so.cc b/llvm_mode/afl-llvm-pass.so.cc
index b77835c5..6b2232f2 100644
--- a/llvm_mode/afl-llvm-pass.so.cc
+++ b/llvm_mode/afl-llvm-pass.so.cc
@@ -287,6 +287,16 @@ bool AFLCoverage::runOnModule(Module &M) {
              Value *HowMuch = IRB.CreateAdd(ConstantInt::get(Int8Ty, 1), cf);
              Incr = IRB.CreateAdd(Counter, HowMuch);
 
+          } else if (neverZero_counters_str[0] == '5') {
+            auto cf = IRB.CreateICmpEQ(Incr, ConstantInt::get(Int8Ty, 0));
+            auto carry = IRB.CreateZExt(cf, Int8Ty);
+            Incr = IRB.CreateAdd(Incr, carry);
+
+          } else if (neverZero_counters_str[0] == '6') {
+            auto cf = IRB.CreateICmpULT(Incr, ConstantInt::get(Int8Ty, 1));
+            auto carry = IRB.CreateZExt(cf, Int8Ty);
+            Incr = IRB.CreateAdd(Incr, carry);
+           
           // no other implementations yet
           } else {
             fprintf(stderr, "Error: unknown value for AFL_NZERO_COUNTS: %s (valid is 1-4)\n", neverZero_counters_str);