Merge pull request #5871 from yandex/fix-clickhouse-copier-heap-use-a…

…fter-free

Fix heap-use-after-free in ClusterCopier

(cherry picked from commit bc1e9df)

dbms/programs/copier/ClusterCopier.cpp

@@ -2196,6 +2196,10 @@ void ClusterCopierApp::mainImpl()
 
     copier->init();
     copier->process(ConnectionTimeouts::getTCPTimeoutsWithoutFailover(context->getSettingsRef()));
+
+    /// Reset ZooKeeper before removing ClusterCopier.
+    /// Otherwise zookeeper watch can call callback which use already removed ClusterCopier object.
+    context->resetZooKeeper();
 }
 
 

dbms/src/Interpreters/Context.cpp

@@ -1458,6 +1458,12 @@ zkutil::ZooKeeperPtr Context::getZooKeeper() const
     return shared->zookeeper;
 }
 
+void Context::resetZooKeeper() const
+{
+    std::lock_guard lock(shared->zookeeper_mutex);
+    shared->zookeeper.reset();
+}
+
 bool Context::hasZooKeeper() const
 {
     return getConfigRef().has("zookeeper");

dbms/src/Interpreters/Context.h

@@ -371,6 +371,8 @@ class Context
     std::shared_ptr<zkutil::ZooKeeper> getZooKeeper() const;
     /// Has ready or expired ZooKeeper
     bool hasZooKeeper() const;
+    /// Reset current zookeeper session. Do not create a new one.
+    void resetZooKeeper() const;
 
     /// Create a cache of uncompressed blocks of specified size. This can be done only once.
     void setUncompressedCache(size_t max_size_in_bytes);