Fix MongoConnectionString errors (#729)

isabelatkinson · web-flow · commit dbc6ef608218 · 2022-01-26T13:06:32.000-05:00
diff --git a/Sources/MongoSwift/ConnectionPool.swift b/Sources/MongoSwift/ConnectionPool.swift
@@ -296,6 +296,22 @@ internal class ConnectionPool {
         }
     }
 
+    /// Retrieves the `defaultAuthDB` configured on the connection string used to create this pool. If SDAM has been
+    /// started in libmongoc, this will return the up-to-date value after SRV lookup.
+    internal func getConnectionStringAuthDB() throws -> String? {
+        try self.withConnection { connection in
+            try connection.withMongocConnection { connPtr in
+                guard let uri = mongoc_client_get_uri(connPtr) else {
+                    throw MongoError.InternalError(message: "Couldn't retrieve client's connection string")
+                }
+                guard let authSource = mongoc_uri_get_database(uri) else {
+                    return nil
+                }
+                return String(cString: authSource)
+            }
+        }
+    }
+
     /// Sets the provided APM callbacks on this pool, using the provided client as the "context" value. **This method
     /// may only be called before any connections are checked out of the pool.** Ideally this code would just live in
     /// `ConnectionPool.init`. However, the client we accept here has to be fully initialized before we can pass it
diff --git a/Sources/MongoSwift/MongoClient.swift b/Sources/MongoSwift/MongoClient.swift
@@ -319,6 +319,7 @@ public class MongoClient {
         if let options = options {
             try connString.applyOptions(options)
         }
+        try connString.validate()
 
         self.operationExecutor = OperationExecutor(
             eventLoopGroup: eventLoopGroup,
diff --git a/Sources/MongoSwift/MongoConnectionString.swift b/Sources/MongoSwift/MongoConnectionString.swift
@@ -120,10 +120,6 @@ public struct MongoConnectionString: Codable, LosslessStringConvertible {
 
         /// Initializes a ServerAddress, using the default localhost:27017 if a host/port is not provided.
         internal init(_ hostAndPort: String = "localhost:27017") throws {
-            guard !hostAndPort.contains("?") else {
-                throw MongoError.InvalidArgumentError(message: "\(hostAndPort) contains invalid characters")
-            }
-
             // Check if host is an IPv6 literal.
             if hostAndPort.first == "[" {
                 let ipLiteralRegex = try NSRegularExpression(pattern: #"^\[(.*)\](?::([0-9]+))?$"#)
@@ -228,7 +224,7 @@ public struct MongoConnectionString: Codable, LosslessStringConvertible {
         fileprivate var wTimeoutMS: Int?
         fileprivate var zlibCompressionLevel: Int?
 
-        fileprivate init(_ uriOptions: String) throws {
+        fileprivate init(_ uriOptions: Substring) throws {
             let options = uriOptions.components(separatedBy: "&")
             // tracks options that have already been set to error on duplicates
             var setOptions: Set<String> = []
@@ -384,30 +380,51 @@ public struct MongoConnectionString: Codable, LosslessStringConvertible {
     /// - Throws:
     ///   - `MongoError.InvalidArgumentError` if the input is invalid.
     public init(string input: String) throws {
+        // Parse the connection string's scheme.
         let schemeAndRest = input.components(separatedBy: "://")
         guard schemeAndRest.count == 2, let scheme = Scheme(schemeAndRest[0]) else {
             throw MongoError.InvalidArgumentError(
                 message: "Invalid connection string scheme, expecting \'mongodb\' or \'mongodb+srv\'"
             )
         }
         guard !schemeAndRest[1].isEmpty else {
-            throw MongoError.InvalidArgumentError(message: "Invalid connection string")
+            throw MongoError.InvalidArgumentError(message: "The connection string must contain host information")
         }
-        let identifiersAndOptions = schemeAndRest[1].components(separatedBy: "/")
-        guard identifiersAndOptions.count <= 2 else {
+
+        // Split the rest of the connection string into its components.
+        let infoAndOptions = schemeAndRest[1].split(separator: "?", maxSplits: 1, omittingEmptySubsequences: false)
+        if infoAndOptions[0].isEmpty {
+            throw MongoError.InvalidArgumentError(message: "The connection string must contain host information")
+        }
+        let userHostsAndAuthDB = infoAndOptions[0].split(separator: "/", omittingEmptySubsequences: false)
+        if userHostsAndAuthDB.count > 2 {
+            throw MongoError.InvalidArgumentError(
+                message: "The user information, host information, and defaultAuthDB in the connection string must not"
+                    + " contain unescaped slashes"
+            )
+        } else if userHostsAndAuthDB.count == 1 && infoAndOptions.count == 2 {
             throw MongoError.InvalidArgumentError(
-                message: "Connection string contains an unescaped slash"
+                message: "The connection string must contain a delimiting slash between the host information and"
+                    + " options"
             )
         }
-        let userAndHost = identifiersAndOptions[0].components(separatedBy: "@")
-        guard userAndHost.count <= 2 else {
-            throw MongoError.InvalidArgumentError(message: "Connection string contains an unescaped @ symbol")
+        let userInfoAndHosts = userHostsAndAuthDB[0].split(separator: "@", omittingEmptySubsequences: false)
+        if userInfoAndHosts.count > 2 {
+            throw MongoError.InvalidArgumentError(
+                message: "The user information and host information in the connection string must not contain"
+                    + " unescaped @ symbols"
+            )
         }
 
-        // do not omit empty subsequences to include an empty password
-        let userInfo = userAndHost.count == 2 ?
-            userAndHost[0].split(separator: ":", maxSplits: 1, omittingEmptySubsequences: false) : nil
-        if let userInfo = userInfo {
+        // Parse user information if present and set the hosts string.
+        let hostsString: Substring
+        if userInfoAndHosts.count == 2 {
+            let userInfo = userInfoAndHosts[0].split(separator: ":", omittingEmptySubsequences: false)
+            if userInfo.count > 2 {
+                throw MongoError.InvalidArgumentError(
+                    message: "Username and password in the connection string must not contain unescaped colons"
+                )
+            }
             var credential = MongoCredential()
             credential.username = try userInfo[0].getValidatedUserInfo(forKey: "username")
             if userInfo.count == 2 {
@@ -416,11 +433,16 @@ public struct MongoConnectionString: Codable, LosslessStringConvertible {
             // If no other authentication options or defaultAuthDB were provided, we should use "admin" as the
             // credential source. This will be overwritten later if a defaultAuthDB or an authSource is provided.
             credential.source = "admin"
+            // Overwrite the sourceFromAuthSource field to false as the source is a default.
+            credential.sourceFromAuthSource = false
             self.credential = credential
+            hostsString = userInfoAndHosts[1]
+        } else {
+            hostsString = userInfoAndHosts[0]
         }
 
-        let hostString = userInfo != nil ? userAndHost[1] : userAndHost[0]
-        let hosts = try hostString.components(separatedBy: ",").map(HostIdentifier.init)
+        // Parse host information.
+        let hosts = try hostsString.components(separatedBy: ",").map(HostIdentifier.init)
         if case .srv = scheme {
             guard hosts.count == 1 else {
                 throw MongoError.InvalidArgumentError(
@@ -432,45 +454,44 @@ public struct MongoConnectionString: Codable, LosslessStringConvertible {
                     message: "A port cannot be specified in a mongodb+srv connection string"
                 )
             }
+            guard hosts[0].host.filter({ $0 == "." }).count >= 2 else {
+                throw MongoError.InvalidArgumentError(
+                    message: "The host specified in a mongodb+srv connection string must contain a host name, a domain"
+                        + " name, and a top-level domain"
+                )
+            }
         }
         self.scheme = scheme
         self.hosts = hosts
 
-        guard identifiersAndOptions.count == 2 else {
-            // no auth DB or options were specified
-            try self.validate()
-            return
-        }
-
-        let authDatabaseAndOptions = identifiersAndOptions[1].components(separatedBy: "?")
-        guard authDatabaseAndOptions.count <= 2 else {
-            throw MongoError.InvalidArgumentError(message: "Connection string contains an unescaped question mark")
-        }
-        if !authDatabaseAndOptions[0].isEmpty {
-            let decoded = try authDatabaseAndOptions[0].getPercentDecoded(forKey: "defaultAuthDB")
+        // Parse the defaultAuthDB if present.
+        if userHostsAndAuthDB.count == 2 && !userHostsAndAuthDB[1].isEmpty {
+            let defaultAuthDB = try userHostsAndAuthDB[1].getPercentDecoded(forKey: "defaultAuthDB")
             for character in Self.forbiddenDBCharacters {
-                if decoded.contains(character) {
+                if defaultAuthDB.contains(character) {
                     throw MongoError.InvalidArgumentError(
                         message: "defaultAuthDB contains invalid character: \(character)"
                     )
                 }
             }
-            self.defaultAuthDB = decoded
+            self.defaultAuthDB = defaultAuthDB
             // If no other authentication options were provided, we should use the defaultAuthDB as the credential
             // source. This will be overwritten later if an authSource is provided.
             if self.credential == nil {
                 self.credential = MongoCredential()
             }
-            self.credential?.source = decoded
+            self.credential?.source = defaultAuthDB
+            // Overwrite the sourceFromAuthSource field to false as the source is a default.
+            self.credential?.sourceFromAuthSource = false
         }
 
-        guard authDatabaseAndOptions.count == 2 else {
-            // no options were specified
+        // Return early if no options were specified.
+        guard infoAndOptions.count == 2 else {
             try self.validate()
             return
         }
 
-        let options = try Options(authDatabaseAndOptions[1])
+        let options = try Options(infoAndOptions[1])
 
         // Validate and set compressors. This validation is only necessary for compressors provided in the URI string
         // and therefore is not included in the general validate method.
@@ -617,7 +638,7 @@ public struct MongoConnectionString: Codable, LosslessStringConvertible {
         try self.validate()
     }
 
-    private mutating func validate() throws {
+    internal mutating func validate() throws {
         func optionError(name: OptionName, violation: String) -> MongoError.InvalidArgumentError {
             MongoError.InvalidArgumentError(
                 message: "Value for \(name) in the connection string must " + violation
@@ -717,6 +738,8 @@ public struct MongoConnectionString: Codable, LosslessStringConvertible {
                     ?? self.defaultAuthDB
                     ?? "admin"
                 self.credential?.source = defaultSource
+                // Overwrite the sourceFromAuthSource field to false as the source is a default.
+                self.credential?.sourceFromAuthSource = false
             }
             if self.credential?.mechanism != nil {
                 // credential cannot be nil within the external conditional
@@ -931,7 +954,9 @@ public struct MongoConnectionString: Codable, LosslessStringConvertible {
             }
             return property
         }.joined(separator: ","))
-        uri.appendOption(name: .authSource, option: self.credential?.source)
+        if self.credential?.sourceFromAuthSource == true {
+            uri.appendOption(name: .authSource, option: self.credential?.source)
+        }
         uri.appendOption(name: .compressors, option: self.compressors?.map {
             switch $0._compressor {
             case let .zlib(level):
diff --git a/Sources/MongoSwift/MongoCredential.swift b/Sources/MongoSwift/MongoCredential.swift
@@ -1,3 +1,4 @@
+import SwiftBSON
 /// Represents an authentication credential.
 public struct MongoCredential: Decodable, Equatable {
     /// A string containing the username. For auth mechanisms that do not utilize a password, this may be the entire
@@ -8,7 +9,17 @@ public struct MongoCredential: Decodable, Equatable {
     public var password: String?
 
     /// A string containing the authentication database.
-    public var source: String?
+    public var source: String? {
+        didSet {
+            self.sourceFromAuthSource = self.source != nil
+        }
+    }
+
+    /// Tracks whether the `source` was set manually (i.e. via an `authSource` or by setting the field's value) or by
+    /// a default (i.e. via the `defaultAuthDB`, the `mechanism`'s default, or the default of "admin"). This
+    /// information is necessary to determine whether the `authSource` field should be set when reconstructing a
+    /// `MongoConnectionString`.
+    internal var sourceFromAuthSource: Bool
 
     /// The authentication mechanism. A nil value for this property indicates that a mechanism wasn't specified and
     /// that mechanism negotiation is required.
@@ -21,6 +32,16 @@ public struct MongoCredential: Decodable, Equatable {
         case username, password, source, mechanism, mechanismProperties = "mechanism_properties"
     }
 
+    public init(from decoder: Decoder) throws {
+        let container = try decoder.container(keyedBy: CodingKeys.self)
+        self.username = try container.decodeIfPresent(String.self, forKey: .username)
+        self.password = try container.decodeIfPresent(String.self, forKey: .password)
+        self.source = try container.decodeIfPresent(String.self, forKey: .source)
+        self.sourceFromAuthSource = self.source != nil
+        self.mechanism = try container.decodeIfPresent(Mechanism.self, forKey: .mechanism)
+        self.mechanismProperties = try container.decodeIfPresent(BSONDocument.self, forKey: .mechanismProperties)
+    }
+
     public init(
         username: String? = nil,
         password: String? = nil,
@@ -31,6 +52,7 @@ public struct MongoCredential: Decodable, Equatable {
         self.username = username
         self.password = password
         self.source = source
+        self.sourceFromAuthSource = self.source != nil
         self.mechanism = mechanism
         self.mechanismProperties = mechanismProperties
     }
diff --git a/Tests/MongoSwiftTests/AuthTests.swift b/Tests/MongoSwiftTests/AuthTests.swift
@@ -89,6 +89,10 @@ final class AuthTests: MongoSwiftTestCase {
                         connStringCredential.mechanismProperties = nil
                     }
 
+                    // this field is only relevant for rebuilding a connection string
+                    credential.sourceFromAuthSource = false
+                    connStringCredential.sourceFromAuthSource = false
+
                     // compare rest of non-document options normally
                     expect(connStringCredential).to(equal(credential), description: testCase.description)
                 }
diff --git a/Tests/MongoSwiftTests/DNSSeedlistTests.swift b/Tests/MongoSwiftTests/DNSSeedlistTests.swift
@@ -90,6 +90,11 @@ final class DNSSeedlistTests: MongoSwiftTestCase {
 
     func runDNSSeedlistTests(_ tests: [(String, DNSSeedlistTestCase)]) throws {
         for (fileName, testCase) in tests {
+            // TODO: SWIFT-1455: unskip these test
+            if fileName == "loadBalanced-no-results.json" || fileName == "loadBalanced-true-multiple-hosts.json" {
+                continue
+            }
+
             // this particular test case requires SSL is disabled. see DRIVERS-1324.
             let requiresTLS = fileName != "txt-record-with-overridden-ssl-option.json"
 
@@ -138,7 +143,9 @@ final class DNSSeedlistTests: MongoSwiftTestCase {
                         case "ssl":
                             expect(connStrOptions["tls"]).to(equal(v))
                         // these values are not returned as part of the options doc
-                        case "authSource", "auth_database":
+                        case "auth_database":
+                            expect(try client.connectionPool.getConnectionStringAuthDB()).to(equal(v.stringValue))
+                        case "authSource":
                             expect(try client.connectionPool.getConnectionStringAuthSource()).to(equal(v.stringValue))
                         case "user":
                             expect(connStr.credential?.username).to(equal(v.stringValue))
diff --git a/Tests/MongoSwiftTests/MongoConnectionStringTests.swift b/Tests/MongoSwiftTests/MongoConnectionStringTests.swift
@@ -621,4 +621,44 @@ final class ConnectionStringTests: MongoSwiftTestCase {
         expect(host.host).to(equal("256.1.2.3"))
         expect(host.type).to(equal(.hostname))
     }
+
+    func testAuthSourceInDescription() throws {
+        // defaultAuthDB
+        let connString1 = try MongoConnectionString(string: "mongodb://localhost:27017/test")
+        expect(connString1.description).toNot(contain("authsource"))
+
+        // authSource
+        let connString2 = try MongoConnectionString(string: "mongodb://localhost:27017/?authSource=test")
+        expect(connString2.description).to(contain("authsource=test"))
+
+        // defaultAuthDB and authSource
+        let connString3 = try MongoConnectionString(string: "mongodb://localhost:27017/admin?authSource=test")
+        expect(connString3.description).to(contain("authsource=test"))
+
+        // set from options
+        let options = MongoClientOptions(credential: MongoCredential(source: "test"))
+        var connString4 = try MongoConnectionString(string: "mongodb://localhost:27017")
+        try connString4.applyOptions(options)
+        expect(connString4.description).to(contain("authsource=test"))
+
+        // set manually
+        let credential1 = MongoCredential(username: "user", source: "test")
+        var connString5 = try MongoConnectionString(string: "mongodb://localhost:27107")
+        connString5.credential = credential1
+        expect(connString5.description).to(contain("authsource=test"))
+
+        // field changed to value
+        let credential2 = MongoCredential(username: "user")
+        var connString6 = try MongoConnectionString(string: "mongodb://localhost:27017")
+        connString6.credential = credential2
+        connString6.credential?.source = "test"
+        expect(connString6.description).to(contain("authsource=test"))
+
+        // field changed to nil
+        let credential3 = MongoCredential(username: "user", source: "test")
+        var connString7 = try MongoConnectionString(string: "mongodb://localhost:27107")
+        connString7.credential = credential3
+        connString7.credential?.source = nil
+        expect(connString7.description).toNot(contain("authsource"))
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -319,6 +319,7 @@ public class MongoClient {`
`319`	`319`	`if let options = options {`
`320`	`320`	`try connString.applyOptions(options)`
`321`	`321`	`}`
	`322`	`+ try connString.validate()`
`322`	`323`
`323`	`324`	`self.operationExecutor = OperationExecutor(`
`324`	`325`	`eventLoopGroup: eventLoopGroup,`
Original file line number	Diff line number	Diff line change
`@@ -89,6 +89,10 @@ final class AuthTests: MongoSwiftTestCase {`
`89`	`89`	`connStringCredential.mechanismProperties = nil`
`90`	`90`	`}`
`91`	`91`
	`92`	`+ // this field is only relevant for rebuilding a connection string`
	`93`	`+ credential.sourceFromAuthSource = false`
	`94`	`+ connStringCredential.sourceFromAuthSource = false`
	`95`	`+`
`92`	`96`	`// compare rest of non-document options normally`
`93`	`97`	`expect(connStringCredential).to(equal(credential), description: testCase.description)`
`94`	`98`	`}`