From 487bc0fba32fcd81a44aeb974c2470dceac8a33f Mon Sep 17 00:00:00 2001
From: kssenii <sumarokovakseniia@mail.ru>
Date: Mon, 5 Sep 2022 23:10:03 +0200
Subject: [PATCH] Fix heap use after free

---
 programs/local/LocalServer.cpp | 11 +++++------
 src/Client/ClientBase.h        |  8 +++++---
 src/Client/LocalConnection.cpp |  3 ---
 src/Client/LocalConnection.h   |  1 -
 src/Interpreters/Context.cpp   | 14 +++++++++++---
 5 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/programs/local/LocalServer.cpp b/programs/local/LocalServer.cpp
index fa43c9f1283..fa84199ea92 100644
--- a/programs/local/LocalServer.cpp
+++ b/programs/local/LocalServer.cpp
@@ -227,6 +227,10 @@ void LocalServer::cleanup()
             global_context.reset();
         }
 
+        /// thread status should be destructed before shared context because it relies on process list.
+        thread_status.reset();
+        shared_context.reset();
+
         status.reset();
 
         // Delete the temporary directory if needed.
@@ -366,12 +370,7 @@ int LocalServer::main(const std::vector<std::string> & /*args*/)
 try
 {
     UseSSL use_ssl;
-    ThreadStatus thread_status;
-    SCOPE_EXIT_SAFE({
-        /// Context should not live longer than thread_status.
-        global_context.reset();
-        shared_context.reset();
-    });
+    thread_status.emplace();
 
     StackTrace::setShowAddresses(config().getBool("show_addresses_in_stack_traces", true));
 
diff --git a/src/Client/ClientBase.h b/src/Client/ClientBase.h
index 6b19c1b8e02..219d35d87cd 100644
--- a/src/Client/ClientBase.h
+++ b/src/Client/ClientBase.h
@@ -176,9 +176,6 @@ protected:
     bool stderr_is_a_tty = false; /// stderr is a terminal.
     uint64_t terminal_width = 0;
 
-    ServerConnectionPtr connection;
-    ConnectionParameters connection_parameters;
-
     String format; /// Query results output format.
     bool select_into_file = false; /// If writing result INTO OUTFILE. It affects progress rendering.
     bool select_into_file_and_stdout = false; /// If writing result INTO OUTFILE AND STDOUT. It affects progress rendering.
@@ -199,6 +196,11 @@ protected:
     SharedContextHolder shared_context;
     ContextMutablePtr global_context;
 
+    std::optional<ThreadStatus> thread_status;
+
+    ServerConnectionPtr connection;
+    ConnectionParameters connection_parameters;
+
     /// Buffer that reads from stdin in batch mode.
     ReadBufferFromFileDescriptor std_in{STDIN_FILENO};
     /// Console output.
diff --git a/src/Client/LocalConnection.cpp b/src/Client/LocalConnection.cpp
index b10e24f1ae4..7ac68324915 100644
--- a/src/Client/LocalConnection.cpp
+++ b/src/Client/LocalConnection.cpp
@@ -31,9 +31,6 @@ LocalConnection::LocalConnection(ContextPtr context_, bool send_progress_, bool
     /// Authenticate and create a context to execute queries.
     session.authenticate("default", "", Poco::Net::SocketAddress{});
     session.makeSessionContext();
-
-    if (!CurrentThread::isInitialized())
-        thread_status.emplace();
 }
 
 LocalConnection::~LocalConnection()
diff --git a/src/Client/LocalConnection.h b/src/Client/LocalConnection.h
index 7967874d11f..7a1a73006ac 100644
--- a/src/Client/LocalConnection.h
+++ b/src/Client/LocalConnection.h
@@ -156,7 +156,6 @@ private:
     String description = "clickhouse-local";
 
     std::optional<LocalQueryState> state;
-    std::optional<ThreadStatus> thread_status;
 
     /// Last "server" packet.
     std::optional<UInt64> next_packet_type;
diff --git a/src/Interpreters/Context.cpp b/src/Interpreters/Context.cpp
index d39c39cdb15..87c2a997afa 100644
--- a/src/Interpreters/Context.cpp
+++ b/src/Interpreters/Context.cpp
@@ -143,7 +143,7 @@ namespace ErrorCodes
 /** Set of known objects (environment), that could be used in query.
   * Shared (global) part. Order of members (especially, order of destruction) is very important.
   */
-struct ContextSharedPart
+struct ContextSharedPart : boost::noncopyable
 {
     Poco::Logger * log = &Poco::Logger::get("Context");
 
@@ -314,11 +314,19 @@ struct ContextSharedPart
 
     ~ContextSharedPart()
     {
-        /// Wait for thread pool for background writes,
-        /// since it may use per-user MemoryTracker which will be destroyed here.
         try
         {
+            /// Wait for thread pool for background writes,
+            /// since it may use per-user MemoryTracker which will be destroyed here.
             IObjectStorage::getThreadPoolWriter().wait();
+            /// Make sure that threadpool is destructed before this->process_list
+            /// because thread_status, which was created for threads inside threadpool,
+            /// relies on it.
+            if (load_marks_threadpool)
+            {
+                load_marks_threadpool->wait();
+                load_marks_threadpool.reset();
+            }
         }
         catch (...)
         {