SERVER-51812 Allow DBClientConnection::connect() to specify helloOk when initializing a connection

6 files changed, 125 insertions, 0 deletions

diff --git a/src/mongo/client/dbclient_connection.cpp b/src/mongo/client/dbclient_connection.cpp
index ae95c3bbe02..49b581d5d2d 100644
--- a/src/mongo/client/dbclient_connection.cpp
+++ b/src/mongo/client/dbclient_connection.cpp
@@ -188,6 +188,12 @@ executor::RemoteCommandResponse initWireVersion(
     BSONObjBuilder bob;
     bob.append("isMaster", 1);
 
+    if (uri.isHelloOk()) {
+        // Attach "helloOk: true" to the initial handshake to indicate that the client supports the
+        // hello command.
+        bob.append("helloOk", true);
+    }
+
     *speculativeAuthType = auth::speculateAuth(&bob, uri, saslClientSession);
     if (!uri.getUser().empty()) {
         UserName user(uri.getUser(), uri.getAuthenticationDatabase());
diff --git a/src/mongo/client/mongo_uri.cpp b/src/mongo/client/mongo_uri.cpp
index e2b0af43b7c..e662d81cb7e 100644
--- a/src/mongo/client/mongo_uri.cpp
+++ b/src/mongo/client/mongo_uri.cpp
@@ -484,6 +484,20 @@ MongoURI MongoURI::parseImpl(StringData url) {
         }
     }
 
+    const auto helloOk = [&options]() -> boost::optional<bool> {
+        if (auto optIter = options.find("helloOk"); optIter != end(options)) {
+            if (auto value = optIter->second; value == "true") {
+                return true;
+            } else if (value == "false") {
+                return false;
+            } else {
+                uasserted(ErrorCodes::FailedToParse,
+                          "helloOk must be either \"true\" or \"false\"");
+            }
+        }
+        return boost::none;
+    }();
+
     transport::ConnectSSLMode sslMode = transport::kGlobalSSLMode;
     auto sslModeIter = std::find_if(options.begin(), options.end(), [](auto pred) {
         return pred.first == CaseInsensitiveString("ssl") ||
@@ -509,6 +523,7 @@ MongoURI MongoURI::parseImpl(StringData url) {
                     database,
                     std::move(retryWrites),
                     sslMode,
+                    helloOk,
                     std::move(options));
 }
 
diff --git a/src/mongo/client/mongo_uri.h b/src/mongo/client/mongo_uri.h
index 4a7ca89126a..0118073f495 100644
--- a/src/mongo/client/mongo_uri.h
+++ b/src/mongo/client/mongo_uri.h
@@ -246,6 +246,15 @@ public:
         return _sslMode;
     }
 
+    bool isHelloOk() const {
+        return _helloOk.get_value_or(false);
+    }
+
+    void setHelloOk(bool helloOk) {
+        invariant(!_helloOk.has_value());
+        _helloOk.emplace(helloOk);
+    }
+
     // If you are trying to clone a URI (including its options/auth information) for a single
     // server (say a member of a replica-set), you can pass in its HostAndPort information to
     // get a new URI with the same info, except type() will be kStandalone and getServers() will
@@ -283,6 +292,7 @@ private:
              const std::string& database,
              boost::optional<bool> retryWrites,
              transport::ConnectSSLMode sslMode,
+             boost::optional<bool> helloOk,
              OptionsMap options)
         : _connectString(std::move(connectString)),
           _user(user),
@@ -290,6 +300,7 @@ private:
           _database(database),
           _retryWrites(std::move(retryWrites)),
           _sslMode(sslMode),
+          _helloOk(helloOk),
           _options(std::move(options)) {}
 
     static MongoURI parseImpl(StringData url);
@@ -300,6 +311,7 @@ private:
     std::string _database;
     boost::optional<bool> _retryWrites;
     transport::ConnectSSLMode _sslMode = transport::kGlobalSSLMode;
+    boost::optional<bool> _helloOk;
     OptionsMap _options;
 };
 
diff --git a/src/mongo/db/repl/replication_info.cpp b/src/mongo/db/repl/replication_info.cpp
index dc94c010f41..f35799db885 100644
--- a/src/mongo/db/repl/replication_info.cpp
+++ b/src/mongo/db/repl/replication_info.cpp
@@ -73,6 +73,8 @@ MONGO_FAIL_POINT_DEFINE(waitInHello);
 // standalones. This failpoint will hang right before doing this sleep when set.
 MONGO_FAIL_POINT_DEFINE(hangWaitingForHelloResponseOnStandalone);
 
+MONGO_FAIL_POINT_DEFINE(appendHelloOkToHelloResponse);
+
 using std::list;
 using std::string;
 using std::stringstream;
@@ -418,6 +420,10 @@ public:
             uassertStatusOK(status);
         }
 
+        if (MONGO_unlikely(appendHelloOkToHelloResponse.shouldFail())) {
+            result.append("clientSupportsHello", client->supportsHello());
+        }
+
         auto currentTopologyVersion = appendReplicationInfo(
             opCtx, &result, 0, useLegacyResponseFields(), clientTopologyVersion, maxAwaitTimeMS);
 
