BIT-1091: Initialize crypto for organizations (#370)

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

@@ -179,9 +179,10 @@ class AuthRepositoryImpl constructor(
                             kdf = userStateJson.activeAccount.profile.toSdkParams(),
                             userKey = loginResponse.key,
                             privateKey = loginResponse.privateKey,
-                            // TODO use actual organization keys BIT-1091
-                            organizationalKeys = emptyMap(),
                             masterPassword = password,
+                            // We can separately unlock the vault for organization data after
+                            // receiving the sync response.
+                            organizationKeys = null,
                         )
                         authDiskSource.userState = userStateJson
                         authDiskSource.storeUserKey(

app/src/main/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSource.kt

@@ -7,6 +7,7 @@ import com.bitwarden.core.Collection
 import com.bitwarden.core.CollectionView
 import com.bitwarden.core.Folder
 import com.bitwarden.core.FolderView
+import com.bitwarden.core.InitOrgCryptoRequest
 import com.bitwarden.core.InitUserCryptoRequest
 import com.bitwarden.core.Send
 import com.bitwarden.core.SendView
@@ -19,11 +20,21 @@ import com.x8bit.bitwarden.data.vault.datasource.sdk.model.InitializeCryptoResul
 interface VaultSdkSource {
 
     /**
-     * Attempts to initialize cryptography functionality for the Bitwarden SDK
-     * with a given [InitCryptoRequest].
+     * Attempts to initialize cryptography functionality for an individual user for the
+     * Bitwarden SDK with a given [InitUserCryptoRequest].
      */
     suspend fun initializeCrypto(request: InitUserCryptoRequest): Result<InitializeCryptoResult>
 
+    /**
+     * Attempts to initialize cryptography functionality for organization data associated with
+     * the current user for the Bitwarden SDK with a given [InitOrgCryptoRequest].
+     *
+     * This should only be called after a successful call to [initializeCrypto].
+     */
+    suspend fun initializeOrganizationCrypto(
+        request: InitOrgCryptoRequest,
+    ): Result<InitializeCryptoResult>
+
     /**
      * Encrypts a [CipherView] returning a [Cipher] wrapped in a [Result].
      */

app/src/main/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSourceImpl.kt

@@ -7,6 +7,7 @@ import com.bitwarden.core.Collection
 import com.bitwarden.core.CollectionView
 import com.bitwarden.core.Folder
 import com.bitwarden.core.FolderView
+import com.bitwarden.core.InitOrgCryptoRequest
 import com.bitwarden.core.InitUserCryptoRequest
 import com.bitwarden.core.Send
 import com.bitwarden.core.SendView
@@ -32,7 +33,20 @@ class VaultSdkSourceImpl(
                 clientCrypto.initializeUserCrypto(req = request)
                 InitializeCryptoResult.Success
             } catch (exception: BitwardenException) {
-                // The only truly expected error from the SDK is an incorrect password.
+                // The only truly expected error from the SDK is an incorrect key/password.
+                InitializeCryptoResult.AuthenticationError
+            }
+        }
+
+    override suspend fun initializeOrganizationCrypto(
+        request: InitOrgCryptoRequest,
+    ): Result<InitializeCryptoResult> =
+        runCatching {
+            try {
+                clientCrypto.initializeOrgCrypto(req = request)
+                InitializeCryptoResult.Success
+            } catch (exception: BitwardenException) {
+                // The only truly expected error from the SDK is for incorrect keys.
                 InitializeCryptoResult.AuthenticationError
             }
         }

app/src/main/java/com/x8bit/bitwarden/data/vault/datasource/sdk/model/InitializeCryptoResult.kt

@@ -11,7 +11,7 @@ sealed class InitializeCryptoResult {
     data object Success : InitializeCryptoResult()
 
     /**
-     * Incorrect password provided.
+     * Incorrect password or key(s) provided.
      */
     data object AuthenticationError : InitializeCryptoResult()
 }

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

@@ -71,6 +71,9 @@ interface VaultRepository {
 
     /**
      * Attempt to unlock the vault with the specified user information.
+     *
+     * Note that when [organizationKeys] is absent, no attempt will be made to unlock the vault
+     * for organization data.
      */
     @Suppress("LongParameterList")
     suspend fun unlockVault(
@@ -80,7 +83,7 @@ interface VaultRepository {
         kdf: Kdf,
         userKey: String,
         privateKey: String,
-        organizationalKeys: Map<String, String>,
+        organizationKeys: Map<String, String>?,
     ): VaultUnlockResult
 
     /**

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

@@ -2,6 +2,7 @@ package com.x8bit.bitwarden.data.vault.repository
 
 import com.bitwarden.core.CipherView
 import com.bitwarden.core.FolderView
+import com.bitwarden.core.InitOrgCryptoRequest
 import com.bitwarden.core.InitUserCryptoMethod
 import com.bitwarden.core.InitUserCryptoRequest
 import com.bitwarden.core.Kdf
@@ -12,11 +13,13 @@ import com.x8bit.bitwarden.data.platform.datasource.network.util.isNoConnectionE
 import com.x8bit.bitwarden.data.platform.manager.dispatcher.DispatcherManager
 import com.x8bit.bitwarden.data.platform.repository.model.DataState
 import com.x8bit.bitwarden.data.platform.repository.util.map
+import com.x8bit.bitwarden.data.platform.util.asSuccess
 import com.x8bit.bitwarden.data.platform.util.flatMap
 import com.x8bit.bitwarden.data.vault.datasource.network.model.SyncResponseJson
 import com.x8bit.bitwarden.data.vault.datasource.network.service.CiphersService
 import com.x8bit.bitwarden.data.vault.datasource.network.service.SyncService
 import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
+import com.x8bit.bitwarden.data.vault.datasource.sdk.model.InitializeCryptoResult
 import com.x8bit.bitwarden.data.vault.repository.model.CreateCipherResult
 import com.x8bit.bitwarden.data.vault.repository.model.SendData
 import com.x8bit.bitwarden.data.vault.repository.model.UpdateCipherResult
@@ -107,6 +110,7 @@ class VaultRepositoryImpl constructor(
                                 syncResponse = syncResponse,
                             )
 
+                        unlockVaultForOrganizationsIfNecessary(syncResponse = syncResponse)
                         storeKeys(syncResponse = syncResponse)
                         decryptSyncResponseAndUpdateVaultDataState(syncResponse = syncResponse)
                         decryptSendsAndUpdateSendDataState(sendList = syncResponse.sends)
@@ -177,6 +181,8 @@ class VaultRepositoryImpl constructor(
             ?: return VaultUnlockResult.InvalidStateError
         val privateKey = authDiskSource.getPrivateKey(userId = userState.activeUserId)
             ?: return VaultUnlockResult.InvalidStateError
+        val organizationKeys = authDiskSource
+            .getOrganizationKeys(userId = userState.activeUserId)
         return unlockVault(
             userId = userState.activeUserId,
             masterPassword = masterPassword,
@@ -184,8 +190,7 @@ class VaultRepositoryImpl constructor(
             kdf = userState.activeAccount.profile.toSdkParams(),
             userKey = userKey,
             privateKey = privateKey,
-            // TODO use actual organization keys BIT-1091
-            organizationalKeys = emptyMap(),
+            organizationKeys = organizationKeys,
         )
             .also {
                 if (it is VaultUnlockResult.Success) {
@@ -201,7 +206,7 @@ class VaultRepositoryImpl constructor(
         kdf: Kdf,
         userKey: String,
         privateKey: String,
-        organizationalKeys: Map<String, String>,
+        organizationKeys: Map<String, String>?,
     ): VaultUnlockResult =
         flow {
             willSyncAfterUnlock = true
@@ -218,6 +223,20 @@ class VaultRepositoryImpl constructor(
                             ),
                         ),
                     )
+                    .flatMap { result ->
+                        // Initialize the SDK for organizations if necessary
+                        if (organizationKeys != null &&
+                            result is InitializeCryptoResult.Success
+                        ) {
+                            vaultSdkSource.initializeOrganizationCrypto(
+                                request = InitOrgCryptoRequest(
+                                    organizationKeys = organizationKeys,
+                                ),
+                            )
+                        } else {
+                            result.asSuccess()
+                        }
+                    }
                     .fold(
                         onFailure = { VaultUnlockResult.GenericError },
                         onSuccess = { initializeCryptoResult ->
@@ -320,6 +339,27 @@ class VaultRepositoryImpl constructor(
         }
     }
 
+    private suspend fun unlockVaultForOrganizationsIfNecessary(
+        syncResponse: SyncResponseJson,
+    ) {
+        val profile = syncResponse.profile ?: return
+        val organizationKeys = profile.organizations
+            .orEmpty()
+            .filter { it.key != null }
+            .associate { it.id to requireNotNull(it.key) }
+            .takeUnless { it.isEmpty() }
+            ?: return
+
+        // There shouldn't be issues when unlocking directly from the syncResponse so we can ignore
+        // the return type here.
+        vaultSdkSource
+            .initializeOrganizationCrypto(
+                request = InitOrgCryptoRequest(
+                    organizationKeys = organizationKeys,
+                ),
+            )
+    }
+
     private suspend fun decryptSendsAndUpdateSendDataState(sendList: List<SyncResponseJson.Send>?) {
         val newState = vaultSdkSource
             .decryptSendList(