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)