Allow device key to be persisted to memory (#1209)

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

@@ -119,8 +119,15 @@ interface AuthDiskSource {
 
     /**
      * Stores the device key for the given [userId].
+     *
+     * When [inMemoryOnly] is `true`, the value will only be available via a call to [getDeviceKey]
+     * during the current app session.
      */
-    fun storeDeviceKey(userId: String, deviceKey: String?)
+    fun storeDeviceKey(
+        userId: String,
+        deviceKey: String?,
+        inMemoryOnly: Boolean = false,
+    )
 
     /**
      * Gets the stored [PendingAuthRequestJson] for the given [userId].

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

@@ -56,6 +56,7 @@ class AuthDiskSourceImpl(
 ),
     AuthDiskSource {
 
+    private val inMemoryDeviceKeys = mutableMapOf<String, String?>()
     private val inMemoryPinProtectedUserKeys = mutableMapOf<String, String?>()
     private val mutableOrganizationsFlowMap =
         mutableMapOf<String, MutableSharedFlow<List<SyncResponseJson.Profile.Organization>?>>()
@@ -200,9 +201,15 @@ class AuthDiskSourceImpl(
 
     override fun getDeviceKey(
         userId: String,
-    ): String? = getEncryptedString(key = "${DEVICE_KEY_KEY}_$userId")
+    ): String? = inMemoryDeviceKeys[userId] ?: getEncryptedString(key = "${DEVICE_KEY_KEY}_$userId")
 
-    override fun storeDeviceKey(userId: String, deviceKey: String?) {
+    override fun storeDeviceKey(
+        userId: String,
+        deviceKey: String?,
+        inMemoryOnly: Boolean,
+    ) {
+        inMemoryDeviceKeys[userId] = deviceKey
+        if (inMemoryOnly) return
         putEncryptedString(key = "${DEVICE_KEY_KEY}_$userId", value = deviceKey)
     }
 

app/src/main/java/com/x8bit/bitwarden/data/auth/manager/TrustedDeviceManagerImpl.kt

@@ -2,7 +2,6 @@ package com.x8bit.bitwarden.data.auth.manager
 
 import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
 import com.x8bit.bitwarden.data.auth.datasource.network.service.DevicesService
-import com.x8bit.bitwarden.data.platform.util.asSuccess
 import com.x8bit.bitwarden.data.platform.util.flatMap
 import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
 
@@ -16,7 +15,19 @@ class TrustedDeviceManagerImpl(
 ) : TrustedDeviceManager {
     override suspend fun trustThisDeviceIfNecessary(userId: String): Result<Boolean> =
         if (!authDiskSource.shouldTrustDevice) {
-            false.asSuccess()
+            // Even though we are not trusting the device, we still store the device key in
+            // memory. This allows the user to be "trusted" for this session but on timeout
+            // or reboot, the "trust" will be gone.
+            vaultSdkSource
+                .getTrustDevice(userId = userId)
+                .onSuccess { trustedDevice ->
+                    authDiskSource.storeDeviceKey(
+                        userId = userId,
+                        deviceKey = trustedDevice.deviceKey,
+                        inMemoryOnly = true,
+                    )
+                }
+                .map { false }
         } else {
             vaultSdkSource
                 .getTrustDevice(userId = userId)

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

@@ -38,7 +38,7 @@ class FakeAuthDiskSource : AuthDiskSource {
         mutableMapOf<String, List<SyncResponseJson.Profile.Organization>?>()
     private val storedOrganizationKeys = mutableMapOf<String, Map<String, String>?>()
     private val storedAccountTokens = mutableMapOf<String, AccountTokensJson?>()
-    private val storedDeviceKey = mutableMapOf<String, String?>()
+    private val storedDeviceKey = mutableMapOf<String, Pair<String?, Boolean>>()
     private val storedPendingAuthRequests = mutableMapOf<String, PendingAuthRequestJson?>()
     private val storedBiometricKeys = mutableMapOf<String, String?>()
     private val storedMasterPasswordHashes = mutableMapOf<String, String?>()
@@ -167,10 +167,14 @@ class FakeAuthDiskSource : AuthDiskSource {
         getMutableOrganizationsFlow(userId = userId).tryEmit(organizations)
     }
 
-    override fun getDeviceKey(userId: String): String? = storedDeviceKey[userId]
+    override fun getDeviceKey(userId: String): String? = storedDeviceKey[userId]?.first
 
-    override fun storeDeviceKey(userId: String, deviceKey: String?) {
-        storedDeviceKey[userId] = deviceKey
+    override fun storeDeviceKey(
+        userId: String,
+        deviceKey: String?,
+        inMemoryOnly: Boolean,
+    ) {
+        storedDeviceKey[userId] = deviceKey to inMemoryOnly
     }
 
     override fun getPendingAuthRequest(userId: String): PendingAuthRequestJson? =
@@ -298,8 +302,8 @@ class FakeAuthDiskSource : AuthDiskSource {
     /**
      * Assert that the [deviceKey] was stored successfully using the [userId].
      */
-    fun assertDeviceKey(userId: String, deviceKey: String?) {
-        assertEquals(deviceKey, storedDeviceKey[userId])
+    fun assertDeviceKey(userId: String, deviceKey: String?, inMemoryOnly: Boolean = false) {
+        assertEquals(deviceKey to inMemoryOnly, storedDeviceKey[userId])
     }
 
     /**

app/src/test/java/com/x8bit/bitwarden/data/auth/manager/TrustedDeviceManagerTests.kt

@@ -32,13 +32,30 @@ class TrustedDeviceManagerTests {
     fun `trustThisDeviceIfNecessary when shouldTrustDevice false should return success with false`() =
         runTest {
             val userId = "userId"
+            val deviceKey = "deviceKey"
+            val trustedDeviceResponse = TrustDeviceResponse(
+                deviceKey = deviceKey,
+                protectedUserKey = "protectedUserKey",
+                protectedDevicePrivateKey = "protectedDevicePrivateKey",
+                protectedDevicePublicKey = "protectedDevicePublicKey",
+            )
             fakeAuthDiskSource.shouldTrustDevice = false
+            coEvery {
+                vaultSdkSource.getTrustDevice(userId = userId)
+            } returns trustedDeviceResponse.asSuccess()
 
             val result = manager.trustThisDeviceIfNecessary(userId = userId)
 
             assertEquals(false.asSuccess(), result)
-            coVerify(exactly = 0) {
+            fakeAuthDiskSource.assertDeviceKey(
+                userId = userId,
+                deviceKey = deviceKey,
+                inMemoryOnly = true,
+            )
+            coVerify(exactly = 1) {
                 vaultSdkSource.getTrustDevice(userId = userId)
+            }
+            coVerify(exactly = 0) {
                 devicesService.trustDevice(
                     appId = any(),
                     encryptedUserKey = any(),