PM-11434: Remember SSO org identifier even if redirected (#882)

bitwarden · Aug 29, 2024 · ceb21b9 · ceb21b9
1 parent c14e27b
commit ceb21b9
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 4 deletions.
diff --git a/BitwardenShared/UI/Auth/Login/SingleSignOn/SingleSignOnProcessor.swift b/BitwardenShared/UI/Auth/Login/SingleSignOn/SingleSignOnProcessor.swift
@@ -90,13 +90,15 @@ final class SingleSignOnProcessor: StateProcessor<SingleSignOnState, SingleSignO
         case ASWebAuthenticationSessionError.canceledLogin:
             break
         case let IdentityTokenRequestError.twoFactorRequired(authMethodsData, _, _):
-            coordinator.navigate(to: .twoFactor(state.email, nil, authMethodsData, state.identifierText))
+            rememberOrgIdentifierAndNavigate(to: .twoFactor(state.email, nil, authMethodsData, state.identifierText))
         case AuthError.requireSetPassword:
-            coordinator.navigate(to: .setMasterPassword(organizationIdentifier: state.identifierText))
+            rememberOrgIdentifierAndNavigate(to: .setMasterPassword(organizationIdentifier: state.identifierText))
         case AuthError.requireUpdatePassword:
-            coordinator.navigate(to: .updateMasterPassword)
+            rememberOrgIdentifierAndNavigate(to: .updateMasterPassword)
         case AuthError.requireDecryptionOptions:
-            coordinator.navigate(to: .showLoginDecryptionOptions(organizationIdentifier: state.identifierText))
+            rememberOrgIdentifierAndNavigate(to: .showLoginDecryptionOptions(
+                organizationIdentifier: state.identifierText
+            ))
         default:
             coordinator.showAlert(.networkResponseError(error, tryAgain))
             services.errorReporter.log(error: error)
@@ -154,6 +156,15 @@ final class SingleSignOnProcessor: StateProcessor<SingleSignOnState, SingleSignO
             services.errorReporter.log(error: error)
         }
     }
+
+    /// Remembers the org identifier for future logins and navigates to the specified route.
+    ///
+    /// - Parameter route: The route to navigate to after saving the org identifier.
+    ///
+    private func rememberOrgIdentifierAndNavigate(to route: AuthRoute) {
+        services.stateService.rememberedOrgIdentifier = state.identifierText
+        coordinator.navigate(to: route)
+    }
 }
 
 // MARK: - SingleSignOnFlowDelegate

diff --git a/BitwardenShared/UI/Auth/Login/SingleSignOn/SingleSignOnProcessorTests.swift b/BitwardenShared/UI/Auth/Login/SingleSignOn/SingleSignOnProcessorTests.swift
@@ -171,6 +171,7 @@ class SingleSignOnProcessorTests: BitwardenTestCase {
         XCTAssertFalse(coordinator.isLoadingOverlayShowing)
         XCTAssertEqual(coordinator.alertShown.last, .networkResponseError(BitwardenTestError.example))
         XCTAssertEqual(errorReporter.errors.last as? BitwardenTestError, .example)
+        XCTAssertNil(stateService.rememberedOrgIdentifier)
     }
 
     /// `singleSignOnCompleted(code:)` navigates to the two-factor view if two-factor authentication is needed.
@@ -186,6 +187,7 @@ class SingleSignOnProcessorTests: BitwardenTestCase {
 
         // Verify the results.
         XCTAssertEqual(coordinator.routes.last, .twoFactor("", nil, AuthMethodsData(), "BestOrganization"))
+        XCTAssertEqual(stateService.rememberedOrgIdentifier, "BestOrganization")
     }
 
     /// `singleSignOnCompleted(code:)` navigates to the set password screen if the user needs
@@ -200,6 +202,37 @@ class SingleSignOnProcessorTests: BitwardenTestCase {
         waitFor(!coordinator.routes.isEmpty)
 
         XCTAssertEqual(coordinator.routes, [.setMasterPassword(organizationIdentifier: "BestOrganization")])
+        XCTAssertEqual(stateService.rememberedOrgIdentifier, "BestOrganization")
+    }
+
+    /// `singleSignOnCompleted(code:)` navigates to the update password screen if the user needs
+    /// to update their master password.
+    @MainActor
+    func test_singleSignOnCompleted_requireUpdatePasswordError() {
+        authService.loginWithSingleSignOnResult = .failure(AuthError.requireUpdatePassword)
+        subject.state.identifierText = "BestOrganization"
+
+        subject.singleSignOnCompleted(code: "CODE")
+
+        waitFor(!coordinator.routes.isEmpty)
+
+        XCTAssertEqual(coordinator.routes, [.updateMasterPassword])
+        XCTAssertEqual(stateService.rememberedOrgIdentifier, "BestOrganization")
+    }
+
+    /// `singleSignOnCompleted(code:)` navigates to the show login decryption options screen if the
+    /// user needs to choose their decryption option for login.
+    @MainActor
+    func test_singleSignOnCompleted_requireDecryptionOptionsError() {
+        authService.loginWithSingleSignOnResult = .failure(AuthError.requireDecryptionOptions)
+        subject.state.identifierText = "BestOrganization"
+
+        subject.singleSignOnCompleted(code: "CODE")
+
+        waitFor(!coordinator.routes.isEmpty)
+
+        XCTAssertEqual(coordinator.routes, [.showLoginDecryptionOptions(organizationIdentifier: "BestOrganization")])
+        XCTAssertEqual(stateService.rememberedOrgIdentifier, "BestOrganization")
     }
 
     /// `singleSignOnCompleted(code:)` navigates to the vault unlock view if the vault is still locked.