From 31ed58674a3bd43c943b962711bbc8e2079c5f50 Mon Sep 17 00:00:00 2001
From: Alex Hultman <alexhultman@gmail.com>
Date: Fri, 7 Aug 2020 19:33:03 +0200
Subject: [PATCH] Improve PerMessageDeflate fuzzing

---
 fuzzing/PerMessageDeflate.cpp | 40 +++++++++++++++++++++++++++--------
 1 file changed, 31 insertions(+), 9 deletions(-)

diff --git a/fuzzing/PerMessageDeflate.cpp b/fuzzing/PerMessageDeflate.cpp
index 2900a88..3a5c96e 100644
--- a/fuzzing/PerMessageDeflate.cpp
+++ b/fuzzing/PerMessageDeflate.cpp
@@ -10,17 +10,39 @@
 
 #include "helpers.h"
 
-struct StaticData {
-    uWS::ZlibContext zlibContext;
-
-    uWS::InflationStream inflationStream;
-    uWS::DeflationStream deflationStream = uWS::DEDICATED_COMPRESSOR_3KB;
-} staticData;
-
 extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
 
+    /* First byte determines what compressor to use */
+    if (size < 1) {
+        return 0;
+    }
+
+    int compressors[] = {
+        uWS::DEDICATED_COMPRESSOR_3KB,
+        uWS::DEDICATED_COMPRESSOR_4KB,
+        uWS::DEDICATED_COMPRESSOR_8KB,
+        uWS::DEDICATED_COMPRESSOR_16KB,
+        uWS::DEDICATED_COMPRESSOR_32KB,
+        uWS::DEDICATED_COMPRESSOR_64KB,
+        uWS::DEDICATED_COMPRESSOR_128KB,
+        uWS::DEDICATED_COMPRESSOR_256KB
+    };
+
+    auto compressor = compressors[data[0] & 8];
+    data++;
+    size--;
+
+    /* If we could specify LARGE_BUFFER_SIZE small here we could force it to inflate in chunks,
+     * triggering more line coverage. Currently it is set to 16kb which is always too much */
+    struct StaticData {
+        uWS::DeflationStream deflationStream;
+        uWS::ZlibContext zlibContext;
+
+        uWS::InflationStream inflationStream;
+    } staticData = {compressor};
+
     /* Why is this padded? */
-    makeChunked(makePadded(data, size), size, [](const uint8_t *data, size_t size) {
+    makeChunked(makePadded(data, size), size, [&staticData](const uint8_t *data, size_t size) {
         auto [inflation, valid] = staticData.inflationStream.inflate(&staticData.zlibContext, std::string_view((char *) data, size), 256);
         if (inflation.length() > 256) {
             /* Cause ASAN to freak out */
@@ -28,7 +50,7 @@ extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
         }
     });
 
-    makeChunked(makePadded(data, size), size, [](const uint8_t *data, size_t size) {
+    makeChunked(makePadded(data, size), size, [&staticData](const uint8_t *data, size_t size) {
         /* Always reset */
         staticData.deflationStream.deflate(&staticData.zlibContext, std::string_view((char *) data, size), true);
     });