[PM-30899] Store account keys on new user creation (#6403)

Co-authored-by: Patrick Honkonen <phonkonen@bitwarden.com>

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryImpl.kt

@@ -2,6 +2,7 @@ package com.x8bit.bitwarden.data.auth.repository
 
 import com.bitwarden.core.AuthRequestMethod
 import com.bitwarden.core.InitUserCryptoMethod
+import com.bitwarden.core.RegisterTdeKeyResponse
 import com.bitwarden.core.WrappedAccountCryptographicState
 import com.bitwarden.core.data.manager.dispatcher.DispatcherManager
 import com.bitwarden.core.data.repository.error.MissingPropertyException
@@ -14,6 +15,7 @@ import com.bitwarden.crypto.Kdf
 import com.bitwarden.data.datasource.disk.ConfigDiskSource
 import com.bitwarden.data.repository.util.toEnvironmentUrls
 import com.bitwarden.data.repository.util.toEnvironmentUrlsOrDefault
+import com.bitwarden.network.model.CreateAccountKeysResponseJson
 import com.bitwarden.network.model.DeleteAccountResponseJson
 import com.bitwarden.network.model.GetTokenResponseJson
 import com.bitwarden.network.model.IdentityTokenAuthModel
@@ -459,42 +461,32 @@ class AuthRepositoryImpl(
                                 .getShouldTrustDevice(userId = userId) == true,
                         )
                     }
-                    .flatMap { keys ->
+                    .flatMap { registerTdeKeyResponse ->
                         accountsService
                             .createAccountKeys(
-                                publicKey = keys.publicKey,
-                                encryptedPrivateKey = keys.privateKey,
+                                publicKey = registerTdeKeyResponse.publicKey,
+                                encryptedPrivateKey = registerTdeKeyResponse.privateKey,
                             )
-                            .map { keys }
+                            .map { createAccountKeysResponse ->
+                                registerTdeKeyResponse to createAccountKeysResponse
+                            }
                     }
-                    .flatMap { keys ->
+                    .flatMap { (registerTdeKeyResponse, createAccountKeysResponse) ->
                         organizationService
                             .organizationResetPasswordEnroll(
                                 organizationId = orgAutoEnrollStatus.organizationId,
                                 userId = userId,
                                 passwordHash = null,
-                                resetPasswordKey = keys.adminReset,
+                                resetPasswordKey = registerTdeKeyResponse.adminReset,
                             )
-                            .map { keys }
+                            .map { registerTdeKeyResponse to createAccountKeysResponse }
                     }
-                    .onSuccess { keys ->
-                        // TDE and SSO user creation still uses crypto-v1. These users are not
-                        // expected to have the AEAD keys so we only store the private key for now.
-                        // See https://github.com/bitwarden/android/pull/5682#discussion_r2273940332
-                        // for more details.
-                        authDiskSource.storePrivateKey(
+                    .onSuccess { (registerTdeKeyResponse, createAccountKeysResponse) ->
+                        createNewSsoUserSuccess(
                             userId = userId,
-                            privateKey = keys.privateKey,
+                            createAccountKeysResponse = createAccountKeysResponse,
+                            registerTdeKeyResponse = registerTdeKeyResponse,
                         )
-                        // Order matters here, we need to make sure that the vault is unlocked
-                        // before we trust the device, to avoid state-base navigation issues.
-                        vaultRepository.syncVaultState(userId = userId)
-                        keys.deviceKey?.let { trustDeviceResponse ->
-                            trustedDeviceManager.trustThisDevice(
-                                userId = userId,
-                                trustDeviceResponse = trustDeviceResponse,
-                            )
-                        }
                     }
             }
             .fold(
@@ -503,6 +495,37 @@ class AuthRepositoryImpl(
             )
     }
 
+    /**
+     * Stores all the relevant data from a successful creation of an SSO user. The data is stored
+     * while in an [UserStateManager.userStateTransaction] to ensure the `UserState` is only
+     * updated once after data stored.
+     */
+    private suspend fun createNewSsoUserSuccess(
+        userId: String,
+        createAccountKeysResponse: CreateAccountKeysResponseJson,
+        registerTdeKeyResponse: RegisterTdeKeyResponse,
+    ): Unit = userStateManager.userStateTransaction {
+        authDiskSource.storeAccountKeys(
+            userId = userId,
+            accountKeys = createAccountKeysResponse.accountKeys,
+        )
+        // TDE and SSO user creation still uses crypto-v1. These users are not
+        // expected to have the AEAD keys so we only store the private key for now.
+        // See https://github.com/bitwarden/android/pull/5682#discussion_r2273940332
+        // for more details.
+        authDiskSource.storePrivateKey(
+            userId = userId,
+            privateKey = registerTdeKeyResponse.privateKey,
+        )
+        vaultRepository.syncVaultState(userId = userId)
+        registerTdeKeyResponse.deviceKey?.let { trustDeviceResponse ->
+            trustedDeviceManager.trustThisDevice(
+                userId = userId,
+                trustDeviceResponse = trustDeviceResponse,
+            )
+        }
+    }
+
     override suspend fun completeTdeLogin(
         requestPrivateKey: String,
         asymmetricalKey: String,

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryTest.kt

@@ -26,6 +26,7 @@ import com.bitwarden.data.datasource.disk.model.ServerConfig
 import com.bitwarden.data.datasource.disk.util.FakeConfigDiskSource
 import com.bitwarden.data.repository.model.Environment
 import com.bitwarden.network.model.ConfigResponseJson
+import com.bitwarden.network.model.CreateAccountKeysResponseJson
 import com.bitwarden.network.model.DeleteAccountResponseJson
 import com.bitwarden.network.model.GetTokenResponseJson
 import com.bitwarden.network.model.IdentityTokenAuthModel
@@ -1137,7 +1138,12 @@ class AuthRepositoryTest {
                 publicKey = userPublicKey,
                 encryptedPrivateKey = userPrivateKey,
             )
-        } returns Unit.asSuccess()
+        } returns CreateAccountKeysResponseJson(
+            key = null,
+            publicKey = userPublicKey,
+            privateKey = userPrivateKey,
+            accountKeys = null,
+        ).asSuccess()
         coEvery {
             organizationService.organizationResetPasswordEnroll(
                 organizationId = orgId,
@@ -1222,7 +1228,12 @@ class AuthRepositoryTest {
                     publicKey = userPublicKey,
                     encryptedPrivateKey = userPrivateKey,
                 )
-            } returns Unit.asSuccess()
+            } returns CreateAccountKeysResponseJson(
+                key = null,
+                publicKey = userPublicKey,
+                privateKey = userPrivateKey,
+                accountKeys = ACCOUNT_KEYS,
+            ).asSuccess()
             coEvery {
                 organizationService.organizationResetPasswordEnroll(
                     organizationId = orgId,
@@ -1236,6 +1247,7 @@ class AuthRepositoryTest {
             val result = repository.createNewSsoUser()
 
             fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = userPrivateKey)
+            fakeAuthDiskSource.assertAccountKeys(userId = USER_ID_1, accountKeys = ACCOUNT_KEYS)
             assertEquals(NewSsoUserResult.Success, result)
             coVerify(exactly = 1) {
                 organizationService.getOrganizationAutoEnrollStatus(orgIdentifier)
@@ -1310,7 +1322,12 @@ class AuthRepositoryTest {
                     publicKey = userPublicKey,
                     encryptedPrivateKey = userPrivateKey,
                 )
-            } returns Unit.asSuccess()
+            } returns CreateAccountKeysResponseJson(
+                key = null,
+                publicKey = userPublicKey,
+                privateKey = userPrivateKey,
+                accountKeys = ACCOUNT_KEYS,
+            ).asSuccess()
             coEvery {
                 organizationService.organizationResetPasswordEnroll(
                     organizationId = orgId,
@@ -1330,6 +1347,7 @@ class AuthRepositoryTest {
             val result = repository.createNewSsoUser()
 
             fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = userPrivateKey)
+            fakeAuthDiskSource.assertAccountKeys(userId = USER_ID_1, accountKeys = ACCOUNT_KEYS)
             assertEquals(NewSsoUserResult.Success, result)
             coVerify(exactly = 1) {
                 organizationService.getOrganizationAutoEnrollStatus(orgIdentifier)

network/src/main/kotlin/com/bitwarden/network/api/AuthenticatedAccountsApi.kt

@@ -1,6 +1,7 @@
 package com.bitwarden.network.api
 
 import com.bitwarden.network.model.CreateAccountKeysRequest
+import com.bitwarden.network.model.CreateAccountKeysResponseJson
 import com.bitwarden.network.model.DeleteAccountRequestJson
 import com.bitwarden.network.model.NetworkResult
 import com.bitwarden.network.model.ResetPasswordRequestJson
@@ -26,7 +27,9 @@ internal interface AuthenticatedAccountsApi {
      * Creates the keys for the current account.
      */
     @POST("/accounts/keys")
-    suspend fun createAccountKeys(@Body body: CreateAccountKeysRequest): NetworkResult<Unit>
+    suspend fun createAccountKeys(
+        @Body body: CreateAccountKeysRequest,
+    ): NetworkResult<CreateAccountKeysResponseJson>
 
     /**
      * Deletes the current account.

network/src/main/kotlin/com/bitwarden/network/model/CreateAccountKeysResponseJson.kt

@@ -0,0 +1,27 @@
+package com.bitwarden.network.model
+
+import kotlinx.serialization.SerialName
+import kotlinx.serialization.Serializable
+
+/**
+ * Response object returned when creating account keys.
+ *
+ * @property key The user key (nullable).
+ * @property publicKey The public key for the account.
+ * @property privateKey The encrypted private key for the account.
+ * @property accountKeys The account keys containing encryption key pairs and security state.
+ */
+@Serializable
+data class CreateAccountKeysResponseJson(
+    @SerialName("key")
+    val key: String?,
+
+    @SerialName("publicKey")
+    val publicKey: String?,
+
+    @SerialName("privateKey")
+    val privateKey: String?,
+
+    @SerialName("accountKeys")
+    val accountKeys: AccountKeysJson?,
+)

network/src/main/kotlin/com/bitwarden/network/service/AccountsService.kt

@@ -1,5 +1,6 @@
 package com.bitwarden.network.service
 
+import com.bitwarden.network.model.CreateAccountKeysResponseJson
 import com.bitwarden.network.model.DeleteAccountResponseJson
 import com.bitwarden.network.model.KeyConnectorKeyRequestJson
 import com.bitwarden.network.model.KeyConnectorMasterKeyResponseJson
@@ -26,7 +27,10 @@ interface AccountsService {
     /**
      * Creates a new account's keys.
      */
-    suspend fun createAccountKeys(publicKey: String, encryptedPrivateKey: String): Result<Unit>
+    suspend fun createAccountKeys(
+        publicKey: String,
+        encryptedPrivateKey: String,
+    ): Result<CreateAccountKeysResponseJson>
 
     /**
      * Make delete account request.

network/src/main/kotlin/com/bitwarden/network/service/AccountsServiceImpl.kt

@@ -5,6 +5,7 @@ import com.bitwarden.network.api.AuthenticatedKeyConnectorApi
 import com.bitwarden.network.api.UnauthenticatedAccountsApi
 import com.bitwarden.network.api.UnauthenticatedKeyConnectorApi
 import com.bitwarden.network.model.CreateAccountKeysRequest
+import com.bitwarden.network.model.CreateAccountKeysResponseJson
 import com.bitwarden.network.model.DeleteAccountRequestJson
 import com.bitwarden.network.model.DeleteAccountResponseJson
 import com.bitwarden.network.model.KeyConnectorKeyRequestJson
@@ -50,7 +51,7 @@ internal class AccountsServiceImpl(
     override suspend fun createAccountKeys(
         publicKey: String,
         encryptedPrivateKey: String,
-    ): Result<Unit> =
+    ): Result<CreateAccountKeysResponseJson> =
         authenticatedAccountsApi
             .createAccountKeys(
                 body = CreateAccountKeysRequest(

network/src/test/kotlin/com/bitwarden/network/service/AccountsServiceTest.kt

@@ -53,11 +53,10 @@ class AccountsServiceTest : BaseServiceTest() {
     }
 
     @Test
-    fun `createAccountKeys with empty response is success`() = runTest {
+    fun `createAccountKeys success response should return Success`() = runTest {
         val publicKey = "publicKey"
         val encryptedPrivateKey = "encryptedPrivateKey"
-        val json = ""
-        val response = MockResponse().setBody(json)
+        val response = MockResponse().setBody(CREATE_ACCOUNT_KEYS_REQUEST_RESPONSE)
         server.enqueue(response)
 
         val result = service.createAccountKeys(
@@ -368,3 +367,22 @@ private val UPDATE_KDF_REQUEST = UpdateKdfJsonRequest(
         salt = "mockSalt",
     ),
 )
+private val CREATE_ACCOUNT_KEYS_REQUEST_RESPONSE = """
+{
+  "key": null,
+  "publicKey": "mockPublicKey-1",
+  "privateKey": "mockPrivateKey-1",
+  "accountKeys": {
+    "signatureKeyPair": null,
+    "publicKeyEncryptionKeyPair": {
+      "wrappedPrivateKey": "mockWrappedPrivateKey-1",
+      "publicKey": "mockPublicKey-1",
+      "signedPublicKey": null,
+      "object": "publicKeyEncryptionKeyPair"
+    },
+    "securityState": null,
+    "object": "privateKeys"
+  },
+  "object": "keys"
+}
+""".trimIndent()