PM-30389: Allow for different auth tab schemes (#6315)

app/src/main/kotlin/com/x8bit/bitwarden/ui/auth/feature/enterprisesignon/EnterpriseSignOnScreen.kt

@@ -64,7 +64,11 @@ fun EnterpriseSignOnScreen(
             EnterpriseSignOnEvent.NavigateBack -> onNavigateBack()
 
             is EnterpriseSignOnEvent.NavigateToSsoLogin -> {
-                intentManager.startAuthTab(uri = event.uri, launcher = authTabLaunchers.sso)
+                intentManager.startAuthTab(
+                    uri = event.uri,
+                    redirectScheme = event.scheme,
+                    launcher = authTabLaunchers.sso,
+                )
             }
 
             is EnterpriseSignOnEvent.NavigateToSetPassword -> {

app/src/main/kotlin/com/x8bit/bitwarden/ui/auth/feature/enterprisesignon/EnterpriseSignOnViewModel.kt

@@ -206,7 +206,12 @@ class EnterpriseSignOnViewModel @Inject constructor(
         action: EnterpriseSignOnAction.Internal.OnGenerateUriForSsoResult,
     ) {
         mutableStateFlow.update { it.copy(dialogState = null) }
-        sendEvent(EnterpriseSignOnEvent.NavigateToSsoLogin(action.uri))
+        sendEvent(
+            EnterpriseSignOnEvent.NavigateToSsoLogin(
+                uri = action.uri,
+                scheme = action.scheme,
+            ),
+        )
     }
 
     private fun handleOnSsoPrevalidationFailure(
@@ -401,7 +406,12 @@ class EnterpriseSignOnViewModel @Inject constructor(
 
         // Hide any dialog since we're about to launch a custom tab and could return without getting
         // a result due to user intervention
-        sendAction(EnterpriseSignOnAction.Internal.OnGenerateUriForSsoResult(uri.toUri()))
+        sendAction(
+            EnterpriseSignOnAction.Internal.OnGenerateUriForSsoResult(
+                uri = uri.toUri(),
+                scheme = "bitwarden",
+            ),
+        )
     }
 
     private fun showError(
@@ -507,7 +517,10 @@ sealed class EnterpriseSignOnEvent {
     /**
      * Navigates to a custom tab for SSO login using [uri].
      */
-    data class NavigateToSsoLogin(val uri: Uri) : EnterpriseSignOnEvent()
+    data class NavigateToSsoLogin(
+        val uri: Uri,
+        val scheme: String,
+    ) : EnterpriseSignOnEvent()
 
     /**
      * Navigates to the set master password screen.
@@ -568,7 +581,7 @@ sealed class EnterpriseSignOnAction {
         /**
          * A [uri] has been generated to request an SSO result.
          */
-        data class OnGenerateUriForSsoResult(val uri: Uri) : Internal()
+        data class OnGenerateUriForSsoResult(val uri: Uri, val scheme: String) : Internal()
 
         /**
          * A login result has been received.

app/src/main/kotlin/com/x8bit/bitwarden/ui/auth/feature/twofactorlogin/TwoFactorLoginScreen.kt

@@ -102,11 +102,19 @@ fun TwoFactorLoginScreen(
             }
 
             is TwoFactorLoginEvent.NavigateToDuo -> {
-                intentManager.startAuthTab(uri = event.uri, launcher = authTabLaunchers.duo)
+                intentManager.startAuthTab(
+                    uri = event.uri,
+                    redirectScheme = event.scheme,
+                    launcher = authTabLaunchers.duo,
+                )
             }
 
             is TwoFactorLoginEvent.NavigateToWebAuth -> {
-                intentManager.startAuthTab(uri = event.uri, launcher = authTabLaunchers.webAuthn)
+                intentManager.startAuthTab(
+                    uri = event.uri,
+                    redirectScheme = event.scheme,
+                    launcher = authTabLaunchers.webAuthn,
+                )
             }
 
             is TwoFactorLoginEvent.ShowSnackbar -> snackbarHostState.showSnackbar(event.data)

app/src/main/kotlin/com/x8bit/bitwarden/ui/auth/feature/twofactorlogin/TwoFactorLoginViewModel.kt

@@ -175,7 +175,7 @@ class TwoFactorLoginViewModel @Inject constructor(
     /**
      * Navigates to the Duo webpage if appropriate, else processes the login.
      */
-    @Suppress("MaxLineLength")
+    @Suppress("LongMethod")
     private fun handleContinueButtonClick() {
         when (state.authMethod) {
             TwoFactorAuthMethod.DUO,
@@ -185,13 +185,21 @@ class TwoFactorLoginViewModel @Inject constructor(
                 // The url should not be empty unless the environment is somehow not supported.
                 authUrl
                     ?.let {
-                        sendEvent(event = TwoFactorLoginEvent.NavigateToDuo(uri = it.toUri()))
+                        sendEvent(
+                            event = TwoFactorLoginEvent.NavigateToDuo(
+                                uri = it.toUri(),
+                                scheme = "bitwarden",
+                            ),
+                        )
                     }
                     ?: mutableStateFlow.update {
+                        @Suppress("MaxLineLength")
                         it.copy(
                             dialogState = TwoFactorLoginState.DialogState.Error(
                                 title = BitwardenString.an_error_has_occurred.asText(),
-                                message = BitwardenString.error_connecting_with_the_duo_service_use_a_different_two_step_login_method_or_contact_duo_for_assistance.asText(),
+                                message = BitwardenString
+                                    .error_connecting_with_the_duo_service_use_a_different_two_step_login_method_or_contact_duo_for_assistance
+                                    .asText(),
                             ),
                         )
                     }
@@ -220,10 +228,12 @@ class TwoFactorLoginViewModel @Inject constructor(
                                     resId = BitwardenString.fido2_return_to_app,
                                 ),
                             )
-                            TwoFactorLoginEvent.NavigateToWebAuth(uri = uri)
+                            TwoFactorLoginEvent.NavigateToWebAuth(uri = uri, scheme = "bitwarden")
                         }
                         ?: TwoFactorLoginEvent.ShowSnackbar(
-                            message = BitwardenString.there_was_an_error_starting_web_authn_two_factor_authentication.asText(),
+                            message = BitwardenString
+                                .there_was_an_error_starting_web_authn_two_factor_authentication
+                                .asText(),
                         ),
                 )
             }
@@ -667,12 +677,12 @@ sealed class TwoFactorLoginEvent {
     /**
      * Navigates to the Duo 2-factor authentication screen.
      */
-    data class NavigateToDuo(val uri: Uri) : TwoFactorLoginEvent()
+    data class NavigateToDuo(val uri: Uri, val scheme: String) : TwoFactorLoginEvent()
 
     /**
      * Navigates to the WebAuth authentication screen.
      */
-    data class NavigateToWebAuth(val uri: Uri) : TwoFactorLoginEvent()
+    data class NavigateToWebAuth(val uri: Uri, val scheme: String) : TwoFactorLoginEvent()
 
     /**
      * Navigates to the recovery code help page.

app/src/test/kotlin/com/x8bit/bitwarden/ui/auth/feature/enterprisesignon/EnterpriseSignOnScreenTest.kt

@@ -45,7 +45,7 @@ class EnterpriseSignOnScreenTest : BitwardenComposeTest() {
     }
 
     private val intentManager: IntentManager = mockk {
-        every { startAuthTab(uri = any(), launcher = any()) } just runs
+        every { startAuthTab(uri = any(), redirectScheme = any(), launcher = any()) } just runs
     }
 
     @Before
@@ -114,9 +114,14 @@ class EnterpriseSignOnScreenTest : BitwardenComposeTest() {
     @Test
     fun `NavigateToSsoLogin should call startCustomTabsActivity`() {
         val ssoUri = Uri.parse("https://identity.bitwarden.com/sso-test")
-        mutableEventFlow.tryEmit(EnterpriseSignOnEvent.NavigateToSsoLogin(ssoUri))
+        val scheme = "bitwarden"
+        mutableEventFlow.tryEmit(EnterpriseSignOnEvent.NavigateToSsoLogin(ssoUri, scheme))
         verify(exactly = 1) {
-            intentManager.startAuthTab(uri = ssoUri, launcher = ssoLauncher)
+            intentManager.startAuthTab(
+                uri = ssoUri,
+                redirectScheme = scheme,
+                launcher = ssoLauncher,
+            )
         }
     }
 

app/src/test/kotlin/com/x8bit/bitwarden/ui/auth/feature/enterprisesignon/EnterpriseSignOnViewModelTest.kt

@@ -172,28 +172,27 @@ class EnterpriseSignOnViewModelTest : BaseViewModelTest() {
             } returns ssoUri
 
             val viewModel = createViewModel(state)
-            viewModel.stateFlow.test {
-                assertEquals(state, awaitItem())
+            viewModel.stateEventFlow(backgroundScope) { stateFlow, eventFlow ->
+                assertEquals(state, stateFlow.awaitItem())
                 viewModel.trySendAction(EnterpriseSignOnAction.LogInClick)
 
                 assertEquals(
                     state.copy(
                         dialogState = EnterpriseSignOnState.DialogState.Loading(
-                            BitwardenString.logging_in.asText(),
+                            message = BitwardenString.logging_in.asText(),
                         ),
                     ),
-                    awaitItem(),
+                    stateFlow.awaitItem(),
                 )
 
                 assertEquals(
                     state.copy(dialogState = null),
-                    awaitItem(),
+                    stateFlow.awaitItem(),
                 )
-            }
-            viewModel.eventFlow.test {
+
                 assertEquals(
-                    EnterpriseSignOnEvent.NavigateToSsoLogin(ssoUri),
-                    awaitItem(),
+                    EnterpriseSignOnEvent.NavigateToSsoLogin(uri = ssoUri, scheme = "bitwarden"),
+                    eventFlow.awaitItem(),
                 )
             }
         }

app/src/test/kotlin/com/x8bit/bitwarden/ui/auth/feature/twofactorlogin/TwoFactorLoginScreenTest.kt

@@ -39,7 +39,7 @@ class TwoFactorLoginScreenTest : BitwardenComposeTest() {
     private val webAuthnLauncher: ActivityResultLauncher<Intent> = mockk()
     private val intentManager = mockk<IntentManager> {
         every { launchUri(uri = any()) } just runs
-        every { startAuthTab(uri = any(), launcher = any()) } just runs
+        every { startAuthTab(uri = any(), redirectScheme = any(), launcher = any()) } just runs
     }
     private val nfcManager: NfcManager = mockk {
         every { start() } just runs
@@ -283,15 +283,29 @@ class TwoFactorLoginScreenTest : BitwardenComposeTest() {
     @Test
     fun `NavigateToDuo should call intentManager startAuthTab`() {
         val mockUri = mockk<Uri>()
-        mutableEventFlow.tryEmit(TwoFactorLoginEvent.NavigateToDuo(mockUri))
-        verify { intentManager.startAuthTab(uri = mockUri, launcher = duoLauncher) }
+        val scheme = "bitwarden"
+        mutableEventFlow.tryEmit(TwoFactorLoginEvent.NavigateToDuo(mockUri, scheme))
+        verify(exactly = 1) {
+            intentManager.startAuthTab(
+                uri = mockUri,
+                redirectScheme = scheme,
+                launcher = duoLauncher,
+            )
+        }
     }
 
     @Test
     fun `NavigateToWebAuth should call intentManager startCustomTabsActivity`() {
         val mockUri = mockk<Uri>()
-        mutableEventFlow.tryEmit(TwoFactorLoginEvent.NavigateToWebAuth(mockUri))
-        verify { intentManager.startAuthTab(uri = mockUri, launcher = webAuthnLauncher) }
+        val scheme = "bitwarden"
+        mutableEventFlow.tryEmit(TwoFactorLoginEvent.NavigateToWebAuth(mockUri, scheme))
+        verify(exactly = 1) {
+            intentManager.startAuthTab(
+                uri = mockUri,
+                redirectScheme = scheme,
+                launcher = webAuthnLauncher,
+            )
+        }
     }
 
     @Test

app/src/test/kotlin/com/x8bit/bitwarden/ui/auth/feature/twofactorlogin/TwoFactorLoginViewModelTest.kt

@@ -427,7 +427,7 @@ class TwoFactorLoginViewModelTest : BaseViewModelTest() {
             viewModel.eventFlow.test {
                 viewModel.trySendAction(TwoFactorLoginAction.ContinueButtonClick)
                 assertEquals(
-                    TwoFactorLoginEvent.NavigateToDuo(mockkUri),
+                    TwoFactorLoginEvent.NavigateToDuo(uri = mockkUri, scheme = "bitwarden"),
                     awaitItem(),
                 )
             }
@@ -512,7 +512,7 @@ class TwoFactorLoginViewModelTest : BaseViewModelTest() {
             viewModel.eventFlow.test {
                 viewModel.trySendAction(TwoFactorLoginAction.ContinueButtonClick)
                 assertEquals(
-                    TwoFactorLoginEvent.NavigateToWebAuth(mockkUri),
+                    TwoFactorLoginEvent.NavigateToWebAuth(uri = mockkUri, scheme = "bitwarden"),
                     awaitItem(),
                 )
             }