[PM-12012] Save region during account creation (#946)

BitwardenShared/Core/Auth/Services/API/Auth/Fixtures/APITestData+Auth.swift

@@ -4,6 +4,7 @@ extension APITestData {
     static let authRequestSuccess = loadFromJsonBundle(resource: "AuthRequest")
     static let authRequestsSuccess = loadFromJsonBundle(resource: "AuthRequests")
     static let emptyResponse = APITestData(data: "{}".data(using: .utf8)!)
+    static let nilResponse = APITestData(data: "".data(using: .utf8)!)
     static let identityTokenSuccess = loadFromJsonBundle(resource: "IdentityTokenSuccess")
     static let identityTokenWithMasterPasswordPolicy = loadFromJsonBundle(
         resource: "IdentityTokenWithMasterPasswordPolicy"

BitwardenShared/Core/Platform/Services/StateService.swift

@@ -231,6 +231,13 @@ protocol StateService: AnyObject {
     ///
     func getPreAuthEnvironmentUrls() async -> EnvironmentUrlData?
 
+    /// Gets the environment URLs for a given email during account creation.
+    ///
+    /// - Parameter email: The email used to start the account creation.
+    /// - Returns: The environment URLs used prior to start the account creation.
+    ///
+    func getAccountCreationEnvironmentUrls(email: String) async -> EnvironmentUrlData?
+
     /// Gets the server config used by the app prior to the user authenticating.
     /// - Returns: The server config used prior to user authentication.
     func getPreAuthServerConfig() async -> ServerConfig?
@@ -520,6 +527,13 @@ protocol StateService: AnyObject {
     ///
     func setPreAuthEnvironmentUrls(_ urls: EnvironmentUrlData) async
 
+    /// Sets the environment URLs for a given email during account creation.
+    /// - Parameters:
+    ///   - urls: The environment urls used to start the account creation.
+    ///   - email: The email used to start the account creation.
+    ///
+    func setAccountCreationEnvironmentUrls(urls: EnvironmentUrlData, email: String) async
+
     /// Sets the server config used prior to user authentication
     /// - Parameter config: The server config to use prior to user authentication.
     func setPreAuthServerConfig(config: ServerConfig) async
@@ -1290,6 +1304,10 @@ actor DefaultStateService: StateService { // swiftlint:disable:this type_body_le
         appSettingsStore.preAuthEnvironmentUrls
     }
 
+    func getAccountCreationEnvironmentUrls(email: String) async -> EnvironmentUrlData? {
+        appSettingsStore.accountCreationEnvironmentUrls(email: email)
+    }
+
     func getPreAuthServerConfig() async -> ServerConfig? {
         appSettingsStore.preAuthServerConfig
     }
@@ -1536,6 +1554,13 @@ actor DefaultStateService: StateService { // swiftlint:disable:this type_body_le
         appSettingsStore.preAuthEnvironmentUrls = urls
     }
 
+    func setAccountCreationEnvironmentUrls(urls: EnvironmentUrlData, email: String) async {
+        appSettingsStore.setAccountCreationEnvironmentUrls(
+            environmentUrlData: urls,
+            email: email
+        )
+    }
+
     func setPreAuthServerConfig(config: ServerConfig) async {
         appSettingsStore.preAuthServerConfig = config
     }

BitwardenShared/Core/Platform/Services/StateServiceTests.swift

@@ -707,6 +707,21 @@ class StateServiceTests: BitwardenTestCase { // swiftlint:disable:this type_body
         XCTAssertNil(urls)
     }
 
+    /// `getAccountCreationEnvironmentUrls` returns the saved pre-auth URLs for a given email.
+    func test_getAccountCreationEnvironmentUrls() async {
+        let email = "[email protected]"
+        let urls = EnvironmentUrlData(base: .example)
+        appSettingsStore.setAccountCreationEnvironmentUrls(environmentUrlData: urls, email: email)
+        let preAuthUrls = await subject.getAccountCreationEnvironmentUrls(email: email)
+        XCTAssertEqual(preAuthUrls, urls)
+    }
+
+    /// `getAccountCreationEnvironmentUrls` returns `nil` if the URLs haven't been set for a given email.
+    func test_getAccountCreationEnvironmentUrls_notSet() async {
+        let urls = await subject.getAccountCreationEnvironmentUrls(email: "[email protected]")
+        XCTAssertNil(urls)
+    }
+
     /// `getPreAuthServerConfig` returns the saved pre-auth server config.
     func test_getPreAuthServerConfig() async {
         let config = ServerConfig(
@@ -1606,6 +1621,14 @@ class StateServiceTests: BitwardenTestCase { // swiftlint:disable:this type_body
         XCTAssertEqual(appSettingsStore.preAuthEnvironmentUrls, urls)
     }
 
+    /// `test_setAccountCreationEnvironmentUrls` saves the pre-auth URLs for email for a given email.
+    func test_setAccountCreationEnvironmentUrls() async {
+        let email = "[email protected]"
+        let urls = EnvironmentUrlData(base: .example)
+        await subject.setAccountCreationEnvironmentUrls(urls: urls, email: email)
+        XCTAssertEqual(appSettingsStore.accountCreationEnvironmentUrls(email: email), urls)
+    }
+
     /// `setPreAuthServerConfig(config:)` saves the pre-auth server config.
     func test_setPreAuthServerConfig() async {
         let config = ServerConfig(

BitwardenShared/Core/Platform/Services/Stores/AppSettingsStore.swift

@@ -190,6 +190,14 @@ protocol AppSettingsStore: AnyObject {
     ///
     func pinProtectedUserKey(userId: String) -> String?
 
+    /// Gets the environment URLs used to start the account creation flow.
+    ///
+    /// - Parameters:
+    ///  - email: The email used to start the account creation.
+    /// - Returns: The environment URLs used prior to start the account creation.
+    ///
+    func accountCreationEnvironmentUrls(email: String) -> EnvironmentUrlData?
+
     /// The server configuration.
     ///
     /// - Parameter userId: The user ID associated with the server config.
@@ -344,6 +352,14 @@ protocol AppSettingsStore: AnyObject {
     ///
     func setPinProtectedUserKey(key: String?, userId: String)
 
+    /// Sets the environment URLs used to start the account creation flow.
+    ///
+    /// - Parameters:
+    ///  - email: The user's email address.
+    ///  - environmentUrlData: The environment data to be saved.
+    ///
+    func setAccountCreationEnvironmentUrls(environmentUrlData: EnvironmentUrlData, email: String)
+
     /// Sets the server config.
     ///
     /// - Parameters:
@@ -607,6 +623,7 @@ extension DefaultAppSettingsStore: AppSettingsStore {
         case passwordGenerationOptions(userId: String)
         case pinProtectedUserKey(userId: String)
         case preAuthEnvironmentUrls
+        case accountCreationEnvironmentUrls(email: String)
         case preAuthServerConfig
         case rememberedEmail
         case rememberedOrgIdentifier
@@ -682,6 +699,8 @@ extension DefaultAppSettingsStore: AppSettingsStore {
                 key = "pinKeyEncryptedUserKey_\(userId)"
             case .preAuthEnvironmentUrls:
                 key = "preAuthEnvironmentUrls"
+            case let .accountCreationEnvironmentUrls(email):
+                key = "accountCreationEnvironmentUrls_\(email)"
             case .preAuthServerConfig:
                 key = "preAuthServerConfig"
             case .rememberedEmail:
@@ -870,6 +889,12 @@ extension DefaultAppSettingsStore: AppSettingsStore {
         fetch(for: .pinProtectedUserKey(userId: userId))
     }
 
+    func accountCreationEnvironmentUrls(email: String) -> EnvironmentUrlData? {
+        fetch(
+            for: .accountCreationEnvironmentUrls(email: email)
+        )
+    }
+
     func serverConfig(userId: String) -> ServerConfig? {
         fetch(for: .serverConfig(userId: userId))
     }
@@ -952,6 +977,10 @@ extension DefaultAppSettingsStore: AppSettingsStore {
         store(key, for: .pinProtectedUserKey(userId: userId))
     }
 
+    func setAccountCreationEnvironmentUrls(environmentUrlData: EnvironmentUrlData, email: String) {
+        store(environmentUrlData, for: .accountCreationEnvironmentUrls(email: email))
+    }
+
     func setServerConfig(_ config: ServerConfig?, userId: String) {
         store(config, for: .serverConfig(userId: userId))
     }

BitwardenShared/Core/Platform/Services/Stores/AppSettingsStoreTests.swift

@@ -616,6 +616,43 @@ class AppSettingsStoreTests: BitwardenTestCase { // swiftlint:disable:this type_
         )
     }
 
+    /// `accountCreationEnvironmentUrls` returns `nil` if there isn't a previously stored value.
+    func test_accountCreationEnvironmentUrls_isInitiallyNil() {
+        XCTAssertNil(subject.accountCreationEnvironmentUrls(email: "[email protected]"))
+    }
+
+    /// `accountCreationEnvironmentUrls` can be used to get and set the persisted value in user defaults.
+    func test_accountCreationEnvironmentUrls_withValue() {
+        let email = "[email protected]"
+        subject.setAccountCreationEnvironmentUrls(environmentUrlData: .defaultUS, email: email)
+        XCTAssertEqual(subject.accountCreationEnvironmentUrls(email: email), .defaultUS)
+        try XCTAssertEqual(
+            JSONDecoder().decode(
+                EnvironmentUrlData.self,
+                from: XCTUnwrap(
+                    userDefaults
+                        .string(forKey: "bwPreferencesStorage:accountCreationEnvironmentUrls_\(email)")?
+                        .data(using: .utf8)
+                )
+            ),
+            .defaultUS
+        )
+
+        subject.setAccountCreationEnvironmentUrls(environmentUrlData: .defaultEU, email: email)
+        XCTAssertEqual(subject.accountCreationEnvironmentUrls(email: email), .defaultEU)
+        try XCTAssertEqual(
+            JSONDecoder().decode(
+                EnvironmentUrlData.self,
+                from: XCTUnwrap(
+                    userDefaults
+                        .string(forKey: "bwPreferencesStorage:accountCreationEnvironmentUrls_\(email)")?
+                        .data(using: .utf8)
+                )
+            ),
+            .defaultEU
+        )
+    }
+
     /// `preAuthServerConfig` is initially `nil`
     func test_preAuthServerConfig_isInitiallyNil() {
         XCTAssertNil(subject.preAuthServerConfig)

BitwardenShared/Core/Platform/Services/Stores/TestHelpers/MockAppSettingsStore.swift

@@ -37,6 +37,7 @@ class MockAppSettingsStore: AppSettingsStore {
     var notificationsLastRegistrationDates = [String: Date]()
     var passwordGenerationOptions = [String: PasswordGenerationOptions]()
     var pinProtectedUserKey = [String: String]()
+    var accountCreationEnvironmentUrls = [String: EnvironmentUrlData]()
     var serverConfig = [String: ServerConfig]()
     var shouldTrustDevice = [String: Bool?]()
     var timeoutAction = [String: Int]()
@@ -118,6 +119,10 @@ class MockAppSettingsStore: AppSettingsStore {
         pinProtectedUserKey[userId]
     }
 
+    func accountCreationEnvironmentUrls(email: String) -> BitwardenShared.EnvironmentUrlData? {
+        accountCreationEnvironmentUrls[email]
+    }
+
     func twoFactorToken(email: String) -> String? {
         twoFactorTokens[email]
     }
@@ -202,6 +207,10 @@ class MockAppSettingsStore: AppSettingsStore {
         pinProtectedUserKey[userId] = key
     }
 
+    func setAccountCreationEnvironmentUrls(environmentUrlData: BitwardenShared.EnvironmentUrlData, email: String) {
+        accountCreationEnvironmentUrls[email] = environmentUrlData
+    }
+
     func setServerConfig(_ config: ServerConfig?, userId: String) {
         serverConfig[userId] = config
     }

BitwardenShared/Core/Platform/Services/TestHelpers/MockStateService.swift

@@ -55,6 +55,7 @@ class MockStateService: StateService { // swiftlint:disable:this type_body_lengt
     var passwordGenerationOptions = [String: PasswordGenerationOptions]()
     var pinProtectedUserKeyValue = [String: String]()
     var preAuthEnvironmentUrls: EnvironmentUrlData?
+    var accountCreationEnvironmentUrls = [String: EnvironmentUrlData]()
     var preAuthServerConfig: ServerConfig?
     var rememberedOrgIdentifier: String?
     var showWebIcons = true
@@ -255,6 +256,10 @@ class MockStateService: StateService { // swiftlint:disable:this type_body_lengt
         preAuthEnvironmentUrls
     }
 
+    func getAccountCreationEnvironmentUrls(email: String) async -> EnvironmentUrlData? {
+        accountCreationEnvironmentUrls[email]
+    }
+
     func getPreAuthServerConfig() async -> BitwardenShared.ServerConfig? {
         preAuthServerConfig
     }
@@ -479,6 +484,10 @@ class MockStateService: StateService { // swiftlint:disable:this type_body_lengt
         preAuthEnvironmentUrls = urls
     }
 
+    func setAccountCreationEnvironmentUrls(urls: BitwardenShared.EnvironmentUrlData, email: String) async {
+        accountCreationEnvironmentUrls[email] = urls
+    }
+
     func setPreAuthServerConfig(config: BitwardenShared.ServerConfig) async {
         preAuthServerConfig = config
     }

BitwardenShared/UI/Auth/CompleteRegistration/CompleteRegistrationProcessor.swift

@@ -17,6 +17,9 @@ enum CompleteRegistrationError: Error {
 
     /// The password does not meet the minimum length requirement.
     case passwordIsTooShort
+
+    /// The environment urls to complete the registration are empty
+    case preAuthUrlsEmpty
 }
 
 // MARK: - CompleteRegistrationProcessor
@@ -244,13 +247,24 @@ class CompleteRegistrationProcessor: StateProcessor<
         }
     }
 
-    /// Sets the URLs to use.
+    /// Sets the URLs to use if user came from an email link.
     ///
     private func setRegion() async {
-        guard state.region != nil,
-              let urls = state.region?.defaultURLs else { return }
+        do {
+            guard state.fromEmail else {
+                return
+            }
+
+            guard let urls = await services.stateService.getAccountCreationEnvironmentUrls(email: state.userEmail) else {
+                throw CompleteRegistrationError.preAuthUrlsEmpty
+            }
 
-        await services.environmentService.setPreAuthURLs(urls: urls)
+            await services.environmentService.setPreAuthURLs(urls: urls)
+        } catch let error as CompleteRegistrationError {
+            showCompleteRegistrationErrorAlert(error)
+        } catch {
+            coordinator.showAlert(.defaultAlert(title: Localizations.anErrorHasOccurred))
+        }
     }
 
     /// Sets the feature flag value to be used.
@@ -274,6 +288,11 @@ class CompleteRegistrationProcessor: StateProcessor<
             coordinator.showAlert(.validationFieldRequired(fieldName: Localizations.masterPassword))
         case .passwordIsTooShort:
             coordinator.showAlert(.passwordIsTooShort)
+        case .preAuthUrlsEmpty:
+            coordinator.showAlert(.defaultAlert(
+                title: Localizations.anErrorHasOccurred,
+                message: Localizations.theRegionForTheGivenEmailCouldNotBeLoaded
+            ))
         }
     }