BIT-2409: Update attachment migration logic when sharing a cipher (#1447)

app/src/main/java/com/x8bit/bitwarden/data/vault/datasource/network/api/CiphersApi.kt

@@ -14,6 +14,7 @@ import retrofit2.http.GET
 import retrofit2.http.POST
 import retrofit2.http.PUT
 import retrofit2.http.Path
+import retrofit2.http.Query
 
 /**
  * Defines raw calls under the /ciphers API with authentication applied.
@@ -72,6 +73,17 @@ interface CiphersApi {
         @Body body: ShareCipherJsonRequest,
     ): Result<SyncResponseJson.Cipher>
 
+    /**
+     * Shares an attachment.
+     */
+    @POST("ciphers/{cipherId}/attachment/{attachmentId}/share")
+    suspend fun shareAttachment(
+        @Path("cipherId") cipherId: String,
+        @Path("attachmentId") attachmentId: String,
+        @Query("organizationId") organizationId: String?,
+        @Body body: MultipartBody,
+    ): Result<Unit>
+
     /**
      * Updates a cipher's collections.
      */

app/src/main/java/com/x8bit/bitwarden/data/vault/datasource/network/service/CiphersService.kt

@@ -1,5 +1,6 @@
 package com.x8bit.bitwarden.data.vault.datasource.network.service
 
+import com.bitwarden.vault.Attachment
 import com.x8bit.bitwarden.data.vault.datasource.network.model.AttachmentJsonRequest
 import com.x8bit.bitwarden.data.vault.datasource.network.model.AttachmentJsonResponse
 import com.x8bit.bitwarden.data.vault.datasource.network.model.CipherJsonRequest
@@ -59,6 +60,16 @@ interface CiphersService {
         body: ShareCipherJsonRequest,
     ): Result<SyncResponseJson.Cipher>
 
+    /**
+     * Attempt to share an attachment.
+     */
+    suspend fun shareAttachment(
+        cipherId: String,
+        attachment: Attachment,
+        organizationId: String,
+        encryptedFile: File,
+    ): Result<Unit>
+
     /**
      * Attempt to update a cipher's collections.
      */

app/src/main/java/com/x8bit/bitwarden/data/vault/datasource/network/service/CiphersServiceImpl.kt

@@ -1,8 +1,10 @@
 package com.x8bit.bitwarden.data.vault.datasource.network.service
 
 import androidx.core.net.toUri
+import com.bitwarden.vault.Attachment
 import com.x8bit.bitwarden.data.platform.datasource.network.model.toBitwardenError
 import com.x8bit.bitwarden.data.platform.datasource.network.util.parseErrorBodyOrNull
+import com.x8bit.bitwarden.data.platform.util.asFailure
 import com.x8bit.bitwarden.data.vault.datasource.network.api.AzureApi
 import com.x8bit.bitwarden.data.vault.datasource.network.api.CiphersApi
 import com.x8bit.bitwarden.data.vault.datasource.network.model.AttachmentJsonRequest
@@ -57,21 +59,13 @@ class CiphersServiceImpl(
                 ciphersApi.uploadAttachment(
                     cipherId = requireNotNull(cipher.id),
                     attachmentId = attachmentJsonResponse.attachmentId,
-                    body = MultipartBody
-                        .Builder(
-                            boundary = "--BWMobileFormBoundary${clock.instant().toEpochMilli()}",
-                        )
-                        .addPart(
-                            part = MultipartBody.Part.createFormData(
-                                body = encryptedFile.asRequestBody(
-                                    contentType = "application/octet-stream".toMediaType(),
-                                ),
-                                name = "data",
-                                filename = cipher
-                                    .attachments
-                                    ?.find { it.id == attachmentJsonResponse.attachmentId }
-                                    ?.fileName,
-                            ),
+                    body = this
+                        .createMultipartBodyBuilder(
+                            encryptedFile = encryptedFile,
+                            filename = cipher
+                                .attachments
+                                ?.find { it.id == attachmentJsonResponse.attachmentId }
+                                ?.fileName,
                         )
                         .build(),
                 )
@@ -111,6 +105,36 @@ class CiphersServiceImpl(
                     ?: throw throwable
             }
 
+    @Suppress("ReturnCount")
+    override suspend fun shareAttachment(
+        cipherId: String,
+        attachment: Attachment,
+        organizationId: String,
+        encryptedFile: File,
+    ): Result<Unit> {
+        val attachmentId = attachment.id
+            ?: return IllegalStateException("Attachment must have ID").asFailure()
+        val attachmentKey = attachment.key
+            ?: return IllegalStateException("Attachment must have Key").asFailure()
+        return ciphersApi.shareAttachment(
+            cipherId = cipherId,
+            attachmentId = attachmentId,
+            organizationId = organizationId,
+            body = this
+                .createMultipartBodyBuilder(
+                    encryptedFile = encryptedFile,
+                    filename = attachment.fileName,
+                )
+                .addPart(
+                    part = MultipartBody.Part.createFormData(
+                        name = "key",
+                        value = attachmentKey,
+                    ),
+                )
+                .build(),
+        )
+    }
+
     override suspend fun shareCipher(
         cipherId: String,
         body: ShareCipherJsonRequest,
@@ -163,4 +187,20 @@ class CiphersServiceImpl(
 
     override suspend fun hasUnassignedCiphers(): Result<Boolean> =
         ciphersApi.hasUnassignedCiphers()
+
+    private fun createMultipartBodyBuilder(
+        encryptedFile: File,
+        filename: String?,
+    ): MultipartBody.Builder =
+        MultipartBody
+            .Builder(boundary = "--BWMobileFormBoundary${clock.instant().toEpochMilli()}")
+            .addPart(
+                part = MultipartBody.Part.createFormData(
+                    body = encryptedFile.asRequestBody(
+                        contentType = "application/octet-stream".toMediaType(),
+                    ),
+                    name = "data",
+                    filename = filename,
+                ),
+            )
 }

app/src/main/java/com/x8bit/bitwarden/data/vault/manager/CipherManagerImpl.kt

@@ -1,7 +1,6 @@
 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
@@ -241,15 +240,19 @@ class CipherManagerImpl(
         collectionIds: List<String>,
     ): ShareCipherResult {
         val userId = activeUserId ?: return ShareCipherResult.Error
-        return migrateAttachments(cipherView = cipherView)
+        return vaultSdkSource
+            .moveToOrganization(
+                userId = userId,
+                organizationId = organizationId,
+                cipherView = cipherView,
+            )
             .flatMap {
-                vaultSdkSource.moveToOrganization(
+                migrateAttachments(
                     userId = userId,
+                    cipherView = it,
                     organizationId = organizationId,
-                    cipherView = cipherView,
                 )
             }
-            .flatMap { vaultSdkSource.encryptCipher(userId = userId, cipherView = it) }
             .flatMap { cipher ->
                 ciphersService.shareCipher(
                     cipherId = cipherId,
@@ -486,15 +489,29 @@ class CipherManagerImpl(
         }
 
     @Suppress("ReturnCount")
-    private suspend fun migrateAttachments(cipherView: CipherView): Result<CipherView> {
+    private suspend fun migrateAttachments(
+        userId: String,
+        cipherView: CipherView,
+        organizationId: String,
+    ): Result<Cipher> {
         // Only run the migrations if we have attachments that do not have their own 'key'
-        val attachmentsToMigrate = cipherView.attachments.orEmpty().filter { it.key == null }
-        if (attachmentsToMigrate.none()) return cipherView.asSuccess()
+        val attachmentViewsToMigrate = cipherView.attachments.orEmpty().filter { it.key == null }
+        if (attachmentViewsToMigrate.none()) {
+            return vaultSdkSource.encryptCipher(userId = userId, cipherView = cipherView)
+        }
 
         val cipherViewId = cipherView.id
             ?: return IllegalStateException("CipherView must have an ID").asFailure()
+        val cipher = vaultSdkSource
+            .encryptCipher(userId = userId, cipherView = cipherView)
+            .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 {
-            attachmentsToMigrate.map { attachmentView ->
+            attachmentViewsToMigrate.map { attachmentView ->
                 async {
                     attachmentView
                         .id
@@ -504,31 +521,42 @@ class CipherManagerImpl(
                                     cipherView = cipherView,
                                     attachmentId = attachmentId,
                                 )
-                                .flatMap {
-                                    createAttachmentForResult(
-                                        cipherId = cipherViewId,
-                                        cipherView = cipherView,
-                                        fileSizeBytes = attachmentView.size,
-                                        fileName = attachmentView.fileName,
-                                        fileUri = it.toUri(),
-                                    )
+                                .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 }
+                                        }
                                 }
-                                .flatMap {
-                                    deleteCipherAttachmentForResult(
-                                        cipherId = cipherViewId,
-                                        attachmentId = attachmentId,
-                                        cipherView = cipherView,
-                                    )
-                                }
-                                .map { cipherView }
                         }
                         ?: IllegalStateException("AttachmentView must have an ID").asFailure()
                 }
             }
         }
 
-        return awaitAll(*migrations.toTypedArray())
-            .firstOrNull { it.isFailure }
-            ?: cipherView.asSuccess()
+        // 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()
     }
 }

app/src/test/java/com/x8bit/bitwarden/data/vault/datasource/network/service/CiphersServiceTest.kt

@@ -14,6 +14,7 @@ import com.x8bit.bitwarden.data.vault.datasource.network.model.createMockAttachm
 import com.x8bit.bitwarden.data.vault.datasource.network.model.createMockAttachmentJsonResponse
 import com.x8bit.bitwarden.data.vault.datasource.network.model.createMockCipher
 import com.x8bit.bitwarden.data.vault.datasource.network.model.createMockCipherJsonRequest
+import com.x8bit.bitwarden.data.vault.datasource.sdk.model.createMockSdkAttachment
 import io.mockk.every
 import io.mockk.mockk
 import io.mockk.mockkStatic
@@ -194,6 +195,58 @@ class CiphersServiceTest : BaseServiceTest() {
         assertEquals(Unit, result.getOrThrow())
     }
 
+    @Test
+    fun `shareAttachment without attachment ID should return an error`() = runTest {
+        val cipherId = "cipherId"
+        val organizationId = "organizationId"
+        val attachment = createMockSdkAttachment(number = 1).copy(id = null)
+        val encryptedFile = File.createTempFile("mockFile", "temp")
+
+        val result = ciphersService.shareAttachment(
+            cipherId = cipherId,
+            attachment = attachment,
+            organizationId = organizationId,
+            encryptedFile = encryptedFile,
+        )
+
+        assertTrue(result.isFailure)
+    }
+
+    @Test
+    fun `shareAttachment without attachment key should return an error`() = runTest {
+        val cipherId = "cipherId"
+        val organizationId = "organizationId"
+        val attachment = createMockSdkAttachment(number = 1, key = null)
+        val encryptedFile = File.createTempFile("mockFile", "temp")
+
+        val result = ciphersService.shareAttachment(
+            cipherId = cipherId,
+            attachment = attachment,
+            organizationId = organizationId,
+            encryptedFile = encryptedFile,
+        )
+
+        assertTrue(result.isFailure)
+    }
+
+    @Test
+    fun `shareAttachment should execute the share attachment API`() = runTest {
+        server.enqueue(MockResponse().setResponseCode(200))
+        val cipherId = "cipherId"
+        val organizationId = "organizationId"
+        val attachment = createMockSdkAttachment(number = 1)
+        val encryptedFile = File.createTempFile("mockFile", "temp")
+
+        val result = ciphersService.shareAttachment(
+            cipherId = cipherId,
+            attachment = attachment,
+            organizationId = organizationId,
+            encryptedFile = encryptedFile,
+        )
+
+        assertEquals(Unit, result.getOrThrow())
+    }
+
     @Test
     fun `shareCipher should execute the share cipher API`() = runTest {
         server.enqueue(

app/src/test/java/com/x8bit/bitwarden/data/vault/datasource/sdk/model/CipherViewUtil.kt

@@ -145,14 +145,14 @@ fun createMockUriView(number: Int): LoginUriView =
 /**
  * Create a mock [AttachmentView] with a given [number].
  */
-fun createMockAttachmentView(number: Int): AttachmentView =
+fun createMockAttachmentView(number: Int, key: String? = "mockKey-$number"): AttachmentView =
     AttachmentView(
         fileName = "mockFileName-$number",
         size = "1",
         sizeName = "mockSizeName-$number",
         id = "mockId-$number",
         url = "mockUrl-$number",
-        key = "mockKey-$number",
+        key = key,
     )
 
 /**

app/src/test/java/com/x8bit/bitwarden/data/vault/datasource/sdk/model/VaultSdkCipherUtil.kt

@@ -128,14 +128,14 @@ fun createMockSdkCard(number: Int): Card =
 /**
  * Create a mock [Attachment] with a given [number].
  */
-fun createMockSdkAttachment(number: Int): Attachment =
+fun createMockSdkAttachment(number: Int, key: String? = "mockKey-$number"): Attachment =
     Attachment(
         fileName = "mockFileName-$number",
         size = "1",
         sizeName = "mockSizeName-$number",
         id = "mockId-$number",
         url = "mockUrl-$number",
-        key = "mockKey-$number",
+        key = key,
     )
 
 /**

app/src/test/java/com/x8bit/bitwarden/data/vault/manager/CipherManagerTest.kt

@@ -2,6 +2,7 @@ package com.x8bit.bitwarden.data.vault.manager
 
 import android.net.Uri
 import com.bitwarden.vault.Attachment
+import com.bitwarden.vault.AttachmentView
 import com.bitwarden.vault.Cipher
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountTokensJson
@@ -16,6 +17,7 @@ import com.x8bit.bitwarden.data.vault.datasource.network.model.ShareCipherJsonRe
 import com.x8bit.bitwarden.data.vault.datasource.network.model.SyncResponseJson
 import com.x8bit.bitwarden.data.vault.datasource.network.model.UpdateCipherCollectionsJsonRequest
 import com.x8bit.bitwarden.data.vault.datasource.network.model.UpdateCipherResponseJson
+import com.x8bit.bitwarden.data.vault.datasource.network.model.createMockAttachment
 import com.x8bit.bitwarden.data.vault.datasource.network.model.createMockAttachmentJsonResponse
 import com.x8bit.bitwarden.data.vault.datasource.network.model.createMockCipher
 import com.x8bit.bitwarden.data.vault.datasource.network.model.createMockCipherJsonRequest
@@ -780,6 +782,107 @@ class CipherManagerTest {
             assertEquals(ShareCipherResult.Success, result)
         }
 
+    @Test
+    fun `shareCipher with attachment migration success should return ShareCipherResultSuccess`() =
+        runTest {
+            fakeAuthDiskSource.userState = MOCK_USER_STATE
+            val userId = "mockId-1"
+            val organizationId = "organizationId"
+            val mockAttachmentView = createMockAttachmentView(number = 1, key = null)
+            val initialCipherView = createMockCipherView(number = 1).copy(
+                attachments = listOf(mockAttachmentView),
+            )
+            val mockAttachment = createMockSdkAttachment(number = 1, key = null)
+            val mockCipher = createMockSdkCipher(number = 1).copy(
+                attachments = listOf(mockAttachment),
+            )
+            val attachment = createMockAttachment(number = 1)
+            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()
+
+            // Handle mocks for migration
+            coEvery {
+                vaultSdkSource.encryptCipher(userId = userId, cipherView = initialCipherView)
+            } returns mockCipher.asSuccess()
+            coEvery {
+                ciphersService.getCipherAttachment(cipherId = "mockId-1", attachmentId = "mockId-1")
+            } returns attachment.asSuccess()
+            coEvery {
+                fileManager.downloadFileToCache(url = "mockUrl-1")
+            } returns DownloadResult.Success(file = encryptedFile)
+            coEvery {
+                vaultSdkSource.decryptFile(
+                    userId = userId,
+                    cipher = mockCipher,
+                    attachment = mockAttachment,
+                    encryptedFilePath = encryptedFile.path,
+                    decryptedFilePath = decryptedFile.path,
+                )
+            } returns Unit.asSuccess()
+            coEvery {
+                vaultSdkSource.encryptAttachment(
+                    userId = userId,
+                    cipher = mockCipher,
+                    attachmentView = AttachmentView(
+                        id = null,
+                        url = null,
+                        size = "1",
+                        sizeName = null,
+                        fileName = "mockFileName-1",
+                        key = null,
+                    ),
+                    decryptedFilePath = any(),
+                    encryptedFilePath = any(),
+                )
+            } returns mockAttachment.asSuccess()
+            coEvery {
+                ciphersService.shareAttachment(
+                    cipherId = "mockId-1",
+                    attachment = mockAttachment,
+                    organizationId = organizationId,
+                    encryptedFile = any(),
+                )
+            } returns Unit.asSuccess()
+            // Done with mocks for migration
+
+            coEvery {
+                vaultSdkSource.encryptCipher(userId = userId, cipherView = mockCipherView)
+            } returns createMockSdkCipher(number = 1, clock = clock).asSuccess()
+            coEvery {
+                ciphersService.shareCipher(
+                    cipherId = "mockId-1",
+                    body = ShareCipherJsonRequest(
+                        cipher = createMockCipherJsonRequest(number = 1, hasNullUri = true),
+                        collectionIds = listOf("mockId-1"),
+                    ),
+                )
+            } returns createMockCipher(number = 1).asSuccess()
+            coEvery {
+                vaultDiskSource.saveCipher(
+                    userId = userId,
+                    cipher = createMockCipher(number = 1).copy(collectionIds = listOf("mockId-1")),
+                )
+            } just runs
+
+            val result = cipherManager.shareCipher(
+                cipherId = "mockId-1",
+                organizationId = organizationId,
+                cipherView = initialCipherView,
+                collectionIds = listOf("mockId-1"),
+            )
+
+            assertEquals(ShareCipherResult.Success, result)
+        }
+
     @Test
     @Suppress("MaxLineLength")
     fun `shareCipher with cipherService shareCipher failure should return ShareCipherResultError`() =