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

7 files changed, 23 insertions, 30 deletions

diff --git a/src/mongo/client/mongo_uri.cpp b/src/mongo/client/mongo_uri.cpp
index f0f8cab0978..d9153cb3c02 100644
--- a/src/mongo/client/mongo_uri.cpp
+++ b/src/mongo/client/mongo_uri.cpp
@@ -320,13 +320,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
@@ -513,7 +511,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 5b9f73039f8..3d155e1123c 100644
--- a/src/mongo/client/mongo_uri.h
+++ b/src/mongo/client/mongo_uri.h
@@ -134,7 +134,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://)
@@ -281,7 +281,7 @@ private:
           _sslMode(sslMode),
           _options(std::move(options)) {}
 
-    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 c430e057dbb..d941e9834af 100644
--- a/src/mongo/scripting/engine.cpp
+++ b/src/mongo/scripting/engine.cpp
@@ -577,7 +577,7 @@ unique_ptr<Scope> ScriptEngine::getPooledScope(OperationContext* opCtx,
     return p;
 }
 
-void (*ScriptEngine::_connectCallback)(DBClientBase&) = nullptr;
+void (*ScriptEngine::_connectCallback)(DBClientBase&, StringData) = nullptr;
 
 ScriptEngine* getGlobalScriptEngine() {
     if (hasGlobalServiceContext())
diff --git a/src/mongo/scripting/engine.h b/src/mongo/scripting/engine.h
index 434b38ff921..a87d4edc2ea 100644
--- a/src/mongo/scripting/engine.h
+++ b/src/mongo/scripting/engine.h
@@ -259,12 +259,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
@@ -281,7 +281,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 cbbcdb82658..9d13c551017 100644
--- a/src/mongo/scripting/mozjs/mongo.cpp
+++ b/src/mongo/scripting/mozjs/mongo.cpp
@@ -834,7 +834,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 a8216de7473..84e4c1ef4fd 100644
--- a/src/mongo/shell/shell_utils.cpp
+++ b/src/mongo/shell/shell_utils.cpp
@@ -541,12 +541,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());
     }
 }
 
@@ -554,15 +553,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;
         }
@@ -620,7 +614,7 @@ void ConnectionRegistry::killOperationsOnAllConnections(bool withPrompt) const {
 
 ConnectionRegistry connectionRegistry;
 
-void onConnect(DBClientBase& c) {
+void onConnect(DBClientBase& c, StringData uri) {
     if (shellGlobalParams.nokillop) {
         return;
     }
@@ -630,7 +624,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 1d3e4099998..ffde5fa5de7 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"
@@ -54,7 +55,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);
@@ -79,7 +80,7 @@ private:
 class ConnectionRegistry {
 public:
     ConnectionRegistry();
-    void registerConnection(DBClientBase& client);
+    void registerConnection(DBClientBase& client, StringData uri);
     void killOperationsOnAllConnections(bool withPrompt) const;
 
 private: