From 176896c8ce07374e8ef2d6755929f2ddf7a5ca36 Mon Sep 17 00:00:00 2001
From: Nikita Mikhaylov <nikitamikhaylov@clickhouse.com>
Date: Tue, 17 Dec 2024 14:23:33 +0000
Subject: [PATCH] More fuzzers

---
 ..._function_state_deserialization_fuzzer.cpp |  4 ---
 .../fuzzers/compressed_buffer_fuzzer.cpp      | 28 +++++++++++++++--
 .../fuzzers/delta_decompress_fuzzer.cpp       | 30 ++++++++++++++++--
 .../double_delta_decompress_fuzzer.cpp        | 28 +++++++++++++++++
 .../fuzzers/encrypted_decompress_fuzzer.cpp   | 31 +++++++++++++++++--
 .../fuzzers/gcd_decompress_fuzzer.cpp         | 30 +++++++++++++++++-
 .../fuzzers/lz4_decompress_fuzzer.cpp         | 30 +++++++++++++++++-
 .../data_type_deserialization_fuzzer.cpp      |  6 ----
 src/Formats/fuzzers/format_fuzzer.cpp         |  4 ---
 .../fuzzers/execute_query_fuzzer.cpp          | 29 ++++++++++++-----
 10 files changed, 190 insertions(+), 30 deletions(-)

diff --git a/src/AggregateFunctions/fuzzers/aggregate_function_state_deserialization_fuzzer.cpp b/src/AggregateFunctions/fuzzers/aggregate_function_state_deserialization_fuzzer.cpp
index c4a96aa1ccc..8abdaf9953b 100644
--- a/src/AggregateFunctions/fuzzers/aggregate_function_state_deserialization_fuzzer.cpp
+++ b/src/AggregateFunctions/fuzzers/aggregate_function_state_deserialization_fuzzer.cpp
@@ -27,10 +27,6 @@ extern "C" int LLVMFuzzerInitialize(int *, char ***)
     if (context)
         return true;
 
-    /// The SharedContext depends on the Logger which is being destroyed by AutoLoggerShutdown (global variable)
-    /// And the GlobalContext depends on the SharedContext. So, this the SharedContext has to be static in order
-    /// to be destroyed last.
-    /// Addditionally, without it being static the shared context is destroyed on this function exit.
     static SharedContextHolder shared_context = Context::createShared();
     context = Context::createGlobal(shared_context.get());
     context->makeGlobalContext();
diff --git a/src/Compression/fuzzers/compressed_buffer_fuzzer.cpp b/src/Compression/fuzzers/compressed_buffer_fuzzer.cpp
index aea749900db..400e2ad0450 100644
--- a/src/Compression/fuzzers/compressed_buffer_fuzzer.cpp
+++ b/src/Compression/fuzzers/compressed_buffer_fuzzer.cpp
@@ -1,12 +1,36 @@
-#include <IO/ReadBufferFromMemory.h>
-#include <Compression/CompressedReadBuffer.h>
+#include <Common/Arena.h>
+#include <Common/CurrentThread.h>
 #include <Common/Exception.h>
+#include <Common/MemoryTracker.h>
+#include <Compression/CompressedReadBuffer.h>
+#include <IO/ReadBufferFromMemory.h>
+#include <Interpreters/Context.h>
 
