From 9bd4616f6f25ea08fb40c7acff22ba293d09b67f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joel=20Ho=CC=88ner?= <athre0z@zyantific.com>
Date: Mon, 27 Nov 2017 00:32:15 +0100
Subject: [PATCH] Added `ZYDIS_FUZZ_AFL_FAST` CMake switch

- Enables AFL persistent mode and reduces prints in `ZydisFuzzIn`
---
 CMakeLists.txt         |  6 ++++++
 examples/ZydisFuzzIn.c | 34 ++++++++++++++++++++++++++++------
 2 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 7ca73fa..7c35624 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -35,6 +35,9 @@ option(ZYDIS_BUILD_EXAMPLES
 option(ZYDIS_BUILD_TOOLS 
     "Build tools" 
     ON)
+option(ZYDIS_FUZZ_AFL_FAST
+	"Enables AFL persistent mode and reduces prints in ZydisFuzzIn"
+	OFF)
 option(ZYDIS_DEV_MODE
     "Enable developer mode (-Wall, -Werror, ...)"
     OFF)
@@ -156,6 +159,9 @@ if (ZYDIS_BUILD_EXAMPLES)
         target_link_libraries("ZydisFuzzIn" "Zydis")
         set_target_properties("FormatterHooks" PROPERTIES FOLDER "Examples")
         target_compile_definitions("ZydisFuzzIn" PRIVATE "_CRT_SECURE_NO_WARNINGS")
+		if (ZYDIS_FUZZ_AFL_FAST)
+			target_compile_definitions("ZydisFuzzIn" PRIVATE "ZYDIS_FUZZ_AFL_FAST")
+		endif ()
 
         add_executable("ZydisPerfTest" "examples/ZydisPerfTest.c")
         target_link_libraries("ZydisPerfTest" "Zydis")
diff --git a/examples/ZydisFuzzIn.c b/examples/ZydisFuzzIn.c
index 52a95bb..7913ef3 100644
--- a/examples/ZydisFuzzIn.c
+++ b/examples/ZydisFuzzIn.c
@@ -53,6 +53,8 @@ typedef struct ZydisFuzzControlBlock_
 /* Entry point                                                                                    */
 /* ============================================================================================== */
 
+int doIteration();
+
 int main()
 {
     if (ZydisGetVersion() != ZYDIS_VERSION)
@@ -61,10 +63,30 @@ int main()
         return EXIT_FAILURE;
     }
 
+#ifdef ZYDIS_FUZZ_AFL_FAST
+    int finalRet;
+    while (__AFL_LOOP(1000))
+    {
+        finalRet = doIteration();
+    }
+    return finalRet;
+#else
+    return doIteration();
+#endif
+}
+
+#ifdef ZYDIS_FUZZ_AFL_FAST
+#   define ZYDIS_MAYBE_FPUTS(x, y)
+#else
+#   define ZYDIS_MAYBE_FPUTS(x, y) fputs(x, y)
+#endif
+
+int doIteration()
+{
     ZydisFuzzControlBlock controlBlock;
     if (fread(&controlBlock, 1, sizeof(controlBlock), stdin) != sizeof(controlBlock))
     {
-        fputs("not enough bytes to fuzz\n", stderr);
+        ZYDIS_MAYBE_FPUTS("not enough bytes to fuzz\n", stderr);
         return EXIT_FAILURE;
     }
     controlBlock.string[ZYDIS_ARRAY_SIZE(controlBlock.string) - 1] = 0;
@@ -73,7 +95,7 @@ int main()
     if (!ZYDIS_SUCCESS(
         ZydisDecoderInit(&decoder, controlBlock.machineMode, controlBlock.addressWidth)))
     {
-        fputs("Failed to initialize decoder\n", stderr);
+        ZYDIS_MAYBE_FPUTS("Failed to initialize decoder\n", stderr);
         return EXIT_FAILURE;
     }
     for (ZydisDecoderMode mode = 0; mode <= ZYDIS_DECODER_MODE_MAX_VALUE; ++mode)
@@ -81,7 +103,7 @@ int main()
         if (!ZYDIS_SUCCESS(
             ZydisDecoderEnableMode(&decoder, mode, controlBlock.decoderMode[mode] ? 1 : 0)))
         {
-            fputs("Failed to adjust decoder-mode\n", stderr);
+            ZYDIS_MAYBE_FPUTS("Failed to adjust decoder-mode\n", stderr);
             return EXIT_FAILURE;
         }
     }
@@ -89,7 +111,7 @@ int main()
     ZydisFormatter formatter;
     if (!ZYDIS_SUCCESS(ZydisFormatterInit(&formatter, controlBlock.formatterStyle)))
     {
-        fputs("Failed to initialize instruction-formatter\n", stderr);
+        ZYDIS_MAYBE_FPUTS("Failed to initialize instruction-formatter\n", stderr);
         return EXIT_FAILURE;
     }
     for (ZydisFormatterProperty prop = 0; prop <= ZYDIS_FORMATTER_PROP_MAX_VALUE; ++prop)
@@ -107,7 +129,7 @@ int main()
         if (!ZYDIS_SUCCESS(ZydisFormatterSetProperty(&formatter, prop, 
             controlBlock.formatterProperties[prop])))
         {
-            fputs("Failed to set formatter-attribute\n", stderr);
+            ZYDIS_MAYBE_FPUTS("Failed to set formatter-attribute\n", stderr);
             return EXIT_FAILURE;
         }
     }
@@ -142,7 +164,7 @@ int main()
         }
     } while (numBytesRead == sizeof(readBuf));
 
-    return 0;
+    return EXIT_SUCCESS;
 }
 
 /* ============================================================================================== */