Fix bad_weak_ptr when close a ClientConnection during construction

BewareMyPower · BewareMyPower · commit ba714ea2630f · 2023-11-17T16:05:41.000+08:00
Fixes #348 Fixes #349 ### Motivation When `close` is called in `ClientConnection`'s constructor, `shared_from_this()` will be called, which results in a `std::bad_weak_ptr` error. This error does not happen before #317 because `shared_from_this()` could only be called when the `producers` or `consumers` field is not empty. ### Modifications Add a `ResultException` and throw it if there is a failure in `ClientConnection`'s constructor and catch it in `ConnectionPool::getConnectionAsync`. Then retrieve the result and return the failed future. Since `close` is now always called on a constructed `ClientConnection` object, remove the 2nd parameter. In addition, check `authentication_` even for non-TLS URLs. Otherwise, the null authentication will be used to construct `CommandConnect`. Add `testInvalidPlugin` and `testTlsConfigError` to verify the changes.
diff --git a/lib/ClientConnection.cc b/lib/ClientConnection.cc
@@ -187,6 +187,10 @@ ClientConnection::ClientConnection(const std::string& logicalAddress, const std:
       pool_(pool),
       poolIndex_(poolIndex) {
     LOG_INFO(cnxString_ << "Create ClientConnection, timeout=" << clientConfiguration.getConnectionTimeout());
+    if (!authentication_) {
+        LOG_ERROR("Invalid authentication plugin");
+        throw ResultException{ResultAuthenticationError};
+    }
     if (clientConfiguration.isUseTls()) {
 #if BOOST_VERSION >= 105400
         boost::asio::ssl::context ctx(boost::asio::ssl::context::tlsv12_client);
@@ -207,20 +211,13 @@ ClientConnection::ClientConnection(const std::string& logicalAddress, const std:
                     ctx.load_verify_file(trustCertFilePath);
                 } else {
                     LOG_ERROR(trustCertFilePath << ": No such trustCertFile");
-                    close(ResultAuthenticationError, false);
-                    return;
+                    throw ResultException{ResultAuthenticationError};
                 }
             } else {
                 ctx.set_default_verify_paths();
             }
         }
 
-        if (!authentication_) {
-            LOG_ERROR("Invalid authentication plugin");
-            close(ResultAuthenticationError, false);
-            return;
-        }
-
         std::string tlsCertificates = clientConfiguration.getTlsCertificateFilePath();
         std::string tlsPrivateKey = clientConfiguration.getTlsPrivateKeyFilePath();
 
@@ -231,13 +228,11 @@ ClientConnection::ClientConnection(const std::string& logicalAddress, const std:
             tlsPrivateKey = authData->getTlsPrivateKey();
             if (!file_exists(tlsCertificates)) {
                 LOG_ERROR(tlsCertificates << ": No such tlsCertificates");
-                close(ResultAuthenticationError, false);
-                return;
+                throw ResultException{ResultAuthenticationError};
             }
             if (!file_exists(tlsCertificates)) {
                 LOG_ERROR(tlsCertificates << ": No such tlsCertificates");
-                close(ResultAuthenticationError, false);
-                return;
+                throw ResultException{ResultAuthenticationError};
             }
             ctx.use_private_key_file(tlsPrivateKey, boost::asio::ssl::context::pem);
             ctx.use_certificate_file(tlsCertificates, boost::asio::ssl::context::pem);
@@ -1260,7 +1255,7 @@ void ClientConnection::handleConsumerStatsTimeout(const boost::system::error_cod
     startConsumerStatsTimer(consumerStatsRequests);
 }
 
-void ClientConnection::close(Result result, bool detach) {
+void ClientConnection::close(Result result) {
     Lock lock(mutex_);
     if (isClosed()) {
         return;
@@ -1320,9 +1315,7 @@ void ClientConnection::close(Result result, bool detach) {
         LOG_INFO(cnxString_ << "Connection disconnected (refCnt: " << refCount << ")");
     }
     // Remove the connection from the pool before completing any promise
-    if (detach) {
-        pool_.remove(logicalAddress_ + "-" + std::to_string(poolIndex_), this);
-    }
+    pool_.remove(logicalAddress_ + "-" + std::to_string(poolIndex_), this);
 
     auto self = shared_from_this();
     for (ProducersMap::iterator it = producers.begin(); it != producers.end(); ++it) {
diff --git a/lib/ClientConnection.h b/lib/ClientConnection.h
@@ -32,6 +32,7 @@
 #include <boost/asio/strand.hpp>
 #include <boost/optional.hpp>
 #include <deque>
+#include <exception>
 #include <functional>
 #include <memory>
 #include <string>
@@ -146,13 +147,8 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
      * Close the connection.
      *
      * @param result all pending futures will complete with this result
-     * @param detach remove it from the pool if it's true
-     *
-     * `detach` should only be false when:
-     * 1. Before the connection is put into the pool, i.e. during the construction.
-     * 2. When the connection pool is closed
      */
-    void close(Result result = ResultConnectError, bool detach = true);
+    void close(Result result = ResultConnectError);
 
     bool isClosed() const;
 
diff --git a/lib/ConnectionPool.cc b/lib/ConnectionPool.cc
@@ -24,6 +24,7 @@
 #include "ClientConnection.h"
 #include "ExecutorService.h"
 #include "LogUtils.h"
+#include "ResultUtils.h"
 
 using boost::asio::ip::tcp;
 namespace ssl = boost::asio::ssl;
@@ -53,21 +54,17 @@ bool ConnectionPool::close() {
     for (auto cnxIt = pool_.begin(); cnxIt != pool_.end(); cnxIt++) {
         auto& cnx = cnxIt->second;
         if (cnx) {
-            // The 2nd argument is false because removing a value during the iteration will cause segfault
-            cnx->close(ResultDisconnected, false);
+            cnx->close(ResultDisconnected);
         }
     }
     pool_.clear();
     return true;
 }
 
-Future<Result, ClientConnectionWeakPtr> ConnectionPool::getConnectionAsync(const std::string& logicalAddress,
-                                                                           const std::string& physicalAddress,
-                                                                           size_t keySuffix) {
+auto ConnectionPool::getConnectionAsync(const std::string& logicalAddress, const std::string& physicalAddress,
+                                        size_t keySuffix) -> ConnectionFuture {
     if (closed_) {
-        Promise<Result, ClientConnectionWeakPtr> promise;
-        promise.setFailed(ResultAlreadyClosed);
-        return promise.getFuture();
+        return ConnectionFuture::failure(ResultAlreadyClosed);
     }
 
     std::unique_lock<std::recursive_mutex> lock(mutex_);
@@ -99,12 +96,12 @@ Future<Result, ClientConnectionWeakPtr> ConnectionPool::getConnectionAsync(const
         cnx.reset(new ClientConnection(logicalAddress, physicalAddress, executorProvider_->get(keySuffix),
                                        clientConfiguration_, authentication_, clientVersion_, *this,
                                        keySuffix));
+    } catch (const ResultException& e) {
+        return ConnectionFuture::failure(e.result());
     } catch (const std::runtime_error& e) {
         lock.unlock();
         LOG_ERROR("Failed to create connection: " << e.what())
-        Promise<Result, ClientConnectionWeakPtr> promise;
-        promise.setFailed(ResultConnectError);
-        return promise.getFuture();
+        return ConnectionFuture::failure(ResultConnectError);
     }
 
     LOG_INFO("Created connection for " << key);
diff --git a/lib/ConnectionPool.h b/lib/ConnectionPool.h
@@ -53,6 +53,8 @@ class PULSAR_PUBLIC ConnectionPool {
 
     void remove(const std::string& key, ClientConnection* value);
 
+    using ConnectionFuture = Future<Result, ClientConnectionWeakPtr>;
+
     /**
      * Get a connection from the pool.
      * <p>
@@ -73,16 +75,15 @@ class PULSAR_PUBLIC ConnectionPool {
      * @param keySuffix the key suffix to choose which connection on the same broker
      * @return a future that will produce the ClientCnx object
      */
-    Future<Result, ClientConnectionWeakPtr> getConnectionAsync(const std::string& logicalAddress,
-                                                               const std::string& physicalAddress,
-                                                               size_t keySuffix);
+    ConnectionFuture getConnectionAsync(const std::string& logicalAddress, const std::string& physicalAddress,
+                                        size_t keySuffix);
 
-    Future<Result, ClientConnectionWeakPtr> getConnectionAsync(const std::string& logicalAddress,
-                                                               const std::string& physicalAddress) {
+    ConnectionFuture getConnectionAsync(const std::string& logicalAddress,
+                                        const std::string& physicalAddress) {
         return getConnectionAsync(logicalAddress, physicalAddress, generateRandomIndex());
     }
 
-    Future<Result, ClientConnectionWeakPtr> getConnectionAsync(const std::string& address) {
+    ConnectionFuture getConnectionAsync(const std::string& address) {
         return getConnectionAsync(address, address);
     }
 
diff --git a/lib/Future.h b/lib/Future.h
@@ -116,6 +116,8 @@ class Future {
 
     Result get(Type &result) { return state_->get(result); }
 
+    static Future<Result, Type> failure(Result result);
+
    private:
     InternalStatePtr<Result, Type> state_;
 
@@ -144,6 +146,13 @@ class Promise {
     const InternalStatePtr<Result, Type> state_;
 };
 
+template <typename Result, typename Type>
+inline Future<Result, Type> Future<Result, Type>::failure(Result result) {
+    Promise<Result, Type> promise;
+    promise.setFailed(result);
+    return promise.getFuture();
+}
+
 }  // namespace pulsar
 
 #endif
diff --git a/lib/ResultUtils.h b/lib/ResultUtils.h
@@ -20,10 +20,22 @@
 
 #include <pulsar/Result.h>
 
+#include <exception>
+
 namespace pulsar {
 
 inline bool isResultRetryable(Result result) {
     return result == ResultRetryable || result == ResultDisconnected;
 }
 
+class ResultException : public std::exception {
+   public:
+    ResultException(Result result) : result_(result) {}
+    Result result() const noexcept { return result_; }
+    const char* what() const noexcept { return strResult(result_); }
+
+   private:
+    const Result result_;
+};
+
 }  // namespace pulsar
diff --git a/tests/AuthPluginTest.cc b/tests/AuthPluginTest.cc
@@ -581,3 +581,19 @@ TEST(AuthPluginTest, testOauth2Failure) {
     ASSERT_EQ(client5.createProducer(topic, producer), ResultAuthenticationError);
     client5.close();
 }
+
+TEST(AuthPluginTest, testInvalidPlugin) {
+    Client client("pulsar://localhost:6650", ClientConfiguration{}.setAuth(AuthFactory::create("invalid")));
+    Producer producer;
+    ASSERT_EQ(ResultAuthenticationError, client.createProducer("my-topic", producer));
+    client.close();
+}
+
+TEST(AuthPluginTest, testTlsConfigError) {
+    Client client(serviceUrlTls, ClientConfiguration{}
+                                     .setAuth(AuthTls::create(clientPublicKeyPath, clientPrivateKeyPath))
+                                     .setTlsTrustCertsFilePath("invalid"));
+    Producer producer;
+    ASSERT_EQ(ResultAuthenticationError, client.createProducer("my-topic", producer));
+    client.close();
+}
