SERVER-46189 Fix shell kill current ops with TLS replica set and down nodes

(cherry picked from commit 2fbd718e0ab222d4b6614c452b0241e325421395)

7 files changed, 23 insertions, 30 deletions

diff --git a/src/mongo/client/mongo_uri.cpp b/src/mongo/client/mongo_uri.cpp
index 433c33c8e61..86476612aaf 100644
--- a/src/mongo/client/mongo_uri.cpp
+++ b/src/mongo/client/mongo_uri.cpp
@@ -319,13 +319,11 @@ std::string MongoURI::redact(StringData url) {
     return out.str();
 }
 
-MongoURI MongoURI::parseImpl(const std::string& url) {
-    const StringData urlSD(url);
-
+MongoURI MongoURI::parseImpl(StringData url) {
     // 1. Validate and remove the scheme prefix `mongodb://` or `mongodb+srv://`
-    const bool isSeedlist = urlSD.startsWith(kURISRVPrefix);
-    if (!(urlSD.startsWith(kURIPrefix) || isSeedlist)) {
-        return MongoURI(uassertStatusOK(ConnectionString::parse(url)));
+    const bool isSeedlist = url.startsWith(kURISRVPrefix);
+    if (!(url.startsWith(kURIPrefix) || isSeedlist)) {
+        return MongoURI(uassertStatusOK(ConnectionString::parse(url.toString())));
     }
 
     // 2. Split up the URI into its components for further parsing and validation
@@ -512,7 +510,7 @@ MongoURI MongoURI::parseImpl(const std::string& url) {
                     std::move(options));
 }
 
-StatusWith<MongoURI> MongoURI::parse(const std::string& url) try {
+StatusWith<MongoURI> MongoURI::parse(StringData url) try {
     return parseImpl(url);
 } catch (const std::exception&) {
     return exceptionToStatus();
diff --git a/src/mongo/client/mongo_uri.h b/src/mongo/client/mongo_uri.h
index c69b5e555c7..2e31022d2c5 100644
--- a/src/mongo/client/mongo_uri.h
+++ b/src/mongo/client/mongo_uri.h
@@ -133,7 +133,7 @@ public:
     // whichever map type is used provides that guarantee.
     using OptionsMap = std::map<CaseInsensitiveString, std::string>;
 
-    static StatusWith<MongoURI> parse(const std::string& url);
+    static StatusWith<MongoURI> parse(StringData url);
 
     /*
      * Returns true if str starts with one of the uri schemes (e.g. mongodb:// or mongodb+srv://)
@@ -280,7 +280,7 @@ private:
     boost::optional<BSONObj> _makeAuthObjFromOptions(
         int maxWireVersion, const std::vector<std::string>& saslMechsForAuth) const;
 
-    static MongoURI parseImpl(const std::string& url);
+    static MongoURI parseImpl(StringData url);
 
     ConnectionString _connectString;
     std::string _user;
diff --git a/src/mongo/scripting/engine.cpp b/src/mongo/scripting/engine.cpp
index 2e5a0e2abd0..12f7fcdaa8b 100644
--- a/src/mongo/scripting/engine.cpp
+++ b/src/mongo/scripting/engine.cpp
@@ -562,7 +562,7 @@ unique_ptr<Scope> ScriptEngine::getPooledScope(OperationContext* opCtx,
     return p;
 }
 
-void (*ScriptEngine::_connectCallback)(DBClientBase&) = 0;
+void (*ScriptEngine::_connectCallback)(DBClientBase&, StringData) = 0;
 
 ScriptEngine* getGlobalScriptEngine() {
     if (hasGlobalServiceContext())
diff --git a/src/mongo/scripting/engine.h b/src/mongo/scripting/engine.h
index 18eb7554e6e..8373f407cbf 100644
--- a/src/mongo/scripting/engine.h
+++ b/src/mongo/scripting/engine.h
@@ -249,12 +249,12 @@ public:
     void setScopeInitCallback(void (*func)(Scope&)) {
         _scopeInitCallback = func;
     }
-    static void setConnectCallback(void (*func)(DBClientBase&)) {
+    static void setConnectCallback(void (*func)(DBClientBase&, StringData)) {
         _connectCallback = func;
     }
-    static void runConnectCallback(DBClientBase& c) {
+    static void runConnectCallback(DBClientBase& c, StringData uri) {
         if (_connectCallback)
-            _connectCallback(c);
+            _connectCallback(c, uri);
     }
 
     // engine implementation may either respond to interrupt events or
@@ -270,7 +270,7 @@ protected:
     void (*_scopeInitCallback)(Scope&);
 
 private:
-    static void (*_connectCallback)(DBClientBase&);
+    static void (*_connectCallback)(DBClientBase&, StringData);
 };
 
 void installGlobalUtils(Scope& scope);
diff --git a/src/mongo/scripting/mozjs/mongo.cpp b/src/mongo/scripting/mozjs/mongo.cpp
index 943d0b04716..5d856af72ae 100644
--- a/src/mongo/scripting/mozjs/mongo.cpp
+++ b/src/mongo/scripting/mozjs/mongo.cpp
@@ -832,7 +832,7 @@ void MongoExternalInfo::construct(JSContext* cx, JS::CallArgs args) {
         uasserted(ErrorCodes::InternalError, errmsg);
     }
 
-    ScriptEngine::runConnectCallback(*conn);
+    ScriptEngine::runConnectCallback(*conn, host);
 
     JS::RootedObject thisv(cx);
     scope->getProto<MongoExternalInfo>().newObject(&thisv);
diff --git a/src/mongo/shell/shell_utils.cpp b/src/mongo/shell/shell_utils.cpp
index df0e8aba43c..6a4747d6c45 100644
--- a/src/mongo/shell/shell_utils.cpp
+++ b/src/mongo/shell/shell_utils.cpp
@@ -439,12 +439,11 @@ bool Prompter::confirm() {
 
 ConnectionRegistry::ConnectionRegistry() = default;
 
-void ConnectionRegistry::registerConnection(DBClientBase& client) {
+void ConnectionRegistry::registerConnection(DBClientBase& client, StringData uri) {
     BSONObj info;
     if (client.runCommand("admin", BSON("whatsmyuri" << 1), info)) {
-        std::string connstr = client.getServerAddress();
         stdx::lock_guard<Latch> lk(_mutex);
-        _connectionUris[connstr].insert(info["you"].str());
+        _connectionUris[uri.toString()].insert(info["you"].str());
     }
 }
 
@@ -452,15 +451,10 @@ void ConnectionRegistry::killOperationsOnAllConnections(bool withPrompt) const {
     Prompter prompter("do you want to kill the current op(s) on the server?");
     stdx::lock_guard<Latch> lk(_mutex);
     for (auto& connection : _connectionUris) {
-        auto status = ConnectionString::parse(connection.first);
-        if (!status.isOK()) {
-            continue;
-        }
-
-        const ConnectionString cs(status.getValue());
-
         std::string errmsg;
-        std::unique_ptr<DBClientBase> conn(cs.connect("MongoDB Shell", errmsg));
+
+        auto uri = uassertStatusOK(MongoURI::parse(connection.first));
+        std::unique_ptr<DBClientBase> conn(uri.connect("MongoDB Shell", errmsg));
         if (!conn) {
             continue;
         }
@@ -518,7 +512,7 @@ void ConnectionRegistry::killOperationsOnAllConnections(bool withPrompt) const {
 
 ConnectionRegistry connectionRegistry;
 
-void onConnect(DBClientBase& c) {
+void onConnect(DBClientBase& c, StringData uri) {
     if (shellGlobalParams.nokillop) {
         return;
     }
@@ -528,7 +522,7 @@ void onConnect(DBClientBase& c) {
         c.setClientRPCProtocols(*shellGlobalParams.rpcProtocols);
     }
 
-    connectionRegistry.registerConnection(c);
+    connectionRegistry.registerConnection(c, uri);
 }
 
 bool fileExists(const std::string& file) {
diff --git a/src/mongo/shell/shell_utils.h b/src/mongo/shell/shell_utils.h
index 5fb2b844eb6..9151faec2ce 100644
--- a/src/mongo/shell/shell_utils.h
+++ b/src/mongo/shell/shell_utils.h
@@ -34,6 +34,7 @@
 #include <set>
 #include <string>
 
+#include "mongo/client/mongo_uri.h"
 #include "mongo/db/jsobj.h"
 #include "mongo/platform/mutex.h"
 #include "mongo/util/concurrency/mutex.h"
@@ -52,7 +53,7 @@ void RecordMyLocation(const char* _argv0);
 void installShellUtils(Scope& scope);
 
 void initScope(Scope& scope);
-void onConnect(DBClientBase& c);
+void onConnect(DBClientBase& c, StringData uri);
 
 boost::filesystem::path getHistoryFilePath();
 void setEnterpriseShellCallback(EnterpriseShellCallback* callback);
@@ -77,7 +78,7 @@ private:
 class ConnectionRegistry {
 public:
     ConnectionRegistry();
-    void registerConnection(DBClientBase& client);
+    void registerConnection(DBClientBase& client, StringData uri);
     void killOperationsOnAllConnections(bool withPrompt) const;
 
 private: