From c0c88f5108ba417bd89120f8fc192ffdb11c51d6 Mon Sep 17 00:00:00 2001
From: David Perez <david@livefront.com>
Date: Thu, 22 Feb 2024 14:05:13 -0600
Subject: [PATCH] BIT-1898: Two-factor auth should have access to device data
 if available (#1054)

---
 .../data/auth/repository/AuthRepositoryImpl.kt   |  9 +++++++++
 .../data/auth/repository/AuthRepositoryTest.kt   | 16 ++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryImpl.kt b/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryImpl.kt
index 7518ee49bf..e5e9840804 100644
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryImpl.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryImpl.kt
@@ -129,6 +129,12 @@ class AuthRepositoryImpl(
      */
     private var identityTokenAuthModel: IdentityTokenAuthModel? = null
 
+    /**
+     * The device auth information to unlock the vault when logging in with device in the case
+     * of two-factor authentication.
+     */
+    private var twoFactorDeviceData: DeviceDataModel? = null
+
     /**
      * The information necessary to resend the verification code email for two-factor login.
      */
@@ -381,6 +387,7 @@ class AuthRepositoryImpl(
                 authModel = it,
                 twoFactorData = twoFactorData,
                 captchaToken = captchaToken ?: twoFactorResponse?.captchaToken,
+                deviceData = twoFactorDeviceData,
             )
         }
         ?: LoginResult.Error(errorMessage = null)
@@ -431,6 +438,7 @@ class AuthRepositoryImpl(
                         // Cache the data necessary for the remaining two-factor auth flow.
                         identityTokenAuthModel = authModel
                         twoFactorResponse = loginResponse
+                        twoFactorDeviceData = deviceData
                         resendEmailRequestJson = ResendEmailRequestJson(
                             deviceIdentifier = authDiskSource.uniqueAppId,
                             email = email,
@@ -468,6 +476,7 @@ class AuthRepositoryImpl(
                         identityTokenAuthModel = null
                         twoFactorResponse = null
                         resendEmailRequestJson = null
+                        twoFactorDeviceData = null
 
                         // Attempt to unlock the vault with password if possible.
                         password?.let {
diff --git a/app/src/test/java/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryTest.kt b/app/src/test/java/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryTest.kt
index 0dc83cc16e..066d8f56f4 100644
--- a/app/src/test/java/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryTest.kt
+++ b/app/src/test/java/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryTest.kt
@@ -1532,6 +1532,22 @@ class AuthRepositoryTest {
                 environmentUrlData = EnvironmentUrlDataJson.DEFAULT_US,
             )
         } returns SINGLE_USER_STATE_1
+        coEvery {
+            vaultRepository.unlockVault(
+                userId = SINGLE_USER_STATE_1.activeUserId,
+                email = SINGLE_USER_STATE_1.activeAccount.profile.email,
+                kdf = SINGLE_USER_STATE_1.activeAccount.profile.toSdkParams(),
+                privateKey = successResponse.privateKey,
+                initUserCryptoMethod = InitUserCryptoMethod.AuthRequest(
+                    requestPrivateKey = DEVICE_REQUEST_PRIVATE_KEY,
+                    method = AuthRequestMethod.MasterKey(
+                        protectedMasterKey = DEVICE_ASYMMETRICAL_KEY,
+                        authRequestKey = successResponse.key,
+                    ),
+                ),
+                organizationKeys = null,
+            )
+        } returns VaultUnlockResult.Success
         val finalResult = repository.login(
             email = EMAIL,
             password = null,