Merge pull request #15437 from ClickHouse/azat-ThreadPool-UAF-fixes

Destroy resources captured by lambda after `ThreadFromGlobalPool::join()`

src/Common/ThreadPool.cpp

@@ -234,10 +234,16 @@ void ThreadPoolImpl<Thread>::worker(typename std::list<Thread>::iterator thread_
                     std::is_same_v<Thread, std::thread> ? CurrentMetrics::GlobalThreadActive : CurrentMetrics::LocalThreadActive);
 
                 job();
+                /// job should be reseted before decrementing scheduled_jobs to
+                /// ensure that the Job destroyed before wait() returns.
                 job = {};
             }
             catch (...)
             {
+                /// job should be reseted before decrementing scheduled_jobs to
+                /// ensure that the Job destroyed before wait() returns.
+                job = {};
+
                 {
                     std::unique_lock lock(mutex);
                     if (!first_exception)

src/Common/ThreadPool.h

@@ -11,6 +11,7 @@
 
 #include <Poco/Event.h>
 #include <Common/ThreadStatus.h>
+#include <ext/scope_guard.h>
 
 
 /** Very simple thread pool similar to boost::threadpool.
@@ -161,21 +162,19 @@ public:
         GlobalThreadPool::instance().scheduleOrThrow([
             state = state,
             func = std::forward<Function>(func),
-            args = std::make_tuple(std::forward<Args>(args)...)]
+            args = std::make_tuple(std::forward<Args>(args)...)]() mutable /// mutable is needed to destroy capture
         {
-            try
-            {
-                /// Thread status holds raw pointer on query context, thus it always must be destroyed
-                /// before sending signal that permits to join this thread.
-                DB::ThreadStatus thread_status;
-                std::apply(func, args);
-            }
-            catch (...)
-            {
-                state->set();
-                throw;
-            }
-            state->set();
+            SCOPE_EXIT(state->set());
+
+            /// This moves are needed to destroy function and arguments before exit.
+            /// It will guarantee that after ThreadFromGlobalPool::join all captured params are destroyed.
+            auto function = std::move(func);
+            auto arguments = std::move(args);
+
+            /// Thread status holds raw pointer on query context, thus it always must be destroyed
+            /// before sending signal that permits to join this thread.
+            DB::ThreadStatus thread_status;
+            std::apply(function, arguments);
         });
     }
 

tests/queries/0_stateless/01505_pipeline_executor_UAF.sh

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+
+CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+. "$CURDIR"/../shell_config.sh
+
+# Regression for UAF in ThreadPool.
+# (Triggered under TSAN)
+for i in {1..10}; do
+    ${CLICKHOUSE_LOCAL} -q 'select * from numbers_mt(100000000) settings max_threads=100 FORMAT Null'
+    # Binding to specific CPU is not required, but this makes the test more reliable.
+    taskset --cpu-list 0 ${CLICKHOUSE_LOCAL} -q 'select * from numbers_mt(100000000) settings max_threads=100 FORMAT Null'
+done