SERVER-28072 Construct mongodb:// URI from RS string and dbname

4 files changed, 182 insertions, 69 deletions

diff --git a/jstests/noPassthroughWithMongod/replset_host_connection_validation.js b/jstests/noPassthroughWithMongod/replset_host_connection_validation.js
index 7e48de363e2..e65387dd705 100644
--- a/jstests/noPassthroughWithMongod/replset_host_connection_validation.js
+++ b/jstests/noPassthroughWithMongod/replset_host_connection_validation.js
@@ -1,5 +1,6 @@
 // Test --host with a replica set.
 (function() {
+    'use strict';
 
     const replSetName = 'hostTestReplSetName';
 
@@ -26,35 +27,56 @@
 
         // Pass the inner test's exit code back as the outer test's exit code
         if (exitCode != 0) {
-            doassert("inner test failed with exit code " + exitcode);
+            doassert("inner test failed with exit code " + exitCode);
         }
         return;
     }
 
-    const testHost = function(host) {
-        const exitCode = runMongoProgram('mongo', '--eval', ';', '--host', host);
-        if (exitCode !== 0) {
-            doassert("failed to connect with `--host " + host +
-                     "`, but expected success. Exit code: " + exitCode);
+    function testHost(host, uri, ok) {
+        const exitCode = runMongoProgram('mongo', '--eval', ';', '--host', host, uri);
+        if (ok) {
+            assert.eq(exitCode, 0, "failed to connect with `--host " + host + "`");
+        } else {
+            assert.neq(exitCode, 0, "unexpectedly succeeded to connect with `--host " + host + "`");
         }
-    };
-
-    const connStrings = [
-        `localhost:${port}`,
-        `${replSetName}/localhost:${port}`,
-        `mongodb://localhost:${port}/admin?replicaSet=${replSetName}`,
-        `mongodb://localhost:${port}`,
-    ];
-
-    function runConnectionStringTestFor(i, connectionString) {
-        print("Testing connection string " + i + "...");
-        print("    * testing " + connectionString);
-        testHost(connectionString);
     }
 
-    for (let i = 0; i < connStrings.length; ++i) {
-        runConnectionStringTestFor(i, connStrings[i]);
+    function runConnectionStringTestFor(connectionString, uri, ok) {
+        print("* Testing: --host " + connectionString + " " + uri);
+        if (!ok) {
+            print("  This should fail");
+        }
+        testHost(connectionString, uri, ok);
+    }
+
+    function expSuccess(str) {
+        runConnectionStringTestFor(str, '', true);
+        if (!str.startsWith('mongodb://')) {
+            runConnectionStringTestFor(str, 'dbname', true);
+        }
     }
 
+    function expFailure(str) {
+        runConnectionStringTestFor(str, '', false);
+    }
+
+    expSuccess(`localhost:${port}`);
+    expSuccess(`${replSetName}/localhost:${port}`);
+    expSuccess(`${replSetName}/localhost:${port},[::1]:${port}`);
+    expSuccess(`${replSetName}/localhost:${port},`);
+    expSuccess(`${replSetName}/localhost:${port},,`);
+    expSuccess(`mongodb://localhost:${port}/admin?replicaSet=${replSetName}`);
+    expSuccess(`mongodb://localhost:${port}`);
+
+    expFailure(',');
+    expFailure(',,');
+    expFailure(`${replSetName}/`);
+    expFailure(`${replSetName}/,`);
+    expFailure(`${replSetName}/,,`);
+    expFailure(`${replSetName}//not/a/socket`);
+    expFailure(`mongodb://localhost:${port}/admin?replicaSet=`);
+    expFailure('mongodb://localhost:');
+    expFailure(`mongodb://:${port}`);
+
     jsTest.log("SUCCESSFUL test completion");
 })();
diff --git a/src/mongo/client/mongo_uri.cpp b/src/mongo/client/mongo_uri.cpp
index 1db9c10d2c0..e2356881290 100644
--- a/src/mongo/client/mongo_uri.cpp
+++ b/src/mongo/client/mongo_uri.cpp
@@ -59,9 +59,10 @@ const mongo::StringData kURIPrefix{"mongodb://"};
  * Encode data elements in a way which will allow them to be embedded
  * into a mongodb:// URI safely.
  */
-void mongo::uriEncode(std::ostream& ss, StringData toEncode) {
+void mongo::uriEncode(std::ostream& ss, StringData toEncode, StringData passthrough) {
     for (const auto& c : toEncode) {
-        if ((c == '-') || (c == '_') || (c == '.') || (c == '~') || isalnum(c)) {
+        if ((c == '-') || (c == '_') || (c == '.') || (c == '~') || isalnum(c) ||
+            (passthrough.find(c) != std::string::npos)) {
             ss << c;
         } else {
             // Encoding anything not included in section 2.3 "Unreserved characters"
diff --git a/src/mongo/client/mongo_uri.h b/src/mongo/client/mongo_uri.h
index d42d5d33d5d..f7a7c4aeb81 100644
--- a/src/mongo/client/mongo_uri.h
+++ b/src/mongo/client/mongo_uri.h
@@ -47,11 +47,13 @@ namespace mongo {
 /**
  * Encode a string for embedding in a URI.
  * Replaces reserved bytes with %xx sequences.
+ *
+ * Optionally allows passthrough characters to remain unescaped.
  */
-void uriEncode(std::ostream& ss, StringData str);
-inline std::string uriEncode(StringData str) {
+void uriEncode(std::ostream& ss, StringData str, StringData passthrough = ""_sd);
+inline std::string uriEncode(StringData str, StringData passthrough = ""_sd) {
     std::ostringstream ss;
-    uriEncode(ss, str);
+    uriEncode(ss, str, passthrough);
     return ss.str();
 }
 
diff --git a/src/mongo/shell/dbshell.cpp b/src/mongo/shell/dbshell.cpp
index 051d5726993..9273e7b6f25 100644
--- a/src/mongo/shell/dbshell.cpp
+++ b/src/mongo/shell/dbshell.cpp
@@ -217,63 +217,151 @@ void setupSignals() {
     signal(SIGINT, quitNicely);
 }
 
-string getURIFromArgs(const std::string& url, const std::string& host, const std::string& port) {
-    if (host.size() == 0 && port.size() == 0) {
-        return url.size() == 0 ? kDefaultMongoURL.toString() : url;
+string getURIFromArgs(const std::string& arg, const std::string& host, const std::string& port) {
+    if (host.empty() && arg.empty() && port.empty()) {
+        // Nothing provided, just play the default.
+        return kDefaultMongoURL.toString();
     }
 
-    // The name URL is misleading; really it's just a positional argument that wasn't a file. The
-    // check for "/" means "this 'URL' is probably a real URL and not the db name (e.g.)".
-    if (url.find("/") != string::npos) {
-        cerr << "if a full URI is provided, you cannot also specify host or port" << endl;
+    if (str::startsWith(arg, "mongodb://") && host.empty() && port.empty()) {
+        // mongo mongodb://blah
+        return arg;
+    }
+    if (str::startsWith(host, "mongodb://") && arg.empty() && port.empty()) {
+        // mongo --host mongodb://blah
+        return host;
+    }
+
+    // We expect a positional arg to be a plain dbname or plain hostname at this point
+    // since we have separate host/port args.
+    if ((arg.find('/') != string::npos) && (host.size() || port.size())) {
+        cerr << "If a full URI is provided, you cannot also specify --host or --port" << endl;
         quickExit(-1);
     }
 
-    bool hostEndsInSock = str::endsWith(host, ".sock");
-    const auto hostHasPort = (host.find(":") != std::string::npos);
+    const auto parseDbHost = [port](const std::string& db, const std::string& host) -> std::string {
+        // Parse --host as a connection string.
+        // e.g. rs0/host0:27000,host1:27001
+        const auto slashPos = host.find('/');
+        const auto hasReplSet = (slashPos > 0) && (slashPos != std::string::npos);
+
+        std::ostringstream ss;
+        ss << "mongodb://";
+
+        // Handle each sub-element of the connection string individually.
+        // Comma separated list of host elements.
+        // Each host element may be:
+        // * /unix/domain.sock
+        // * hostname
+        // * hostname:port
+        // If --port is specified and port is included in connection string,
+        // then they must match exactly.
+        auto start = hasReplSet ? slashPos + 1 : 0;
+        while (start < host.size()) {
+            // Encode each host component.
+            auto end = host.find(',', start);
+            if (end == std::string::npos) {
+                end = host.size();
+            }
+            if ((end - start) == 0) {
+                // Ignore empty components.
+                start = end + 1;
+                continue;
+            }
+
+            const auto hostElem = host.substr(start, end - start);
+            if ((hostElem.find('/') != std::string::npos) && str::endsWith(hostElem, ".sock")) {
+                // Unix domain socket, ignore --port.
+                ss << uriEncode(hostElem);
+
+            } else {
+                auto colon = hostElem.find(':');
+                if ((colon != std::string::npos) &&
+                    (hostElem.find(':', colon + 1) != std::string::npos)) {
+                    // Looks like an IPv6 numeric address.
+                    const auto close = hostElem.find(']');
+                    if ((hostElem[0] == '[') && (close != std::string::npos)) {
+                        // Encapsulated already.
+                        ss << '[' << uriEncode(hostElem.substr(1, close - 1), ":") << ']';
+                        colon = hostElem.find(':', close + 1);
+                    } else {
+                        // Not encapsulated yet.
+                        ss << '[' << uriEncode(hostElem, ":") << ']';
+                        colon = std::string::npos;
+                    }
+                } else if (colon != std::string::npos) {
+                    // Not IPv6 numeric, but does have a port.
+                    ss << uriEncode(hostElem.substr(0, colon));
+                } else {
+                    // Raw hostname/IPv4 without port.
+                    ss << uriEncode(hostElem);
+                }
+
+                if (colon != std::string::npos) {
+                    // Have a port in our host element, verify it.
+                    const auto myport = hostElem.substr(colon + 1);
+                    if (port.size() && (port != myport)) {
+                        cerr << "connection string bears different port than provided by --port"
+                             << endl;
+                        quickExit(-1);
+                    }
+                    ss << ':' << uriEncode(myport);
+                } else if (port.size()) {
+                    ss << ':' << uriEncode(port);
+                } else {
+                    ss << ":27017";
+                }
+            }
+            start = end + 1;
+            if (start < host.size()) {
+                ss << ',';
+            }
+        }
+
+        ss << '/' << uriEncode(db);
 
-    // If host looks like a full URI (i.e. has a slash and isn't a unix socket) and the other fields
-    // are empty, then just return host.
-    std::string::size_type slashPos;
-    if (url.size() == 0 && port.size() == 0 &&
-        (!hostEndsInSock && ((slashPos = host.find("/")) != string::npos))) {
-        if (str::startsWith(host, "mongodb://")) {
-            return host;
+        if (hasReplSet) {
+            // Remap included replica set name to URI option
+            ss << "?replicaSet=" << uriEncode(host.substr(0, slashPos));
         }
-        // If there's a slash in the host field, then it's the replica set name, not a database name
-        stringstream ss;
-        ss << "mongodb://" << uriEncode(host.substr(slashPos + 1))
-           << "/?replicaSet=" << uriEncode(host.substr(0, slashPos));
+
         return ss.str();
-    }
+    };
+
+    if (host.size()) {
+        // --host provided, treat it as the connect string and get db from positional arg.
+        return parseDbHost(arg, host);
+    } else if (arg.size()) {
+        // --host missing, but we have a potential db/host positional arg.
+        const auto slashPos = arg.find('/');
+        if (slashPos != std::string::npos) {
+            // db/host pair.
+            return parseDbHost(arg.substr(0, slashPos), arg.substr(slashPos + 1));
+        }
 
-    stringstream ss;
-    if (host.size() == 0) {
-        ss << "mongodb://127.0.0.1";
-    } else {
-        if (str::startsWith(host, "mongodb://")) {
-            ss << host;
-        } else {
-            ss << "mongodb://" << uriEncode(host);
+        // Compatability formats.
+        // * Any arg with a dot is assumed to be a hostname or IPv4 numeric address.
+        // * Any arg with a colon followed by a digit assumed to be host or IP followed by port.
+        // * Anything else is assumed to be a db.
+
+        if (arg.find('.') != std::string::npos) {
+            // Assume IPv4 or hostnameish.
+            return parseDbHost("test", arg);
         }
-    }
 
-    if (!hostEndsInSock) {
-        if (port.size() > 0) {
-            if (hostHasPort) {
-                std::cerr << "Cannot specify a port in --host and also with --port" << std::endl;
-                quickExit(-1);
-            }
-            ss << ":" << port;
-        } else if (!hostHasPort || str::endsWith(host, "]")) {
-            // Default the port to 27017 if the host did not provide one (i.e. the host has no
-            // colons or ends in ']' like an IPv6 address).
-            ss << ":27017";
+        const auto colonPos = arg.find(':');
+        if ((colonPos != std::string::npos) && ((colonPos + 1) < arg.size()) &&
+            isdigit(arg[colonPos + 1])) {
+            // Assume IPv4 or hostname with port.
+            return parseDbHost("test", arg);
         }
+
+        // db, assume localhost.
+        return parseDbHost(arg, "127.0.0.1");
     }
 
-    ss << "/" << uriEncode(url);
-    return ss.str();
+    // --host empty, position arg empty, fallback on localhost without a dbname.
+    return parseDbHost("", "127.0.0.1");
 }
 
 static string OpSymbols = "~!%^&*-+=|:,<>/?.";