BIT-2407: Update cipher on server when SDK adds a cipher key (#667)

Author: matt-livefront

BitwardenShared/Core/Vault/Repositories/VaultRepository.swift

@@ -375,6 +375,23 @@ class DefaultVaultRepository { // swiftlint:disable:this type_body_length
 
     // MARK: Private
 
+    /// Encrypts the cipher. If the cipher was migrated by the SDK (e.g. added a cipher key), the
+    /// cipher will be updated locally and on the server.
+    ///
+    /// - Parameter cipherView: The cipher to encrypt.
+    /// - Returns: The encrypted cipher.
+    ///
+    private func encryptAndUpdateCipher(_ cipherView: CipherView) async throws -> Cipher {
+        let cipher = try await clientService.vault().ciphers().encrypt(cipherView: cipherView)
+
+        let didAddCipherKey = cipherView.key == nil && cipher.key != nil
+        if didAddCipherKey {
+            try await cipherService.updateCipherWithServer(cipher)
+        }
+
+        return cipher
+    }
+
     /// Returns a list of `VaultListItem`s for the folders within a nested tree. By default, this
     /// will return the list items for the folders at the root of the tree. Specifying a
     /// `nestedFolderId` will return the list items for the children of the folder with the
@@ -952,7 +969,7 @@ extension DefaultVaultRepository: VaultRepository {
         else { throw BitwardenError.dataError("Missing data") }
 
         // Get the encrypted cipher and attachment, then download the actual data of the attachment.
-        let encryptedCipher = try await clientService.vault().ciphers().encrypt(cipherView: cipher)
+        let encryptedCipher = try await encryptAndUpdateCipher(cipher)
         guard let attachment = encryptedCipher.attachments?.first(where: { $0.id == attachmentId }),
               let downloadedUrl = try await cipherService.downloadAttachment(withId: attachmentId, cipherId: cipherId)
         else { return nil }
@@ -1032,7 +1049,7 @@ extension DefaultVaultRepository: VaultRepository {
     func restoreCipher(_ cipher: BitwardenSdk.CipherView) async throws {
         guard let id = cipher.id else { throw CipherAPIServiceError.updateMissingId }
         let restoredCipher = cipher.update(deletedDate: nil)
-        let encryptCipher = try await clientService.vault().ciphers().encrypt(cipherView: restoredCipher)
+        let encryptCipher = try await encryptAndUpdateCipher(restoredCipher)
         try await cipherService.restoreCipherWithServer(id: id, encryptCipher)
     }
 
@@ -1048,7 +1065,7 @@ extension DefaultVaultRepository: VaultRepository {
         )
 
         // Encrypt the attachment.
-        let cipher = try await clientService.vault().ciphers().encrypt(cipherView: cipherView)
+        let cipher = try await encryptAndUpdateCipher(cipherView)
         let attachment = try await clientService.vault().attachments().encryptBuffer(
             cipher: cipher,
             attachment: attachmentView,
@@ -1091,8 +1108,8 @@ extension DefaultVaultRepository: VaultRepository {
     func softDeleteCipher(_ cipher: CipherView) async throws {
         guard let id = cipher.id else { throw CipherAPIServiceError.updateMissingId }
         let softDeletedCipher = cipher.update(deletedDate: timeProvider.presentTime)
-        let encryptCipher = try await clientService.vault().ciphers().encrypt(cipherView: softDeletedCipher)
-        try await cipherService.softDeleteCipherWithServer(id: id, encryptCipher)
+        let encryptedCipher = try await encryptAndUpdateCipher(softDeletedCipher)
+        try await cipherService.softDeleteCipherWithServer(id: id, encryptedCipher)
     }
 
     func updateCipher(_ cipherView: CipherView) async throws {
@@ -1101,7 +1118,7 @@ extension DefaultVaultRepository: VaultRepository {
     }
 
     func updateCipherCollections(_ cipherView: CipherView) async throws {
-        let cipher = try await clientService.vault().ciphers().encrypt(cipherView: cipherView)
+        let cipher = try await encryptAndUpdateCipher(cipherView)
         try await cipherService.updateCipherCollectionsWithServer(cipher)
     }
 

BitwardenShared/Core/Vault/Repositories/VaultRepositoryTests.swift

@@ -228,6 +228,29 @@ class VaultRepositoryTests: BitwardenTestCase { // swiftlint:disable:this type_b
         }
     }
 
+    /// `downloadAttachment(_:cipher:)` updates the cipher on the server if the SDK adds a cipher key.
+    func test_downloadAttachment_updatesMigratedCipher() async throws {
+        stateService.activeAccount = .fixture()
+        let downloadUrl = FileManager.default.temporaryDirectory.appendingPathComponent("sillyGoose.txt")
+        try Data("🪿".utf8).write(to: downloadUrl)
+        cipherService.downloadAttachmentResult = .success(downloadUrl)
+        let attachment = AttachmentView.fixture(fileName: "sillyGoose.txt")
+        let cipherView = CipherView.fixture(attachments: [attachment])
+        let cipher = Cipher.fixture(
+            attachments: [Attachment(attachmentView: attachment)],
+            key: "new key"
+        )
+        clientCiphers.encryptCipherResult = .success(cipher)
+
+        let resultUrl = try await subject.downloadAttachment(attachment, cipher: cipherView)
+
+        XCTAssertEqual(clientService.mockVault.clientCiphers.encryptedCiphers.last, cipherView)
+        XCTAssertEqual(cipherService.downloadAttachmentId, attachment.id)
+        XCTAssertEqual(clientService.mockVault.clientAttachments.encryptedFilePaths.last, downloadUrl.path)
+        XCTAssertEqual(cipherService.updateCipherWithServerCiphers, [cipher])
+        XCTAssertEqual(resultUrl?.lastPathComponent, "sillyGoose.txt")
+    }
+
     /// `fetchCipher(withId:)` returns the cipher if it exists and `nil` otherwise.
     func test_fetchCipher() async throws {
         var cipher = try await subject.fetchCipher(withId: "1")
@@ -1235,6 +1258,19 @@ class VaultRepositoryTests: BitwardenTestCase { // swiftlint:disable:this type_b
         XCTAssertEqual(cipherService.restoredCipherId, "123")
     }
 
+    /// `restoreCipher(_:cipher:)` updates the cipher on the server if the SDK adds a cipher key.
+    func test_restoreCipher_updatesMigratedCipher() async throws {
+        stateService.activeAccount = .fixture()
+        let cipherView = CipherView.fixture(deletedDate: .now)
+        let cipher = Cipher.fixture(key: "new key")
+        clientCiphers.encryptCipherResult = .success(cipher)
+
+        try await subject.restoreCipher(cipherView)
+
+        XCTAssertEqual(cipherService.restoredCipher, cipher)
+        XCTAssertEqual(cipherService.updateCipherWithServerCiphers, [cipher])
+    }
+
     /// `saveAttachment(cipherView:fileData:fileName:)` saves the attachment to the cipher.
     func test_saveAttachment() async throws {
         cipherService.saveAttachmentWithServerResult = .success(.fixture(id: "42"))
@@ -1253,6 +1289,25 @@ class VaultRepositoryTests: BitwardenTestCase { // swiftlint:disable:this type_b
         XCTAssertEqual(updatedCipher.id, "42")
     }
 
+    /// `saveAttachment(cipherView:fileData:fileName:)` updates the cipher on the server if the SDK adds a cipher key.
+    func test_saveAttachment_updatesMigratedCipher() async throws {
+        cipherService.saveAttachmentWithServerResult = .success(.fixture(id: "42"))
+        let cipher = Cipher.fixture(key: "new key")
+        clientCiphers.encryptCipherResult = .success(cipher)
+
+        let updatedCipher = try await subject.saveAttachment(
+            cipherView: .fixture(),
+            fileData: Data(),
+            fileName: "Pineapple on pizza"
+        )
+
+        XCTAssertEqual(clientService.mockVault.clientCiphers.encryptedCiphers, [.fixture()])
+        XCTAssertEqual(clientService.mockVault.clientAttachments.encryptedBuffers, [Data()])
+        XCTAssertEqual(cipherService.updateCipherWithServerCiphers, [cipher])
+        XCTAssertEqual(cipherService.saveAttachmentWithServerCipher, cipher)
+        XCTAssertEqual(updatedCipher.id, "42")
+    }
+
     /// `softDeleteCipher()` throws on id errors.
     func test_softDeleteCipher_idError_nil() async throws {
         stateService.accounts = [.fixtureAccountLogin()]
@@ -1275,6 +1330,19 @@ class VaultRepositoryTests: BitwardenTestCase { // swiftlint:disable:this type_b
         XCTAssertEqual(cipherService.softDeleteCipherId, "123")
     }
 
+    /// `softDeleteCipher(_:cipher:)` updates the cipher on the server if the SDK adds a cipher key.
+    func test_softDeleteCipher_updatesMigratedCipher() async throws {
+        stateService.activeAccount = .fixture()
+        let cipherView = CipherView.fixture(deletedDate: .now)
+        let cipher = Cipher.fixture(key: "new key")
+        clientCiphers.encryptCipherResult = .success(cipher)
+
+        try await subject.softDeleteCipher(cipherView)
+
+        XCTAssertEqual(cipherService.softDeleteCipher, cipher)
+        XCTAssertEqual(cipherService.updateCipherWithServerCiphers, [cipher])
+    }
+
     /// `vaultListPublisher(group:filter:)` returns a publisher for the vault list items.
     func test_vaultListPublisher_groups_card() async throws {
         let cipher = Cipher.fixture(id: "1", type: .card)

BitwardenShared/Core/Vault/Services/SyncService.swift

@@ -272,9 +272,8 @@ extension DefaultSyncService {
         try await folderService.deleteFolderWithLocalStorage(id: data.id)
 
         let updatedCiphers = try await cipherService.fetchAllCiphers()
-            .asyncMap { try await clientService.vault().ciphers().decrypt(cipher: $0) }
+            .filter { $0.folderId == data.id }
             .map { $0.update(folderId: nil) }
-            .asyncMap { try await clientService.vault().ciphers().encrypt(cipherView: $0) }
 
         for cipher in updatedCiphers {
             try await cipherService.updateCipherWithLocalStorage(cipher)

BitwardenShared/Core/Vault/Services/SyncServiceTests.swift

@@ -606,7 +606,10 @@ class SyncServiceTests: BitwardenTestCase {
     func test_deleteFolder() async throws {
         stateService.activeAccount = .fixture()
         folderService.deleteFolderWithLocalStorageResult = .success(())
-        cipherService.fetchAllCiphersResult = .success([.fixture(folderId: "id")])
+        cipherService.fetchAllCiphersResult = .success([
+            .fixture(folderId: "id", id: "1"),
+            .fixture(folderId: "other", id: "2"),
+        ])
         cipherService.updateCipherWithLocalStorageResult = .success(())
 
         let notification = SyncFolderNotification(
@@ -616,7 +619,10 @@ class SyncServiceTests: BitwardenTestCase {
         )
         try await subject.deleteFolder(data: notification)
         XCTAssertEqual(folderService.deleteFolderWithLocalStorageId, "id")
-        XCTAssertEqual(cipherService.updateCipherWithLocalStorageCipher, .fixture(folderId: nil))
+        XCTAssertEqual(
+            cipherService.updateCipherWithLocalStorageCiphers,
+            [.fixture(folderId: nil, id: "1")]
+        )
     }
 
     func test_deleteSend() async throws {

BitwardenShared/Core/Vault/Services/TestHelpers/MockCipherService.swift

@@ -51,7 +51,7 @@ class MockCipherService: CipherService {
     var syncCipherWithServerId: String?
     var syncCipherWithServerResult: Result<Void, Error> = .success(())
 
-    var updateCipherWithLocalStorageCipher: BitwardenSdk.Cipher?
+    var updateCipherWithLocalStorageCiphers = [BitwardenSdk.Cipher]()
     var updateCipherWithLocalStorageResult: Result<Void, Error> = .success(())
 
     var updateCipherWithServerCiphers = [Cipher]()
@@ -135,7 +135,7 @@ class MockCipherService: CipherService {
     }
 
     func updateCipherWithLocalStorage(_ cipher: Cipher) async throws {
-        updateCipherWithLocalStorageCipher = cipher
+        updateCipherWithLocalStorageCiphers.append(cipher)
         return try updateCipherWithLocalStorageResult.get()
     }
 

BitwardenShared/Core/Vault/Services/TestHelpers/MockClientVaultService.swift

@@ -86,6 +86,7 @@ class MockClientAttachments: ClientAttachmentsProtocol {
 // MARK: - MockClientCiphers
 
 class MockClientCiphers: ClientCiphersProtocol {
+    var encryptCipherResult: Result<Cipher, Error>?
     var encryptError: Error?
     var encryptedCiphers = [CipherView]()
     var moveToOrganizationCipher: CipherView?
@@ -106,7 +107,7 @@ class MockClientCiphers: ClientCiphersProtocol {
         if let encryptError {
             throw encryptError
         }
-        return Cipher(cipherView: cipherView)
+        return try encryptCipherResult?.get() ?? Cipher(cipherView: cipherView)
     }
 
     func moveToOrganization(

BitwardenShared/UI/Vault/VaultItem/ViewItem/ViewLoginItem/Extensions/Cipher+Update.swift

@@ -34,4 +34,38 @@ extension Cipher {
             revisionDate: revisionDate
         )
     }
+
+    /// Returns a copy of the existing cipher with an updated folder ID.
+    ///
+    /// - Parameter folderId: The folder ID to update in the cipher.
+    /// - Returns: A copy of the existing cipher, with the specified properties updated.
+    ///
+    func update(folderId: String?) -> Cipher {
+        Cipher(
+            id: id,
+            organizationId: organizationId,
+            folderId: folderId,
+            collectionIds: collectionIds,
+            key: key,
+            name: name,
+            notes: notes,
+            type: type,
+            login: login,
+            identity: identity,
+            card: card,
+            secureNote: secureNote,
+            favorite: favorite,
+            reprompt: reprompt,
+            organizationUseTotp: organizationUseTotp,
+            edit: edit,
+            viewPassword: viewPassword,
+            localData: localData,
+            attachments: attachments,
+            fields: fields,
+            passwordHistory: passwordHistory,
+            creationDate: creationDate,
+            deletedDate: deletedDate,
+            revisionDate: revisionDate
+        )
+    }
 }