+using namespace DB;
+ContextMutablePtr context;
+extern "C" int LLVMFuzzerInitialize(int *, char ***)
+{
+    if (context)
+        return true;
+
+    static SharedContextHolder shared_context = Context::createShared();
+    context = Context::createGlobal(shared_context.get());
+    context->makeGlobalContext();
+
+    MainThreadStatus::getInstance();
+
+    return 0;
+}
 
 extern "C" int LLVMFuzzerTestOneInput(const uint8_t * data, size_t size)
 {
     try
     {
+        total_memory_tracker.resetCounters();
+        total_memory_tracker.setHardLimit(1_GiB);
+        CurrentThread::get().memory_tracker.resetCounters();
+        CurrentThread::get().memory_tracker.setHardLimit(1_GiB);
+
         DB::ReadBufferFromMemory from(data, size);
         DB::CompressedReadBuffer in{from};
 
diff --git a/src/Compression/fuzzers/delta_decompress_fuzzer.cpp b/src/Compression/fuzzers/delta_decompress_fuzzer.cpp
index 267c6014461..ad4ad7c7eae 100644
--- a/src/Compression/fuzzers/delta_decompress_fuzzer.cpp
+++ b/src/Compression/fuzzers/delta_decompress_fuzzer.cpp
@@ -1,7 +1,12 @@
 #include <string>
 
-#include <Compression/ICompressionCodec.h>
-#include <IO/BufferWithOwnMemory.h>
+#include <Common/Arena.h>
+#include <Common/CurrentThread.h>
+#include <Common/Exception.h>
+#include <Common/MemoryTracker.h>
+#include <Compression/CompressedReadBuffer.h>
+#include <IO/ReadBufferFromMemory.h>
+#include <Interpreters/Context.h>
 
 namespace DB
 {
@@ -14,10 +19,31 @@ struct AuxiliaryRandomData
     size_t decompressed_size;
 };
 
+using namespace DB;
+ContextMutablePtr context;
+extern "C" int LLVMFuzzerInitialize(int *, char ***)
+{
+    if (context)
+        return true;
+
+    static SharedContextHolder shared_context = Context::createShared();
+    context = Context::createGlobal(shared_context.get());
+    context->makeGlobalContext();
+
+    MainThreadStatus::getInstance();
+
+    return 0;
+}
+
 extern "C" int LLVMFuzzerTestOneInput(const uint8_t * data, size_t size)
 {
     try
     {
+        total_memory_tracker.resetCounters();
+        total_memory_tracker.setHardLimit(1_GiB);
+        CurrentThread::get().memory_tracker.resetCounters();
+        CurrentThread::get().memory_tracker.setHardLimit(1_GiB);
+
         if (size < sizeof(AuxiliaryRandomData))
             return 0;
 
diff --git a/src/Compression/fuzzers/double_delta_decompress_fuzzer.cpp b/src/Compression/fuzzers/double_delta_decompress_fuzzer.cpp
index 22d89a1282b..4c463969115 100644
--- a/src/Compression/fuzzers/double_delta_decompress_fuzzer.cpp
+++ b/src/Compression/fuzzers/double_delta_decompress_fuzzer.cpp
@@ -1,7 +1,14 @@
 #include <string>
 
+#include <Common/Arena.h>
+#include <Common/CurrentThread.h>
+#include <Common/Exception.h>
+#include <Common/MemoryTracker.h>
+#include <Compression/CompressedReadBuffer.h>
 #include <Compression/ICompressionCodec.h>
 #include <IO/BufferWithOwnMemory.h>
+#include <IO/ReadBufferFromMemory.h>
+#include <Interpreters/Context.h>
 
 namespace DB
 {
@@ -14,10 +21,31 @@ struct AuxiliaryRandomData
     size_t decompressed_size;
 };
 
+using namespace DB;
+ContextMutablePtr context;
+extern "C" int LLVMFuzzerInitialize(int *, char ***)
+{
+    if (context)
+        return true;
+
+    static SharedContextHolder shared_context = Context::createShared();
+    context = Context::createGlobal(shared_context.get());
+    context->makeGlobalContext();
+
+    MainThreadStatus::getInstance();
+
+    return 0;
+}
+
 extern "C" int LLVMFuzzerTestOneInput(const uint8_t * data, size_t size)
 {
     try
     {
+        total_memory_tracker.resetCounters();
+        total_memory_tracker.setHardLimit(1_GiB);
+        CurrentThread::get().memory_tracker.resetCounters();
+        CurrentThread::get().memory_tracker.setHardLimit(1_GiB);
+
         if (size < sizeof(AuxiliaryRandomData))
             return 0;
 
diff --git a/src/Compression/fuzzers/encrypted_decompress_fuzzer.cpp b/src/Compression/fuzzers/encrypted_decompress_fuzzer.cpp
index 4131ab43c59..68583198646 100644
--- a/src/Compression/fuzzers/encrypted_decompress_fuzzer.cpp
+++ b/src/Compression/fuzzers/encrypted_decompress_fuzzer.cpp
@@ -3,9 +3,16 @@
 #include <memory>
 #include <string>
 
-#include <Compression/ICompressionCodec.h>
+#include <Common/Arena.h>
+#include <Common/CurrentThread.h>
+#include <Common/Exception.h>
+#include <Common/MemoryTracker.h>
+#include <Compression/CompressedReadBuffer.h>
 #include <Compression/CompressionCodecEncrypted.h>
+#include <Compression/ICompressionCodec.h>
 #include <IO/BufferWithOwnMemory.h>
+#include <IO/ReadBufferFromMemory.h>
+#include <Interpreters/Context.h>
 #include <Poco/DOM/AutoPtr.h>
 #include <Poco/DOM/Document.h>
 #include <Poco/DOM/Element.h>
@@ -13,13 +20,28 @@
 #include <Poco/NumericString.h>
 #include <Poco/Util/AbstractConfiguration.h>
 #include <Poco/Util/XMLConfiguration.h>
-#include "Common/Exception.h"
 
 inline DB::CompressionCodecPtr getCompressionCodecEncrypted(DB::EncryptionMethod Method)
 {
     return std::make_shared<DB::CompressionCodecEncrypted>(Method);
 }
 
+using namespace DB;
+ContextMutablePtr context;
+extern "C" int LLVMFuzzerInitialize(int *, char ***)
+{
+    if (context)
+        return true;
+
+    static SharedContextHolder shared_context = Context::createShared();
+    context = Context::createGlobal(shared_context.get());
+    context->makeGlobalContext();
+
+    MainThreadStatus::getInstance();
+
+    return 0;
+}
+
 namespace
 {
 
@@ -274,6 +296,11 @@ extern "C" int LLVMFuzzerTestOneInput(const uint8_t * data, size_t size)
 {
     try
     {
+        total_memory_tracker.resetCounters();
+        total_memory_tracker.setHardLimit(1_GiB);
+        CurrentThread::get().memory_tracker.resetCounters();
+        CurrentThread::get().memory_tracker.setHardLimit(1_GiB);
+
         XMLGenerator generator(data, size);
 
         generator.generate();
diff --git a/src/Compression/fuzzers/gcd_decompress_fuzzer.cpp b/src/Compression/fuzzers/gcd_decompress_fuzzer.cpp
index f8b8823ed09..d810e3cf8ad 100644
--- a/src/Compression/fuzzers/gcd_decompress_fuzzer.cpp
+++ b/src/Compression/fuzzers/gcd_decompress_fuzzer.cpp
@@ -1,9 +1,32 @@
 #include <iostream>
 #include <string>
 
+#include <Common/Arena.h>
+#include <Common/CurrentThread.h>
+#include <Common/Exception.h>
+#include <Common/MemoryTracker.h>
+#include <Compression/CompressedReadBuffer.h>
 #include <Compression/ICompressionCodec.h>
 #include <IO/BufferWithOwnMemory.h>
-#include "base/types.h"
+#include <IO/ReadBufferFromMemory.h>
+#include <Interpreters/Context.h>
+#include <base/types.h>
+
+using namespace DB;
+ContextMutablePtr context;
+extern "C" int LLVMFuzzerInitialize(int *, char ***)
+{
+    if (context)
+        return true;
+
+    static SharedContextHolder shared_context = Context::createShared();
+    context = Context::createGlobal(shared_context.get());
+    context->makeGlobalContext();
+
+    MainThreadStatus::getInstance();
+
+    return 0;
+}
 
 namespace DB
 {
@@ -19,6 +42,11 @@ struct AuxiliaryRandomData
 extern "C" int LLVMFuzzerTestOneInput(const uint8_t * data, size_t size)
 try
 {
+    total_memory_tracker.resetCounters();
+    total_memory_tracker.setHardLimit(1_GiB);
+    CurrentThread::get().memory_tracker.resetCounters();
+    CurrentThread::get().memory_tracker.setHardLimit(1_GiB);
+
     if (size < sizeof(AuxiliaryRandomData))
         return 0;
 
diff --git a/src/Compression/fuzzers/lz4_decompress_fuzzer.cpp b/src/Compression/fuzzers/lz4_decompress_fuzzer.cpp
index e2275990b72..d249ba9605d 100644
--- a/src/Compression/fuzzers/lz4_decompress_fuzzer.cpp
+++ b/src/Compression/fuzzers/lz4_decompress_fuzzer.cpp
@@ -1,8 +1,31 @@
 #include <string>
 
+#include <Common/Arena.h>
+#include <Common/CurrentThread.h>
+#include <Common/Exception.h>
+#include <Common/MemoryTracker.h>
+#include <Compression/CompressedReadBuffer.h>
 #include <Compression/ICompressionCodec.h>
-#include <IO/BufferWithOwnMemory.h>
 #include <Compression/LZ4_decompress_faster.h>
+#include <IO/BufferWithOwnMemory.h>
+#include <IO/ReadBufferFromMemory.h>
+#include <Interpreters/Context.h>
+
+using namespace DB;
+ContextMutablePtr context;
+extern "C" int LLVMFuzzerInitialize(int *, char ***)
+{
+    if (context)
+        return true;
+
+    static SharedContextHolder shared_context = Context::createShared();
+    context = Context::createGlobal(shared_context.get());
+    context->makeGlobalContext();
+
+    MainThreadStatus::getInstance();
+
+    return 0;
+}
 
 namespace DB
 {
@@ -19,6 +42,11 @@ extern "C" int LLVMFuzzerTestOneInput(const uint8_t * data, size_t size)
 {
     try
     {
+        total_memory_tracker.resetCounters();
+        total_memory_tracker.setHardLimit(1_GiB);
+        CurrentThread::get().memory_tracker.resetCounters();
+        CurrentThread::get().memory_tracker.setHardLimit(1_GiB);
+
         if (size < sizeof(AuxiliaryRandomData) + LZ4::ADDITIONAL_BYTES_AT_END_OF_BUFFER)
             return 0;
 
diff --git a/src/DataTypes/fuzzers/data_type_deserialization_fuzzer.cpp b/src/DataTypes/fuzzers/data_type_deserialization_fuzzer.cpp
index 79b593140bd..a82d3c7ccc9 100644
--- a/src/DataTypes/fuzzers/data_type_deserialization_fuzzer.cpp
+++ b/src/DataTypes/fuzzers/data_type_deserialization_fuzzer.cpp
@@ -15,18 +15,12 @@
 
 using namespace DB;
 
-
 ContextMutablePtr context;
-
 extern "C" int LLVMFuzzerInitialize(int *, char ***)
 {
     if (context)
         return true;
 
-    /// The SharedContext depends on the Logger which is being destroyed by AutoLoggerShutdown (global variable)
-    /// And the GlobalContext depends on the SharedContext. So, this the SharedContext has to be static in order
-    /// to be destroyed last.
-    /// Addditionally, without it being static the shared context is destroyed on this function exit.
     static SharedContextHolder shared_context = Context::createShared();
     context = Context::createGlobal(shared_context.get());
     context->makeGlobalContext();
diff --git a/src/Formats/fuzzers/format_fuzzer.cpp b/src/Formats/fuzzers/format_fuzzer.cpp
index d81cb922d6d..ab79c0dbbd5 100644
--- a/src/Formats/fuzzers/format_fuzzer.cpp
+++ b/src/Formats/fuzzers/format_fuzzer.cpp
@@ -29,10 +29,6 @@ extern "C" int LLVMFuzzerInitialize(int *, char ***)
     if (context)
         return true;
 
-    /// The SharedContext depends on the Logger which is being destroyed by AutoLoggerShutdown (global variable)
-    /// And the GlobalContext depends on the SharedContext. So, this the SharedContext has to be static in order
-    /// to be destroyed last.
-    /// Addditionally, without it being static the shared context is destroyed on this function exit.
     static SharedContextHolder shared_context = Context::createShared();
     context = Context::createGlobal(shared_context.get());
     context->makeGlobalContext();
diff --git a/src/Interpreters/fuzzers/execute_query_fuzzer.cpp b/src/Interpreters/fuzzers/execute_query_fuzzer.cpp
index 1ce9b85dacd..f5e86f581c5 100644
--- a/src/Interpreters/fuzzers/execute_query_fuzzer.cpp
+++ b/src/Interpreters/fuzzers/execute_query_fuzzer.cpp
@@ -1,7 +1,8 @@
 #include <Interpreters/executeQuery.h>
 #include <Interpreters/Context.h>
 #include <Interpreters/registerInterpreters.h>
-#include "Processors/Executors/PullingPipelineExecutor.h"
+#include <Processors/Executors/PullingPipelineExecutor.h>
+#include <Processors/Executors/PushingPipelineExecutor.h>
 
 #include <Databases/registerDatabases.h>
 #include <Functions/registerFunctions.h>
@@ -22,14 +23,12 @@ extern "C" int LLVMFuzzerInitialize(int *, char ***)
     if (context)
         return true;
 
-    /// The SharedContext depends on the Logger which is being destroyed by AutoLoggerShutdown (global variable)
-    /// And the GlobalContext depends on the SharedContext. So, this the SharedContext has to be static in order
-    /// to be destroyed last.
-    /// Addditionally, without it being static the shared context is destroyed on this function exit.
     static SharedContextHolder shared_context = Context::createShared();
     context = Context::createGlobal(shared_context.get());
     context->makeGlobalContext();
 
+    MainThreadStatus::getInstance();
+
     registerInterpreters();
     registerFunctions();
     registerAggregateFunctions();
@@ -47,13 +46,27 @@ extern "C" int LLVMFuzzerTestOneInput(const uint8_t * data, size_t size)
 {
     try
     {
+        total_memory_tracker.resetCounters();
+        total_memory_tracker.setHardLimit(1_GiB);
+        CurrentThread::get().memory_tracker.resetCounters();
+        CurrentThread::get().memory_tracker.setHardLimit(1_GiB);
+
         std::string input = std::string(reinterpret_cast<const char*>(data), size);
 
         auto io = DB::executeQuery(input, context, QueryFlags{ .internal = true }, QueryProcessingStage::Complete).second;
 
-        PullingPipelineExecutor executor(io.pipeline);
-        Block res;
-        while (!res && executor.pull(res));
+        /// Execute only SELECTs
+        if (io.pipeline.pulling())
+        {
+            PullingPipelineExecutor executor(io.pipeline);
+            Block res;
+            while (!res && executor.pull(res));
+        }
+        /// We don't want to execute it and thus need to finish it properly.
+        else
+        {
+            io.onCancelOrConnectionLoss();
+        }
     }
     catch (...)
     {