BIT-2409: Update the attachment migration process (#1454)

This commit is contained in:
David Perez 2024-06-17 11:28:03 -05:00 committed by Álison Fernandes
parent e860560df4
commit dc633f0c0a
2 changed files with 108 additions and 81 deletions

View file

@ -1,6 +1,7 @@
package com.x8bit.bitwarden.data.vault.manager
import android.net.Uri
import androidx.core.net.toUri
import com.bitwarden.vault.AttachmentView
import com.bitwarden.vault.Cipher
import com.bitwarden.vault.CipherView
@ -28,9 +29,6 @@ import com.x8bit.bitwarden.data.vault.repository.util.toEncryptedNetworkCipher
import com.x8bit.bitwarden.data.vault.repository.util.toEncryptedNetworkCipherResponse
import com.x8bit.bitwarden.data.vault.repository.util.toEncryptedSdkCipher
import com.x8bit.bitwarden.data.vault.repository.util.toNetworkAttachmentRequest
import kotlinx.coroutines.async
import kotlinx.coroutines.awaitAll
import kotlinx.coroutines.coroutineScope
import java.io.File
import java.time.Clock
@ -240,19 +238,15 @@ class CipherManagerImpl(
collectionIds: List<String>,
): ShareCipherResult {
val userId = activeUserId ?: return ShareCipherResult.Error
return vaultSdkSource
.moveToOrganization(
userId = userId,
organizationId = organizationId,
cipherView = cipherView,
)
return migrateAttachments(userId = userId, cipherView = cipherView)
.flatMap {
migrateAttachments(
vaultSdkSource.moveToOrganization(
userId = userId,
cipherView = it,
organizationId = organizationId,
cipherView = it,
)
}
.flatMap { vaultSdkSource.encryptCipher(userId = userId, cipherView = it) }
.flatMap { cipher ->
ciphersService.shareCipher(
cipherId = cipherId,
@ -492,71 +486,58 @@ class CipherManagerImpl(
private suspend fun migrateAttachments(
userId: String,
cipherView: CipherView,
organizationId: String,
): Result<Cipher> {
): Result<CipherView> {
// Only run the migrations if we have attachments that do not have their own 'key'
val attachmentViewsToMigrate = cipherView.attachments.orEmpty().filter { it.key == null }
if (attachmentViewsToMigrate.none()) {
return vaultSdkSource.encryptCipher(userId = userId, cipherView = cipherView)
}
if (attachmentViewsToMigrate.none()) return cipherView.asSuccess()
val cipherViewId = cipherView.id
?: return IllegalStateException("CipherView must have an ID").asFailure()
val cipher = vaultSdkSource
.encryptCipher(userId = userId, cipherView = cipherView)
var migratedCipherView = cipherView
.encryptCipherAndCheckForMigration(userId = userId, cipherId = cipherViewId)
.flatMap { vaultSdkSource.decryptCipher(userId = userId, cipher = it) }
.getOrElse { return it.asFailure() }
// Gets a list of all the attachments that do not require migration
// We will combine this with all migrated attachments at the end
val attachmentsWithKeys = cipher.attachments.orEmpty().filter { it.key != null }
val migrations = coroutineScope {
attachmentViewsToMigrate.map { attachmentView ->
async {
attachmentView
.id
?.let { attachmentId ->
this@CipherManagerImpl
.downloadAttachmentForResult(
attachmentViewsToMigrate
.map { attachmentView ->
attachmentView
.id
?.let { attachmentId ->
// This process downloads the attachment file and creates an entirely
// new attachment before deleting the original one
this@CipherManagerImpl
.downloadAttachmentForResult(
cipherView = migratedCipherView,
attachmentId = attachmentId,
)
.flatMap { file ->
this@CipherManagerImpl
.createAttachmentForResult(
cipherId = cipherViewId,
cipherView = migratedCipherView,
fileSizeBytes = attachmentView.size,
fileName = attachmentView.fileName,
fileUri = file.toUri(),
)
.onSuccess { fileManager.delete(file) }
}
.flatMap { cipherView ->
deleteCipherAttachmentForResult(
cipherView = cipherView,
attachmentId = attachmentId,
cipherId = cipherViewId,
)
.flatMap { decryptedFile ->
val encryptedFile = File("${decryptedFile.absolutePath}.enc")
// Re-encrypting the attachment will generate the `key` and
// we need to encrypt the associated file with that `key`
vaultSdkSource
.encryptAttachment(
userId = userId,
cipher = cipher,
attachmentView = attachmentView,
decryptedFilePath = decryptedFile.absolutePath,
encryptedFilePath = encryptedFile.absolutePath,
)
.onSuccess { fileManager.delete(decryptedFile) }
.flatMap { attachment ->
ciphersService
.shareAttachment(
cipherId = cipherViewId,
attachment = attachment,
organizationId = organizationId,
encryptedFile = encryptedFile,
)
.onSuccess { fileManager.delete(encryptedFile) }
.map { attachment }
}
}
}
?: IllegalStateException("AttachmentView must have an ID").asFailure()
}
}
.flatMap { vaultSdkSource.decryptCipher(userId = userId, cipher = it) }
.onSuccess { migratedCipherView = it }
}
?: IllegalStateException("AttachmentView must have an ID").asFailure()
}
}
// We are collecting the migrated attachments to combine with the un-migrated attachments
// If anything fails, we consider the entire process to be a failure
val migratedAttachments = awaitAll(*migrations.toTypedArray()).map {
it.getOrElse { error -> return error.asFailure() }
}
return cipher.copy(attachments = attachmentsWithKeys + migratedAttachments).asSuccess()
.onEach { result ->
// If anything fails, we consider the entire process to be a failure
// The attachments will be partially migrated and that is OK
result.onFailure { return it.asFailure() }
}
return migratedCipherView.asSuccess()
}
}

View file

@ -1,6 +1,7 @@
package com.x8bit.bitwarden.data.vault.manager
import android.net.Uri
import androidx.core.net.toUri
import com.bitwarden.vault.Attachment
import com.bitwarden.vault.AttachmentView
import com.bitwarden.vault.Cipher
@ -39,6 +40,7 @@ import com.x8bit.bitwarden.data.vault.repository.model.UpdateCipherResult
import com.x8bit.bitwarden.data.vault.repository.util.toEncryptedNetworkCipher
import com.x8bit.bitwarden.data.vault.repository.util.toEncryptedNetworkCipherResponse
import com.x8bit.bitwarden.data.vault.repository.util.toEncryptedSdkCipher
import com.x8bit.bitwarden.data.vault.repository.util.toNetworkAttachmentRequest
import io.mockk.coEvery
import io.mockk.coVerify
import io.mockk.every
@ -787,7 +789,7 @@ class CipherManagerTest {
runTest {
fakeAuthDiskSource.userState = MOCK_USER_STATE
val userId = "mockId-1"
val organizationId = "organizationId"
val organizationId = "mockOrganizationId"
val mockAttachmentView = createMockAttachmentView(number = 1, key = null)
val initialCipherView = createMockCipherView(number = 1).copy(
attachments = listOf(mockAttachmentView),
@ -800,19 +802,16 @@ class CipherManagerTest {
val encryptedFile = File("path/to/encrypted/file")
val decryptedFile = File("path/to/encrypted/file_decrypted")
val mockCipherView = createMockCipherView(number = 1)
coEvery {
vaultSdkSource.moveToOrganization(
userId = userId,
organizationId = organizationId,
cipherView = initialCipherView,
)
} returns mockCipherView.asSuccess()
val mockAttachmentJsonResponse = createMockAttachmentJsonResponse(number = 1)
val mockNetworkCipher = createMockCipher(number = 1)
// Handle mocks for migration
coEvery {
vaultSdkSource.encryptCipher(userId = userId, cipherView = initialCipherView)
} returns mockCipher.asSuccess()
coEvery {
vaultSdkSource.decryptCipher(userId = userId, cipher = mockCipher)
} returns initialCipherView.asSuccess()
coEvery {
ciphersService.getCipherAttachment(cipherId = "mockId-1", attachmentId = "mockId-1")
} returns attachment.asSuccess()
@ -828,6 +827,10 @@ class CipherManagerTest {
decryptedFilePath = decryptedFile.path,
)
} returns Unit.asSuccess()
val cacheFile = File("path/to/cache/file")
coEvery {
fileManager.writeUriToCache(fileUri = decryptedFile.toUri())
} returns cacheFile.asSuccess()
coEvery {
vaultSdkSource.encryptAttachment(
userId = userId,
@ -840,20 +843,63 @@ class CipherManagerTest {
fileName = "mockFileName-1",
key = null,
),
decryptedFilePath = any(),
encryptedFilePath = any(),
decryptedFilePath = cacheFile.absolutePath,
encryptedFilePath = "${cacheFile.absolutePath}.enc",
)
} returns mockAttachment.asSuccess()
coEvery {
ciphersService.shareAttachment(
ciphersService.createAttachment(
cipherId = "mockId-1",
attachment = mockAttachment,
organizationId = organizationId,
encryptedFile = any(),
body = mockAttachment.toNetworkAttachmentRequest(),
)
} returns mockAttachmentJsonResponse.asSuccess()
coEvery {
ciphersService.uploadAttachment(
attachmentJsonResponse = mockAttachmentJsonResponse,
encryptedFile = File("${cacheFile.absolutePath}.enc"),
)
} returns mockNetworkCipher.asSuccess()
coEvery {
vaultDiskSource.saveCipher(
userId = userId,
cipher = mockNetworkCipher.copy(collectionIds = listOf("mockId-1")),
)
} just runs
coEvery {
vaultSdkSource.decryptCipher(
userId = userId,
cipher = mockNetworkCipher
.copy(collectionIds = listOf("mockId-1"))
.toEncryptedSdkCipher(),
)
} returns mockCipherView.asSuccess()
coEvery {
ciphersService.deleteCipherAttachment(
cipherId = "mockId-1",
attachmentId = "mockId-1",
)
} returns Unit.asSuccess()
coEvery {
vaultSdkSource.encryptCipher(
userId = userId,
cipherView = mockCipherView.copy(attachments = emptyList()),
)
} returns mockCipher.asSuccess()
coEvery {
vaultDiskSource.saveCipher(
userId = userId,
cipher = mockCipher.toEncryptedNetworkCipherResponse(),
)
} just runs
// Done with mocks for migration
coEvery {
vaultSdkSource.moveToOrganization(
userId = userId,
organizationId = organizationId,
cipherView = initialCipherView,
)
} returns mockCipherView.asSuccess()
coEvery {
vaultSdkSource.encryptCipher(userId = userId, cipherView = mockCipherView)
} returns createMockSdkCipher(number = 1, clock = clock).asSuccess()