From 6f9147b2b25ebcf4dfe4fec944535981300f144c Mon Sep 17 00:00:00 2001
From: David Perez <david@livefront.com>
Date: Thu, 25 Jan 2024 10:33:54 -0600
Subject: [PATCH] Add method for encrypting an attachment (#774)

---
 .../vault/datasource/sdk/VaultSdkSource.kt    | 16 +++++++++
 .../datasource/sdk/VaultSdkSourceImpl.kt      | 19 ++++++++++
 .../datasource/sdk/VaultSdkSourceTest.kt      | 36 +++++++++++++++++++
 3 files changed, 71 insertions(+)

diff --git a/app/src/main/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSource.kt b/app/src/main/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSource.kt
index b8b2386663..34060fc359 100644
--- a/app/src/main/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSource.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSource.kt
@@ -1,5 +1,7 @@
 package com.x8bit.bitwarden.data.vault.datasource.sdk
 
+import com.bitwarden.core.AttachmentEncryptResult
+import com.bitwarden.core.AttachmentView
 import com.bitwarden.core.Cipher
 import com.bitwarden.core.CipherListView
 import com.bitwarden.core.CipherView
@@ -87,6 +89,20 @@ interface VaultSdkSource {
         request: InitOrgCryptoRequest,
     ): Result<InitializeCryptoResult>
 
+    /**
+     * Encrypts a [AttachmentView] for the user with the given [userId], returning an
+     * [AttachmentEncryptResult] wrapped in a [Result].
+     *
+     * This should only be called after a successful call to [initializeCrypto] for the associated
+     * user.
+     */
+    suspend fun encryptAttachment(
+        userId: String,
+        cipher: Cipher,
+        attachmentView: AttachmentView,
+        fileBuffer: ByteArray,
+    ): Result<AttachmentEncryptResult>
+
     /**
      * Encrypts a [CipherView] for the user with the given [userId], returning a [Cipher] wrapped
      * in a [Result].
diff --git a/app/src/main/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSourceImpl.kt b/app/src/main/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSourceImpl.kt
index d75e491549..9fe6971643 100644
--- a/app/src/main/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSourceImpl.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSourceImpl.kt
@@ -1,5 +1,7 @@
 package com.x8bit.bitwarden.data.vault.datasource.sdk
 
+import com.bitwarden.core.AttachmentEncryptResult
+import com.bitwarden.core.AttachmentView
 import com.bitwarden.core.Cipher
 import com.bitwarden.core.CipherListView
 import com.bitwarden.core.CipherView
@@ -121,6 +123,23 @@ class VaultSdkSourceImpl(
                 )
         }
 
+    override suspend fun encryptAttachment(
+        userId: String,
+        cipher: Cipher,
+        attachmentView: AttachmentView,
+        fileBuffer: ByteArray,
+    ): Result<AttachmentEncryptResult> =
+        runCatching {
+            getClient(userId = userId)
+                .vault()
+                .attachments()
+                .encryptBuffer(
+                    cipher = cipher,
+                    attachment = attachmentView,
+                    buffer = fileBuffer,
+                )
+        }
+
     override suspend fun encryptCipher(
         userId: String,
         cipherView: CipherView,
diff --git a/app/src/test/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSourceTest.kt b/app/src/test/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSourceTest.kt
index fe14236ece..dbc4593478 100644
--- a/app/src/test/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSourceTest.kt
+++ b/app/src/test/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSourceTest.kt
@@ -1,5 +1,7 @@
 package com.x8bit.bitwarden.data.vault.datasource.sdk
 
+import com.bitwarden.core.AttachmentEncryptResult
+import com.bitwarden.core.AttachmentView
 import com.bitwarden.core.Cipher
 import com.bitwarden.core.CipherListView
 import com.bitwarden.core.CipherView
@@ -37,6 +39,7 @@ import kotlinx.coroutines.test.runTest
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.Test
 
+@Suppress("LargeClass")
 class VaultSdkSourceTest {
     private val clientCrypto = mockk<ClientCrypto>()
     private val clientPasswordHistory = mockk<ClientPasswordHistory>()
@@ -472,6 +475,39 @@ class VaultSdkSourceTest {
             verify { sdkClientManager.getOrCreateClient(userId = userId) }
         }
 
+    @Test
+    fun `encryptAttachment should call SDK and return correct data wrapped in a Result`() =
+        runBlocking {
+            val userId = "userId"
+            val expectedResult = mockk<AttachmentEncryptResult>()
+            val mockCipher = mockk<Cipher>()
+            val mockAttachmentView = mockk<AttachmentView>()
+            val fileBuffer = byteArrayOf(1, 2)
+            coEvery {
+                clientVault.attachments().encryptBuffer(
+                    cipher = mockCipher,
+                    attachment = mockAttachmentView,
+                    buffer = fileBuffer,
+                )
+            } returns expectedResult
+
+            val result = vaultSdkSource.encryptAttachment(
+                userId = userId,
+                cipher = mockCipher,
+                attachmentView = mockAttachmentView,
+                fileBuffer = fileBuffer,
+            )
+
+            assertEquals(expectedResult.asSuccess(), result)
+            coVerify {
+                clientVault.attachments().encryptBuffer(
+                    cipher = mockCipher,
+                    attachment = mockAttachmentView,
+                    buffer = fileBuffer,
+                )
+            }
+        }
+
     @Test
     fun `encryptSend should call SDK and return correct data wrapped in a Result`() = runBlocking {
         val userId = "userId"