ctx and ngram can be used together now

5 files changed, 150 insertions, 100 deletions

diff --git a/docs/Changelog.md b/docs/Changelog.md
index cae99681..8c0624b6 100644
--- a/docs/Changelog.md
+++ b/docs/Changelog.md
@@ -20,6 +20,7 @@ sending a mail to <afl-users+subscribe@googlegroups.com>.
        address for the shared memory map is used as this increases the
        fuzzing speed
      - fixes to LTO mode if instrumented edges > MAP_SIZE
+     - CTX and NGRAM can now be used together
      - added AFL_LLVM_SKIP_NEVERZERO to skip the never zero coverage counter
        implmentation. For targets with little or no loops or heavy called
        functions. Gives a small performance boost.
diff --git a/docs/env_variables.md b/docs/env_variables.md
index bdbb8520..ab5808ec 100644
--- a/docs/env_variables.md
+++ b/docs/env_variables.md
@@ -97,12 +97,14 @@ Then there are a few specific features that are only available in llvm_mode:
 
     - AFL_LLVM_INSTRUMENT - this configures the instrumentation mode. 
       Available options:
-        DEFAULT - classic AFL (map[cur_loc ^ prev_loc >> 1]++)
+        CLASSIC - classic AFL (map[cur_loc ^ prev_loc >> 1]++) (default)
         CFG - InsTrim instrumentation (see below)
         LTO - LTO instrumentation (see below)
         CTX - context sensitive instrumentation (see below)
         NGRAM-x - deeper previous location coverage (from NGRAM-2 up to NGRAM-16)
-      Only one can be used.
+      In CLASSIC (default) can can also specify CTX and/nor NGRAM, seperate
+      the options with a comma "," then, e.g.:
+        AFL_LLVM_INSTRUMENT=CLASSIC,CTX,NGRAM-4
 
 ### LTO
 
diff --git a/examples/afl_network_proxy/afl-network-client.c b/examples/afl_network_proxy/afl-network-client.c
index 781ea328..cf09b2ad 100644
--- a/examples/afl_network_proxy/afl-network-client.c
+++ b/examples/afl_network_proxy/afl-network-client.c
@@ -321,7 +321,7 @@ int main(int argc, char *argv[]) {
 
     // fprintf(stderr, "Sending testcase with len %u\n", *lenptr);
 #ifdef USE_DEFLATE
-  #ifdef COMPRESS_TESTCASES
+#ifdef COMPRESS_TESTCASES
     // we only compress the testcase if it does not fit in the TCP packet
     if (*lenptr > 1500 - 20 - 32 - 4) {
 
@@ -331,25 +331,27 @@ int main(int argc, char *argv[]) {
                                                   buf2 + 8, buf2_len);
       if (send(s, buf2, *lenptr2 + 8, 0) != *lenptr2 + 8)
         PFATAL("sending test data failed");
-      //fprintf(stderr, "COMPRESS (%u->%u):\n", *lenptr, *lenptr2);
-      //for (u32 i = 0; i < *lenptr; i++)
+      // fprintf(stderr, "COMPRESS (%u->%u):\n", *lenptr, *lenptr2);
+      // for (u32 i = 0; i < *lenptr; i++)
       //  fprintf(stderr, "%02x", buf[i + 4]);
-      //fprintf(stderr, "\n");
-      //for (u32 i = 0; i < *lenptr2; i++)
+      // fprintf(stderr, "\n");
+      // for (u32 i = 0; i < *lenptr2; i++)
       //  fprintf(stderr, "%02x", buf2[i + 8]);
-      //fprintf(stderr, "\n");
+      // fprintf(stderr, "\n");
 
     } else {
-  #endif
+
+#endif
 #endif
       if (send(s, buf, *lenptr + 4, 0) != *lenptr + 4)
         PFATAL("sending test data failed");
 #ifdef USE_DEFLATE
-  #ifdef COMPRESS_TESTCASES
+#ifdef COMPRESS_TESTCASES
       // fprintf(stderr, "unCOMPRESS (%u)\n", *lenptr);
 
     }
-  #endif
+
+#endif
 #endif
 
     received = 0;
@@ -381,9 +383,9 @@ int main(int argc, char *argv[]) {
                                       &decompress_len) != LIBDEFLATE_SUCCESS ||
         decompress_len != __afl_map_size)
       FATAL("decompression failed");