diff --git a/src/mongo/rpc/op_msg_integration_test.cpp b/src/mongo/rpc/op_msg_integration_test.cpp
index 3c3c46ed63b..c3f285e46c5 100644
--- a/src/mongo/rpc/op_msg_integration_test.cpp
+++ b/src/mongo/rpc/op_msg_integration_test.cpp
@@ -29,6 +29,8 @@
 
 #define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kTest
 
+#include <fmt/format.h>
+
 #include "mongo/platform/basic.h"
 
 #include "mongo/client/dbclient_connection.h"
@@ -1283,4 +1285,79 @@ TEST(OpMsg, ServerHandlesReallyLargeMessagesGracefully) {
     ASSERT_NOT_OK(replyStatus);
     ASSERT_EQ(replyStatus, ErrorCodes::BSONObjectTooLarge);
 }
+
+class HelloOkTest final {
+public:
+    auto connect(boost::optional<bool> helloOk = boost::none) const {
+        auto connStr = unittest::getFixtureConnectionString();
+
+        auto swURI = MongoURI::parse(connStr.toString());
+        ASSERT_OK(swURI.getStatus());
+
+        auto uri = swURI.getValue();
+        if (helloOk.has_value()) {
+            uri.setHelloOk(helloOk.get());
+        }
+
+        std::string errMsg;
+        auto conn = connStr.connect(_appName, errMsg, 0, &uri);
+        uassert(ErrorCodes::SocketException, errMsg, conn);
+
+        _configureFailPoint(conn.get());
+        return conn;
+    }
+
+    auto checkIfClientSupportsHello(DBClientBase* conn) const {
+        auto checkHelloSupport = [conn](const std::string& helloCommand) {
+            auto response =
+                conn->runCommand(OpMsgRequest::fromDBAndBody("admin", BSON(helloCommand << 1)))
+                    ->getCommandReply()
+                    .getOwned();
+            auto helloOk = response.getField("clientSupportsHello");
+            ASSERT(!helloOk.eoo());
+            return helloOk.Bool();
+        };
+
+        auto helloOk = checkHelloSupport("hello");
+        ASSERT_EQ(helloOk, checkHelloSupport("isMaster"));
+        ASSERT_EQ(helloOk, checkHelloSupport("ismaster"));
+        return helloOk;
+    }
+
+private:
+    void _configureFailPoint(DBClientBase* conn) const {
+        const auto threadName = getThreadNameByAppName(conn, _appName);
+        const auto failPointObj = BSON("configureFailPoint"
+                                       << "appendHelloOkToHelloResponse"
+                                       << "mode"
+                                       << "alwaysOn"
+                                       << "data" << BSON("threadName" << threadName));
+        auto response = conn->runCommand(OpMsgRequest::fromDBAndBody("admin", failPointObj));
+        ASSERT_OK(getStatusFromCommandResult(response->getCommandReply()));
+    }
+
+    static constexpr auto _appName = "integration_test";
+};
+
+TEST(OpMsg, HelloOkIsDisabledByDefault) {
+    HelloOkTest instance;
+    auto conn = instance.connect();
+    auto isHelloOk = instance.checkIfClientSupportsHello(conn.get());
+    ASSERT(!isHelloOk);
+}
+
+TEST(OpMsg, HelloOkCanBeEnabled) {
+    HelloOkTest instance;
+    auto conn = instance.connect(true);
+    auto isHelloOk = instance.checkIfClientSupportsHello(conn.get());
+    ASSERT(isHelloOk);
+}
+
+TEST(OpMsg, HelloOkCanBeDisabled) {
+    HelloOkTest instance;
+    auto conn = instance.connect(false);
+    auto isHelloOk = instance.checkIfClientSupportsHello(conn.get());
+    ASSERT(!isHelloOk);
+}
+
 }  // namespace mongo
diff --git a/src/mongo/s/commands/cluster_hello_cmd.cpp b/src/mongo/s/commands/cluster_hello_cmd.cpp
index 279103faafb..7a50b0ff5cb 100644
--- a/src/mongo/s/commands/cluster_hello_cmd.cpp
+++ b/src/mongo/s/commands/cluster_hello_cmd.cpp
@@ -53,6 +53,8 @@ namespace mongo {
 // Hangs in the beginning of each hello command when set.
 MONGO_FAIL_POINT_DEFINE(waitInHello);
 
+MONGO_FAIL_POINT_DEFINE(appendHelloOkToHelloResponse);
+
 namespace {
 
 constexpr auto kHelloString = "hello"_sd;
@@ -150,6 +152,9 @@ public:
         bool helloOk;
         Status status = bsonExtractBooleanField(cmdObj, "helloOk", &helloOk);
         if (status.isOK()) {
+            // If the hello request contains a "helloOk" field, set _supportsHello on the Client
+            // to the value.
+            client->setSupportsHello(helloOk);
             // Attach helloOk: true to the response so that the client knows the server supports
             // the hello command.
             result.append("helloOk", true);
@@ -157,6 +162,10 @@ public:
             uassertStatusOK(status);
         }
 
+        if (MONGO_unlikely(appendHelloOkToHelloResponse.shouldFail())) {
+            result.append("clientSupportsHello", client->supportsHello());
+        }
+
         result.appendNumber("maxBsonObjectSize", BSONObjMaxUserSize);
         result.appendNumber("maxMessageSizeBytes", MaxMessageSizeBytes);
         result.appendNumber("maxWriteBatchSize", write_ops::kMaxWriteBatchSize);