[PM-24569] Save accountKeys to AuthDiskSource (#5679)

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

@@ -126,13 +126,34 @@ interface AuthDiskSource : AppIdProvider {
     /**
      * Retrieves a private key using a [userId].
      */
+    @Deprecated(
+        message = "Use getAccountKeys instead.",
+        replaceWith = ReplaceWith("getAccountKeys"),
+    )
     fun getPrivateKey(userId: String): String?
 
     /**
      * Stores a private key using a [userId].
      */
+    @Deprecated(
+        message = "Use storeAccountKeys instead.",
+        replaceWith = ReplaceWith("storeAccountKeys"),
+    )
     fun storePrivateKey(userId: String, privateKey: String?)
 
+    /**
+     * Returns the profile account keys for the given [userId].
+     */
+    fun getAccountKeys(userId: String): SyncResponseJson.Profile.AccountKeys?
+
+    /**
+     * Stores the profile account keys for the given [userId].
+     */
+    fun storeAccountKeys(
+        userId: String,
+        accountKeys: SyncResponseJson.Profile.AccountKeys?,
+    )
+
     /**
      * Retrieves a user auto-unlock key for the given [userId].
      */

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

@@ -48,6 +48,7 @@ private const val USES_KEY_CONNECTOR = "usesKeyConnector"
 private const val ONBOARDING_STATUS_KEY = "onboardingStatus"
 private const val SHOW_IMPORT_LOGINS_KEY = "showImportLogins"
 private const val LAST_LOCK_TIMESTAMP = "lastLockTimestamp"
+private const val PROFILE_ACCOUNT_KEYS_KEY = "profileAccountKeys"
 
 /**
  * Primary implementation of [AuthDiskSource].
@@ -142,6 +143,7 @@ class AuthDiskSourceImpl(
         storePinProtectedUserKey(userId = userId, pinProtectedUserKey = null)
         storeEncryptedPin(userId = userId, encryptedPin = null)
         storePrivateKey(userId = userId, privateKey = null)
+        storeAccountKeys(userId = userId, accountKeys = null)
         storeOrganizationKeys(userId = userId, organizationKeys = null)
         storeOrganizations(userId = userId, organizations = null)
         storeUserBiometricInitVector(userId = userId, iv = null)
@@ -228,9 +230,11 @@ class AuthDiskSourceImpl(
         )
     }
 
+    @Deprecated("Use getAccountKeys instead.", replaceWith = ReplaceWith("getAccountKeys"))
     override fun getPrivateKey(userId: String): String? =
         getString(key = MASTER_KEY_ENCRYPTION_PRIVATE_KEY.appendIdentifier(userId))
 
+    @Deprecated("Use storeAccountKeys instead.", replaceWith = ReplaceWith("storeAccountKeys"))
     override fun storePrivateKey(userId: String, privateKey: String?) {
         putString(
             key = MASTER_KEY_ENCRYPTION_PRIVATE_KEY.appendIdentifier(userId),
@@ -238,6 +242,20 @@ class AuthDiskSourceImpl(
         )
     }
 
+    override fun getAccountKeys(userId: String): SyncResponseJson.Profile.AccountKeys? =
+        getEncryptedString(key = PROFILE_ACCOUNT_KEYS_KEY.appendIdentifier(userId))
+            ?.let { json.decodeFromStringOrNull(it) }
+
+    override fun storeAccountKeys(
+        userId: String,
+        accountKeys: SyncResponseJson.Profile.AccountKeys?,
+    ) {
+        putEncryptedString(
+            key = PROFILE_ACCOUNT_KEYS_KEY.appendIdentifier(userId),
+            value = accountKeys?.let { json.encodeToString(it) },
+        )
+    }
+
     override fun getUserAutoUnlockKey(userId: String): String? =
         getEncryptedString(
             key = USER_AUTO_UNLOCK_KEY_KEY.appendIdentifier(userId),

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

@@ -12,6 +12,7 @@ import com.bitwarden.network.model.TrustedDeviceUserDecryptionOptionsJson
 import com.bitwarden.network.model.UserDecryptionOptionsJson
 import com.bitwarden.network.model.createMockOrganization
 import com.bitwarden.network.model.createMockPolicy
+import com.bitwarden.network.model.createMockPrivateKeys
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountTokensJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.ForcePasswordResetReason
@@ -280,6 +281,10 @@ class AuthDiskSourceTest {
             userAutoUnlockKey = "userAutoUnlockKey",
         )
         authDiskSource.storePrivateKey(userId = userId, privateKey = "privateKey")
+        authDiskSource.storeAccountKeys(
+            userId = userId,
+            accountKeys = createMockPrivateKeys(number = 1),
+        )
         authDiskSource.storeOrganizationKeys(
             userId = userId,
             organizationKeys = mapOf("organizationId" to "key"),
@@ -330,6 +335,7 @@ class AuthDiskSourceTest {
         assertNull(authDiskSource.getUserKey(userId = userId))
         assertNull(authDiskSource.getUserAutoUnlockKey(userId = userId))
         assertNull(authDiskSource.getPrivateKey(userId = userId))
+        assertNull(authDiskSource.getAccountKeys(userId = userId))
         assertNull(authDiskSource.getOrganizationKeys(userId = userId))
         assertNull(authDiskSource.getOrganizations(userId = userId))
         assertNull(authDiskSource.getPolicies(userId = userId))
@@ -476,6 +482,43 @@ class AuthDiskSourceTest {
         )
     }
 
+    @Test
+    fun `getAccountKeys should pull from SharedPreferences`() {
+        val accountKeysBaseKey = "bwSecureStorage:profileAccountKeys"
+        val mockUserId = "mockUserId"
+        val mockAccountKeys = createMockPrivateKeys(number = 1)
+        fakeEncryptedSharedPreferences.edit {
+            putString(
+                "${accountKeysBaseKey}_$mockUserId",
+                json.encodeToString(mockAccountKeys),
+            )
+        }
+        val actual = authDiskSource.getAccountKeys(userId = mockUserId)
+        assertEquals(
+            mockAccountKeys,
+            actual,
+        )
+    }
+
+    @Test
+    fun `storeAccountKeys should update sharedPreferences`() {
+        val accountKeysBaseKey = "bwSecureStorage:profileAccountKeys"
+        val mockUserId = "mockUserId"
+        val mockAccountKeys = createMockPrivateKeys(number = 1)
+        authDiskSource.storeAccountKeys(
+            userId = mockUserId,
+            accountKeys = mockAccountKeys,
+        )
+        val actual = fakeEncryptedSharedPreferences.getString(
+            "${accountKeysBaseKey}_$mockUserId",
+            null,
+        )
+        assertEquals(
+            json.encodeToJsonElement(mockAccountKeys),
+            json.parseToJsonElement(requireNotNull(actual)),
+        )
+    }
+
     @Test
     fun `getTwoFactorToken should pull from SharedPreferences`() {
         val twoFactorTokenBaseKey = "bwPreferencesStorage:twoFactorToken"

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

@@ -63,6 +63,7 @@ class FakeAuthDiskSource : AuthDiskSource {
     private val storedOnboardingStatus = mutableMapOf<String, OnboardingStatus?>()
     private val storedShowImportLogins = mutableMapOf<String, Boolean?>()
     private val storedLastLockTimestampState = mutableMapOf<String, Instant?>()
+    private val storedAccountKeys = mutableMapOf<String, SyncResponseJson.Profile.AccountKeys?>()
 
     override var userState: UserStateJson? = null
         set(value) {
@@ -137,12 +138,24 @@ class FakeAuthDiskSource : AuthDiskSource {
         storedUserKeys[userId] = userKey
     }
 
+    @Deprecated("Use getAccountKeys instead.", replaceWith = ReplaceWith("getAccountKeys"))
     override fun getPrivateKey(userId: String): String? = storedPrivateKeys[userId]
 
+    @Deprecated("Use storeAccountKeys instead.", replaceWith = ReplaceWith("storeAccountKeys"))
     override fun storePrivateKey(userId: String, privateKey: String?) {
         storedPrivateKeys[userId] = privateKey
     }
 
+    override fun getAccountKeys(userId: String): SyncResponseJson.Profile.AccountKeys? =
+        storedAccountKeys[userId]
+
+    override fun storeAccountKeys(
+        userId: String,
+        accountKeys: SyncResponseJson.Profile.AccountKeys?,
+    ) {
+        storedAccountKeys[userId] = accountKeys
+    }
+
     override fun getTwoFactorToken(email: String): String? = storedTwoFactorTokens[email]
 
     override fun storeTwoFactorToken(email: String, twoFactorToken: String?) {