-    // fprintf(stderr, "DECOMPRESS (%u->%u): ", compress_len, decompress_len);
-    // for (u32 i = 0; i < __afl_map_size; i++) fprintf(stderr, "%02x",
-    // __afl_area_ptr[i]); fprintf(stderr, "\n");
+      // fprintf(stderr, "DECOMPRESS (%u->%u): ", compress_len, decompress_len);
+      // for (u32 i = 0; i < __afl_map_size; i++) fprintf(stderr, "%02x",
+      // __afl_area_ptr[i]); fprintf(stderr, "\n");
 #else
     while (received < __afl_map_size &&
            (ret = recv(s, __afl_area_ptr + received, __afl_map_size - received,
diff --git a/llvm_mode/afl-clang-fast.c b/llvm_mode/afl-clang-fast.c
index c59b814d..9f85e5c6 100644
--- a/llvm_mode/afl-clang-fast.c
+++ b/llvm_mode/afl-clang-fast.c
@@ -43,7 +43,7 @@ static u8 * obj_path;                  /* Path to runtime libraries         */
 static u8 **cc_params;                 /* Parameters passed to the real CC  */
 static u32  cc_par_cnt = 1;            /* Param count, including argv0      */
 static u8   llvm_fullpath[PATH_MAX];
-static u8   instrument_mode;
+static u8   instrument_mode, instrument_opt_mode, ngram_size;
 static u8 * lto_flag = AFL_CLANG_FLTO;
 static u8 * march_opt = CFLAGS_OPT;
 static u8   debug;
@@ -60,14 +60,15 @@ enum {
   INSTRUMENT_INSTRIM = 2,
   INSTRUMENT_CFG = 2,
   INSTRUMENT_LTO = 3,
-  INSTRUMENT_CTX = 4,
-  INSTRUMENT_NGRAM = 5  // + ngram value of 2-16 = 7 - 21
+  INSTRUMENT_OPT_CTX = 4,
+  INSTRUMENT_OPT_NGRAM = 8
 
 };
 
-char instrument_mode_string[6][16] = {
+char instrument_mode_string[10][16] = {
 
-    "DEFAULT", "PCGUARD", "CFG", "LTO", "CTX",
+    "CLASSIC", "PCGUARD", "CFG",   "LTO", "CTX", "",
+    "",        "",        "NGRAM", ""
 
 };
 
@@ -562,59 +563,6 @@ int main(int argc, char **argv, char **envp) {
   instrument_mode = INSTRUMENT_PCGUARD;
 #endif
 
-  if (getenv("AFL_LLVM_SKIP_NEVERZERO") && getenv("AFL_LLVM_NOT_ZERO"))
-    FATAL(
-        "AFL_LLVM_NOT_ZERO and AFL_LLVM_SKIP_NEVERZERO can not be set "
-        "together");
-
-  if ((ptr = getenv("AFL_LLVM_INSTRUMENT")) != NULL) {
-
-    if (strncasecmp(ptr, "default", strlen("default")) == 0 ||
-        strncasecmp(ptr, "afl", strlen("afl")) == 0 ||
-        strncasecmp(ptr, "classic", strlen("classic")) == 0)
-      instrument_mode = INSTRUMENT_DEFAULT;
-    if (strncasecmp(ptr, "cfg", strlen("cfg")) == 0 ||
-        strncasecmp(ptr, "instrim", strlen("instrim")) == 0)
-      instrument_mode = INSTRUMENT_CFG;
-    else if (strncasecmp(ptr, "pc-guard", strlen("pc-guard")) == 0 ||
-             strncasecmp(ptr, "pcguard", strlen("pcgard")) == 0)
-      instrument_mode = INSTRUMENT_PCGUARD;
-    else if (strncasecmp(ptr, "lto", strlen("lto")) == 0)
-      instrument_mode = INSTRUMENT_LTO;
-    else if (strncasecmp(ptr, "ctx", strlen("ctx")) == 0) {
-
-      instrument_mode = INSTRUMENT_CTX;
-      setenv("AFL_LLVM_CTX", "1", 1);
-
-    } else if (strncasecmp(ptr, "ngram", strlen("ngram")) == 0) {
-
-      ptr += strlen("ngram");
-      while (*ptr && (*ptr < '0' || *ptr > '9'))
-        ptr++;
-      if (!*ptr)
-        if ((ptr = getenv("AFL_LLVM_NGRAM_SIZE")) != NULL)
-          FATAL(
-              "you must set the NGRAM size with (e.g. for value 2) "
-              "AFL_LLVM_INSTRUMENT=ngram-2");
-      instrument_mode = INSTRUMENT_NGRAM + atoi(ptr);
-      if (instrument_mode < INSTRUMENT_NGRAM + 2 ||
-          instrument_mode > INSTRUMENT_NGRAM + NGRAM_SIZE_MAX)
-        FATAL(
-            "NGRAM instrumentation mode must be between 2 and NGRAM_SIZE_MAX "
-            "(%u)",
-            NGRAM_SIZE_MAX);
-
-      ptr = alloc_printf("%u", instrument_mode - INSTRUMENT_NGRAM);
-      setenv("AFL_LLVM_NGRAM_SIZE", ptr, 1);
-
-    } else if (strncasecmp(ptr, "classic", strlen("classic")) != 0 ||
-
-               strncasecmp(ptr, "default", strlen("default")) != 0 ||
-               strncasecmp(ptr, "afl", strlen("afl")) != 0)
-      FATAL("unknown AFL_LLVM_INSTRUMENT value: %s", ptr);
-
-  }
-
   if (getenv("USE_TRACE_PC") || getenv("AFL_USE_TRACE_PC") ||
       getenv("AFL_LLVM_USE_TRACE_PC") || getenv("AFL_TRACE_PC")) {
 
@@ -636,39 +584,116 @@ int main(int argc, char **argv, char **envp) {
 
   }
 
-  if (getenv("AFL_LLVM_CTX")) {
+  if (getenv("AFL_LLVM_CTX")) instrument_opt_mode |= INSTRUMENT_OPT_CTX;
 
-    if (instrument_mode == 0)
-      instrument_mode = INSTRUMENT_CTX;
-    else if (instrument_mode != INSTRUMENT_CTX)
-      FATAL("you can not set AFL_LLVM_INSTRUMENT and AFL_LLVM_CTX together");
+  if (getenv("AFL_LLVM_NGRAM_SIZE")) {
+
+    instrument_opt_mode |= INSTRUMENT_OPT_NGRAM;
+    ngram_size = atoi(getenv("AFL_LLVM_NGRAM_SIZE"));
+    if (ngram_size < 2 || ngram_size > NGRAM_SIZE_MAX)
+      FATAL(
+          "NGRAM instrumentation mode must be between 2 and NGRAM_SIZE_MAX "
+          "(%u)",
+          NGRAM_SIZE_MAX);
 
   }
 
-  if (getenv("AFL_LLVM_NGRAM_SIZE")) {
+  if (getenv("AFL_LLVM_INSTRUMENT")) {
 
-    if (instrument_mode == 0) {
+    u8 *ptr = strtok(getenv("AFL_LLVM_INSTRUMENT"), ":,;");
 
-      instrument_mode = INSTRUMENT_NGRAM + atoi(getenv("AFL_LLVM_NGRAM_SIZE"));
-      if (instrument_mode < INSTRUMENT_NGRAM + 2 ||
-          instrument_mode > INSTRUMENT_NGRAM + NGRAM_SIZE_MAX)
-        FATAL(
-            "NGRAM instrumentation mode must be between 2 and NGRAM_SIZE_MAX "
-            "(%u)",
-            NGRAM_SIZE_MAX);
+    while (ptr) {
 
-    } else if (instrument_mode != INSTRUMENT_NGRAM)
+      if (strncasecmp(ptr, "default", strlen("default")) == 0 ||
+          strncasecmp(ptr, "afl", strlen("afl")) == 0 ||
+          strncasecmp(ptr, "classic", strlen("classic")) == 0) {
 
-      FATAL(
-          "you can not set AFL_LLVM_INSTRUMENT and AFL_LLVM_NGRAM_SIZE "
-          "together");
+        if (!instrument_mode || instrument_mode == INSTRUMENT_DEFAULT)
+          instrument_mode = INSTRUMENT_DEFAULT;
+        else
+          FATAL("main instrumentation mode already set with %s",
+                instrument_mode_string[instrument_mode]);
+
+      }
+
+      if (strncasecmp(ptr, "pc-guard", strlen("pc-guard")) == 0 ||
+          strncasecmp(ptr, "pcguard", strlen("pcgard")) == 0) {
+
+        if (!instrument_mode || instrument_mode == INSTRUMENT_PCGUARD)
+          instrument_mode = INSTRUMENT_PCGUARD;
+        else
+          FATAL("main instrumentation mode already set with %s",
+                instrument_mode_string[instrument_mode]);
+
+      }
+
+      if (strncasecmp(ptr, "cfg", strlen("cfg")) == 0 ||
+          strncasecmp(ptr, "instrim", strlen("instrim")) == 0) {
+
+        if (!instrument_mode || instrument_mode == INSTRUMENT_CFG)
+          instrument_mode = INSTRUMENT_CFG;
+        else
+          FATAL("main instrumentation mode already set with %s",
+                instrument_mode_string[instrument_mode]);
+
+      }
+
+      if (strncasecmp(ptr, "lto", strlen("lto")) == 0) {
+
+        if (!instrument_mode || instrument_mode == INSTRUMENT_LTO)
+          instrument_mode = INSTRUMENT_LTO;
+        else
+          FATAL("main instrumentation mode already set with %s",
+                instrument_mode_string[instrument_mode]);
+
+      }
+
+      if (strncasecmp(ptr, "ctx", strlen("ctx")) == 0) {
+
+        instrument_opt_mode |= INSTRUMENT_OPT_CTX;
+        setenv("AFL_LLVM_CTX", "1", 1);
+
+      }
+
+      if (strncasecmp(ptr, "ngram", strlen("ngram")) == 0) {
+
+        ptr += strlen("ngram");
+        while (*ptr && (*ptr < '0' || *ptr > '9'))
+          ptr++;
+        if (!*ptr)
+          if ((ptr = getenv("AFL_LLVM_NGRAM_SIZE")) != NULL)
+            FATAL(
+                "you must set the NGRAM size with (e.g. for value 2) "
+                "AFL_LLVM_INSTRUMENT=ngram-2");
+        ngram_size = atoi(ptr);
+        if (ngram_size < 2 || ngram_size > NGRAM_SIZE_MAX)
+          FATAL(
+              "NGRAM instrumentation option must be between 2 and "
+              "NGRAM_SIZE_MAX "
+              "(%u)",
+              NGRAM_SIZE_MAX);
+        instrument_opt_mode |= (INSTRUMENT_OPT_NGRAM);
+        ptr = alloc_printf("%u", ngram_size);
+        setenv("AFL_LLVM_NGRAM_SIZE", ptr, 1);
+
+      }
+
+      ptr = strtok(NULL, ":,;");
+
+    }
 
   }
 
-  if (instrument_mode < INSTRUMENT_NGRAM)
+  if (!instrument_opt_mode)
     ptr = instrument_mode_string[instrument_mode];
+  else if (instrument_opt_mode == INSTRUMENT_OPT_CTX)
+    ptr = alloc_printf("%s + CTX", instrument_mode_string[instrument_mode]);
+  else if (instrument_opt_mode == INSTRUMENT_OPT_NGRAM)
+    ptr = alloc_printf("%s + NGRAM-%u", instrument_mode_string[instrument_mode],
+                       ngram_size);
   else
-    ptr = alloc_printf("NGRAM-%u", instrument_mode - INSTRUMENT_NGRAM);
+    ptr = alloc_printf("%s + CTX + NGRAM-%u",
+                       instrument_mode_string[instrument_mode], ngram_size);
 
   if (strstr(argv[0], "afl-clang-lto") != NULL) {
 
@@ -690,9 +715,28 @@ int main(int argc, char **argv, char **envp) {
 
 #ifndef AFL_CLANG_FLTO
   if (instrument_mode == INSTRUMENT_LTO)
-    FATAL("instrumentation mode LTO specified but LLVM support not available");
+    FATAL(
+        "instrumentation mode LTO specified but LLVM support not available "
+        "(requires LLVM 11)");
 #endif
 
+  if (instrument_opt_mode && instrument_mode != INSTRUMENT_CLASSIC)
+    /*&& instrument_mode != INSTRUMENT_CFG*/
+    FATAL(
+        "CTX and NGRAM instrumentation options can only be used with the "
+        "CLASSIC instrumentation mode!");
+
+  if (getenv("AFL_LLVM_SKIP_NEVERZERO") && getenv("AFL_LLVM_NOT_ZERO"))
+    FATAL(
+        "AFL_LLVM_NOT_ZERO and AFL_LLVM_SKIP_NEVERZERO can not be set "
+        "together");
+
+  if (instrument_mode == INSTRUMENT_CFG &&
+      getenv("AFL_LLVM_INSTRIM_SKIPSINGLEBLOCK") == NULL && ngram_size)
+    FATAL(
+        "NGRAM option together with CFG/INSTRIM instrumentation mode can only "
+        "be used if AFL_LLVM_INSTRIM_SKIPSINGLEBLOCK is set");
+
   if (argc < 2 || strcmp(argv[1], "-h") == 0) {
 
     if (instrument_mode != INSTRUMENT_LTO)
diff --git a/llvm_mode/afl-llvm-pass.so.cc b/llvm_mode/afl-llvm-pass.so.cc
index 3e9026c8..42a2f3af 100644
--- a/llvm_mode/afl-llvm-pass.so.cc
+++ b/llvm_mode/afl-llvm-pass.so.cc
@@ -182,7 +182,7 @@ bool AFLCoverage::runOnModule(Module &M) {
 #endif
   skip_nozero = getenv("AFL_LLVM_SKIP_NEVERZERO");
 
-  unsigned PrevLocSize;
+  unsigned PrevLocSize = 0;
 
   char *ngram_size_str = getenv("AFL_LLVM_NGRAM_SIZE");
   if (!ngram_size_str) ngram_size_str = getenv("AFL_NGRAM_SIZE");
@@ -216,9 +216,6 @@ bool AFLCoverage::runOnModule(Module &M) {
   if (ngram_size) PrevLocTy = VectorType::get(IntLocTy, PrevLocVecSize);
 #endif
 
-  if (ctx_str && ngram_size_str)
-    FATAL("you must decide between NGRAM and CTX instrumentation");
-
   /* Get globals for the SHM region and the previous location. Note that
      __afl_prev_loc is thread-local. */
 
@@ -437,8 +434,10 @@ bool AFLCoverage::runOnModule(Module &M) {
         PrevLocTrans = IRB.CreateXorReduce(PrevLoc);
       else
 #endif
-          if (ctx_str)
-        PrevLocTrans = IRB.CreateZExt(IRB.CreateXor(PrevLoc, PrevCtx), Int32Ty);
+        PrevLocTrans = PrevLoc;
+      if (ctx_str)
+        PrevLocTrans =
+            IRB.CreateZExt(IRB.CreateXor(PrevLocTrans, PrevCtx), Int32Ty);
       else
         PrevLocTrans = IRB.CreateZExt(PrevLoc, IRB.getInt32Ty());
 
@@ -452,7 +451,9 @@ bool AFLCoverage::runOnModule(Module &M) {
       if (ngram_size)
         MapPtrIdx = IRB.CreateGEP(
             MapPtr,
-            IRB.CreateZExt(IRB.CreateXor(PrevLocTrans, CurLoc), Int32Ty));
+            IRB.CreateZExt(
+                IRB.CreateXor(PrevLocTrans, IRB.CreateZExt(CurLoc, Int32Ty)),
+                Int32Ty));
       else
 #endif
         MapPtrIdx = IRB.CreateGEP(MapPtr, IRB.CreateXor(PrevLocTrans, CurLoc));