PM-10909: Add persistance layer for usersKeyConnector (#3740)

app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/disk/AuthDiskSource.kt

@@ -45,6 +45,16 @@ interface AuthDiskSource {
      */
     fun clearData(userId: String)
 
+    /**
+     * Retrieves the state indicating that the user should use a key connector.
+     */
+    fun getShouldUseKeyConnector(userId: String): Boolean?
+
+    /**
+     * Stores the boolean indicating that the user should use a key connector.
+     */
+    fun storeShouldUseKeyConnector(userId: String, shouldUseKeyConnector: Boolean?)
+
     /**
      * Retrieves the state indicating that the user has chosen to trust this device.
      *

app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/disk/AuthDiskSourceImpl.kt

@@ -39,6 +39,7 @@ private const val TWO_FACTOR_TOKEN_KEY = "twoFactorToken"
 private const val MASTER_PASSWORD_HASH_KEY = "keyHash"
 private const val POLICIES_KEY = "policies"
 private const val SHOULD_TRUST_DEVICE_KEY = "shouldTrustDevice"
+private const val USES_KEY_CONNECTOR = "usesKeyConnector"
 
 /**
  * Primary implementation of [AuthDiskSource].
@@ -122,11 +123,23 @@ class AuthDiskSourceImpl(
         storeMasterPasswordHash(userId = userId, passwordHash = null)
         storePolicies(userId = userId, policies = null)
         storeAccountTokens(userId = userId, accountTokens = null)
+        storeShouldUseKeyConnector(userId = userId, shouldUseKeyConnector = null)
 
         // Do not remove the DeviceKey or PendingAuthRequest on logout, these are persisted
         // indefinitely unless the TDE flow explicitly removes them.
     }
 
+    override fun getShouldUseKeyConnector(
+        userId: String,
+    ): Boolean? = getBoolean(key = USES_KEY_CONNECTOR.appendIdentifier(userId))
+
+    override fun storeShouldUseKeyConnector(userId: String, shouldUseKeyConnector: Boolean?) {
+        putBoolean(
+            key = USES_KEY_CONNECTOR.appendIdentifier(userId),
+            value = shouldUseKeyConnector,
+        )
+    }
+
     override fun getShouldTrustDevice(userId: String): Boolean =
         requireNotNull(
             getBoolean(key = SHOULD_TRUST_DEVICE_KEY.appendIdentifier(userId), default = false),

app/src/main/java/com/x8bit/bitwarden/data/vault/repository/VaultRepositoryImpl.kt

@@ -898,27 +898,29 @@ class VaultRepositoryImpl(
     ) {
         val profile = syncResponse.profile
         val userId = profile.id
-        val userKey = profile.key
-        val privateKey = profile.privateKey
         authDiskSource.apply {
             storeUserKey(
                 userId = userId,
-                userKey = userKey,
+                userKey = profile.key,
             )
             storePrivateKey(
                 userId = userId,
-                privateKey = privateKey,
+                privateKey = profile.privateKey,
             )
             storeOrganizationKeys(
-                userId = profile.id,
+                userId = userId,
                 organizationKeys = profile.organizations
                     .orEmpty()
                     .filter { it.key != null }
                     .associate { it.id to requireNotNull(it.key) },
             )
+            storeShouldUseKeyConnector(
+                userId = userId,
+                shouldUseKeyConnector = profile.shouldUseKeyConnector,
+            )
             storeOrganizations(
-                userId = profile.id,
+                userId = userId,
-                organizations = syncResponse.profile.organizations,
+                organizations = profile.organizations,
             )
         }
     }

app/src/test/java/com/x8bit/bitwarden/data/auth/datasource/disk/AuthDiskSourceTest.kt

@@ -121,6 +121,24 @@ class AuthDiskSourceTest {
         assertNull(authDiskSource.rememberedOrgIdentifier)
     }
 
+    @Test
+    fun `shouldUseKeyConnector should pull from and update SharedPreferences`() {
+        val userId = "userId"
+        val shouldUseKeyConnectorKey = "bwPreferencesStorage:usesKeyConnector_$userId"
+
+        // Shared preferences and the disk source start with the same value.
+        assertNull(authDiskSource.getShouldUseKeyConnector(userId = userId))
+        assertFalse(fakeSharedPreferences.getBoolean(shouldUseKeyConnectorKey, false))
+
+        // Updating the disk source updates shared preferences
+        authDiskSource.storeShouldUseKeyConnector(userId = userId, shouldUseKeyConnector = true)
+        assertTrue(fakeSharedPreferences.getBoolean(shouldUseKeyConnectorKey, false))
+
+        // Update SharedPreferences updates the disk source
+        fakeSharedPreferences.edit { putBoolean(shouldUseKeyConnectorKey, false) }
+        assertFalse(authDiskSource.getShouldUseKeyConnector(userId = userId) ?: true)
+    }
+
     @Test
     fun `shouldTrustDevice should pull from and update SharedPreferences`() {
         val userId = "userId"
@@ -191,6 +209,7 @@ class AuthDiskSourceTest {
             userId = userId,
             pendingAuthRequest = pendingAuthRequestJson,
         )
+        authDiskSource.storeShouldUseKeyConnector(userId = userId, shouldUseKeyConnector = true)
         val shouldTrustDevice = true
         authDiskSource.storeShouldTrustDevice(
             userId = userId,
@@ -258,6 +277,7 @@ class AuthDiskSourceTest {
         assertNull(authDiskSource.getAccountTokens(userId = userId))
         assertNull(authDiskSource.getEncryptedPin(userId = userId))
         assertNull(authDiskSource.getMasterPasswordHash(userId = userId))
+        assertNull(authDiskSource.getShouldUseKeyConnector(userId = userId))
     }
 
     @Test

app/src/test/java/com/x8bit/bitwarden/data/auth/datasource/disk/util/FakeAuthDiskSource.kt

@@ -26,6 +26,7 @@ class FakeAuthDiskSource : AuthDiskSource {
         mutableMapOf<String, MutableSharedFlow<AccountTokensJson?>>()
     private val mutableUserStateFlow = bufferedMutableSharedFlow<UserStateJson?>(replay = 1)
 
+    private val storedShouldUseKeyConnector = mutableMapOf<String, Boolean?>()
     private val storedShouldTrustDevice = mutableMapOf<String, Boolean?>()
     private val storedInvalidUnlockAttempts = mutableMapOf<String, Int?>()
     private val storedUserKeys = mutableMapOf<String, String?>()
@@ -72,6 +73,14 @@ class FakeAuthDiskSource : AuthDiskSource {
         mutableAccountTokensFlowMap.remove(userId)
     }
 
+    override fun getShouldUseKeyConnector(
+        userId: String,
+    ): Boolean = storedShouldUseKeyConnector[userId] ?: false
+
+    override fun storeShouldUseKeyConnector(userId: String, shouldUseKeyConnector: Boolean?) {
+        storedShouldUseKeyConnector[userId] = shouldUseKeyConnector
+    }
+
     override fun getShouldTrustDevice(userId: String): Boolean =
         storedShouldTrustDevice[userId] ?: false
 
@@ -214,6 +223,13 @@ class FakeAuthDiskSource : AuthDiskSource {
         getMutableAccountTokensFlow(userId = userId).tryEmit(accountTokens)
     }
 
+    /**
+     * Assert the the [shouldUseKeyConnector] was stored successfully using the [userId].
+     */
+    fun assertShouldUseKeyConnector(userId: String, shouldUseKeyConnector: Boolean?) {
+        assertEquals(shouldUseKeyConnector, storedShouldUseKeyConnector[userId])
+    }
+
     /**
      * Assert that the given [userState] matches the currently tracked value.
      */

app/src/test/java/com/x8bit/bitwarden/data/vault/repository/VaultRepositoryTest.kt

@@ -826,6 +826,10 @@ class VaultRepositoryTest {
                 userId = "mockId-1",
                 policies = listOf(createMockPolicy(number = 1)),
             )
+            fakeAuthDiskSource.assertShouldUseKeyConnector(
+                userId = "mockId-1",
+                shouldUseKeyConnector = false,
+            )
             coVerify {
                 vaultDiskSource.replaceVaultData(
                     userId = MOCK_USER_STATE.activeUserId,