diff --git a/Session/Settings/DeveloperSettingsViewModel.swift b/Session/Settings/DeveloperSettingsViewModel.swift index 9b2572427..654dbb9a8 100644 --- a/Session/Settings/DeveloperSettingsViewModel.swift +++ b/Session/Settings/DeveloperSettingsViewModel.swift @@ -200,10 +200,12 @@ class DeveloperSettingsViewModel: SessionTableViewModel, NavigatableStateHolder, explanation: NSAttributedString( string: """ This will generate a file encrypted using the provided password includes all app data, attachments, settings and keys. - - We've generated a secure password for you but feel free to provide your own. + + This exported file can only be imported by Session iOS. Use at your own risk! + + We've generated a secure password for you but feel free to provide your own. """ ), placeholder: "Enter a password", @@ -234,25 +236,6 @@ class DeveloperSettingsViewModel: SessionTableViewModel, NavigatableStateHolder, } private func importDatabase(_ targetView: UIView?) { - func showError(_ error: Error) { - self.transitionToScreen( - ConfirmationModal( - info: ConfirmationModal.Info( - title: "theError".localized(), - body: { - switch error { - case CryptoKitError.incorrectKeySize: - return .text("The password must be between 6 and 32 characters (padded to 32 bytes)") - - default: return .text("Failed to export database") - } - }() - ) - ), - transitionType: .present - ) - } - self.databaseKeyEncryptionPassword = "" self.transitionToScreen( ConfirmationModal( @@ -262,6 +245,8 @@ class DeveloperSettingsViewModel: SessionTableViewModel, NavigatableStateHolder, explanation: NSAttributedString( string: """ Importing a database will result in the loss of all data stored locally. + + This can only import backup files exported by Session iOS. Use at your own risk! """ @@ -277,57 +262,7 @@ class DeveloperSettingsViewModel: SessionTableViewModel, NavigatableStateHolder, dismissOnConfirm: false, onConfirm: { [weak self] modal in modal.dismiss(animated: true) { - guard - let password: String = self?.databaseKeyEncryptionPassword, - password.count >= 6 - else { return showError(CryptoKitError.incorrectKeySize) } - - let documentPickerResult: DocumentPickerResult = DocumentPickerResult { url in - guard let url: URL = url else { return } - - let viewController: UIViewController = ModalActivityIndicatorViewController(canCancel: false) { modalActivityIndicator in - do { - let tmpUnencryptPath: String = "\(Singleton.appContext.temporaryDirectory)/new_session.bak" - let extraFilePaths: [String] = try DirectoryArchiver.unarchiveDirectory( - archivePath: url.path, - destinationPath: tmpUnencryptPath, - password: password, - progressChanged: { fileProgress, fileSize in - let percentage: Int = { - guard fileSize > 0 else { return 0 } - - return Int((Double(fileProgress) / Double(fileSize)) * 100) - }() - - DispatchQueue.main.async { - modalActivityIndicator.setMessage( - "Decryption progress: \(percentage)%" - ) - } - } - ) - - // TODO: Need to actually replace the current content then kill the app - // TODO: Might be nice to validate that we have database access to the new database with the key - print("RAWR") - modalActivityIndicator.dismiss { - print("RAWR2") - } - } - catch { showError(error) } - } - - self?.transitionToScreen(viewController, transitionType: .present) - } - self?.documentPickerResult = documentPickerResult - - // UIDocumentPickerModeImport copies to a temp file within our container. - // It uses more memory than "open" but lets us avoid working with security scoped URLs. - let documentPickerVC = UIDocumentPickerViewController(forOpeningContentTypes: [.item], asCopy: true) - documentPickerVC.delegate = documentPickerResult - documentPickerVC.modalPresentationStyle = .fullScreen - - self?.transitionToScreen(documentPickerVC, transitionType: .present) + self?.performImport() } } ) @@ -349,6 +284,8 @@ class DeveloperSettingsViewModel: SessionTableViewModel, NavigatableStateHolder, switch error { case CryptoKitError.incorrectKeySize: return .text("The password must be between 6 and 32 characters (padded to 32 bytes)") + case is DatabaseError: + return .text("An error occurred finalising pending changes in the database") default: return .text("Failed to export database") } @@ -365,6 +302,9 @@ class DeveloperSettingsViewModel: SessionTableViewModel, NavigatableStateHolder, let backupFile: String = "\(Singleton.appContext.temporaryDirectory)/session.bak" do { + /// Perform a full checkpoint to ensure any pending changes are written to the main database file + try dependencies.storage.checkpoint(.truncate) + let secureDbKey: String = try dependencies.storage.secureExportKey( password: databaseKeyEncryptionPassword ) @@ -372,6 +312,11 @@ class DeveloperSettingsViewModel: SessionTableViewModel, NavigatableStateHolder, try DirectoryArchiver.archiveDirectory( sourcePath: FileManager.default.appSharedDataDirectoryPath, destinationPath: backupFile, + filenamesToExclude: [ + ".DS_Store", + "\(Storage.dbFileName)-wal", + "\(Storage.dbFileName)-shm" + ], additionalPaths: [secureDbKey], password: databaseKeyEncryptionPassword, progressChanged: { fileIndex, totalFiles, currentFileProgress, currentFileSize in @@ -394,7 +339,12 @@ class DeveloperSettingsViewModel: SessionTableViewModel, NavigatableStateHolder, } ) } - catch { return showError(error) } + catch { + modalActivityIndicator.dismiss { + showError(error) + } + return + } modalActivityIndicator.dismiss { switch viaShareSheet { @@ -431,6 +381,173 @@ class DeveloperSettingsViewModel: SessionTableViewModel, NavigatableStateHolder, self.transitionToScreen(viewController, transitionType: .present) } + + private func performImport() { + func showError(_ error: Error) { + self.transitionToScreen( + ConfirmationModal( + info: ConfirmationModal.Info( + title: "theError".localized(), + body: { + switch error { + case CryptoKitError.incorrectKeySize: + return .text("The password must be between 6 and 32 characters (padded to 32 bytes)") + + case is DatabaseError: return .text("Database key in backup file was invalid.") + default: return .text("\(error)") + } + }() + ) + ), + transitionType: .present + ) + } + + guard databaseKeyEncryptionPassword.count >= 6 else { return showError(CryptoKitError.incorrectKeySize) } + + let documentPickerResult: DocumentPickerResult = DocumentPickerResult { url in + guard let url: URL = url else { return } + + let viewController: UIViewController = ModalActivityIndicatorViewController(canCancel: false) { [weak self, password = self.databaseKeyEncryptionPassword, dependencies = self.dependencies] modalActivityIndicator in + do { + let tmpUnencryptPath: String = "\(Singleton.appContext.temporaryDirectory)/new_session.bak" + let (paths, additionalFilePaths): ([String], [String]) = try DirectoryArchiver.unarchiveDirectory( + archivePath: url.path, + destinationPath: tmpUnencryptPath, + password: password, + progressChanged: { filesSaved, totalFiles, fileProgress, fileSize in + let percentage: Int = { + guard fileSize > 0 else { return 0 } + + return Int((Double(fileProgress) / Double(fileSize)) * 100) + }() + + DispatchQueue.main.async { + modalActivityIndicator.setMessage([ + "Decryption progress: \(percentage)%", + "Files imported: \(filesSaved)/\(totalFiles)" + ].compactMap { $0 }.joined(separator: "\n")) + } + } + ) + + /// Test that we actually have valid access to the database + guard + let encKeyPath: String = additionalFilePaths + .first(where: { $0.hasSuffix(Storage.encKeyFilename) }), + let databasePath: String = paths + .first(where: { $0.hasSuffix(Storage.dbFileName) }) + else { throw ArchiveError.unableToFindDatabaseKey } + + DispatchQueue.main.async { + modalActivityIndicator.setMessage( + "Checking for valid database..." + ) + } + + let testStorage: Storage = try Storage( + testAccessTo: databasePath, + encryptedKeyPath: encKeyPath, + encryptedKeyPassword: password + ) + + guard testStorage.isValid else { + throw ArchiveError.decryptionFailed + } + + /// Now that we have confirmed access to the replacement database we need to + /// stop the current account from doing anything + DispatchQueue.main.async { + modalActivityIndicator.setMessage( + "Clearing current account data..." + ) + + (UIApplication.shared.delegate as? AppDelegate)?.stopPollers() + } + + dependencies.jobRunner.stopAndClearPendingJobs(using: dependencies) + LibSession.suspendNetworkAccess() + dependencies.storage.suspendDatabaseAccess() + try dependencies.storage.closeDatabase() + + let deleteEnumerator: FileManager.DirectoryEnumerator? = FileManager.default.enumerator( + at: URL( + fileURLWithPath: FileManager.default.appSharedDataDirectoryPath + ), + includingPropertiesForKeys: [.isRegularFileKey] + ) + let fileUrls: [URL] = (deleteEnumerator?.allObjects.compactMap { $0 as? URL } ?? []) + try fileUrls.forEach { url in + /// The database `wal` and `shm` files might not exist anymore at this point + /// so we should only remove files which exist to prevent errors + guard FileManager.default.fileExists(atPath: url.path) else { return } + + try FileManager.default.removeItem(atPath: url.path) + } + + /// Current account data has been removed, we now need to copy over the + /// newly imported data + DispatchQueue.main.async { + modalActivityIndicator.setMessage( + "Moving imported data..." + ) + } + + try paths.forEach { path in + /// Need to ensure the destination directry + let targetPath: String = [ + FileManager.default.appSharedDataDirectoryPath, + path.replacingOccurrences(of: tmpUnencryptPath, with: "") + ].joined() // Already has '/' after 'appSharedDataDirectoryPath' + + try FileManager.default.createDirectory( + atPath: URL(fileURLWithPath: targetPath) + .deletingLastPathComponent() + .path, + withIntermediateDirectories: true + ) + try FileManager.default.moveItem(atPath: path, toPath: targetPath) + } + + /// All of the main files have been moved across, we now need to replace the current database key with + /// the one included in the backup + try dependencies.storage.replaceDatabaseKey(path: encKeyPath, password: password) + + /// The import process has completed so we need to restart the app + DispatchQueue.main.async { + self?.transitionToScreen( + ConfirmationModal( + info: ConfirmationModal.Info( + title: "Import Complete", + body: .text("The import completed successfully, Session must be reopened in order to complete the process."), + cancelTitle: "Exit", + cancelStyle: .alert_text, + onCancel: { _ in exit(0) } + ) + ), + transitionType: .present + ) + } + } + catch { + modalActivityIndicator.dismiss { + showError(error) + } + } + } + + self.transitionToScreen(viewController, transitionType: .present) + } + self.documentPickerResult = documentPickerResult + + // UIDocumentPickerModeImport copies to a temp file within our container. + // It uses more memory than "open" but lets us avoid working with security scoped URLs. + let documentPickerVC = UIDocumentPickerViewController(forOpeningContentTypes: [.item], asCopy: true) + documentPickerVC.delegate = documentPickerResult + documentPickerVC.modalPresentationStyle = .fullScreen + + self.transitionToScreen(documentPickerVC, transitionType: .present) + } } private class DocumentPickerResult: NSObject, UIDocumentPickerDelegate { diff --git a/Session/Settings/HelpViewModel.swift b/Session/Settings/HelpViewModel.swift index 89062aa02..2530fd329 100644 --- a/Session/Settings/HelpViewModel.swift +++ b/Session/Settings/HelpViewModel.swift @@ -232,26 +232,3 @@ class HelpViewModel: SessionTableViewModel, NavigatableStateHolder, ObservableTa } } } - -private class DocumentPickerResult: NSObject, UIDocumentPickerDelegate { - private let onResult: (URL?) -> Void - - init(onResult: @escaping (URL?) -> Void) { - self.onResult = onResult - } - - // MARK: - UIDocumentPickerDelegate - - func documentPicker(_ controller: UIDocumentPickerViewController, didPickDocumentsAt urls: [URL]) { - guard let url: URL = urls.first else { - self.onResult(nil) - return - } - - self.onResult(url) - } - - func documentPickerWasCancelled(_ controller: UIDocumentPickerViewController) { - self.onResult(nil) - } -} diff --git a/SessionUtilitiesKit/Database/Storage.swift b/SessionUtilitiesKit/Database/Storage.swift index 57013e1b3..65f87eaef 100644 --- a/SessionUtilitiesKit/Database/Storage.swift +++ b/SessionUtilitiesKit/Database/Storage.swift @@ -15,7 +15,7 @@ public extension KeychainStorage.DataKey { static let dbCipherKeySpec: Self = "G open class Storage { public static let queuePrefix: String = "SessionDatabase" - private static let dbFileName: String = "Session.sqlite" + public static let dbFileName: String = "Session.sqlite" private static let kSQLCipherKeySpecLength: Int = 48 /// If a transaction takes longer than this duration a warning will be logged but the transaction will continue to run @@ -69,6 +69,18 @@ open class Storage { configureDatabase(customWriter: customWriter) } + public init( + testAccessTo databasePath: String, + encryptedKeyPath: String, + encryptedKeyPassword: String + ) throws { + try testAccess( + databasePath: databasePath, + encryptedKeyPath: encryptedKeyPath, + encryptedKeyPassword: encryptedKeyPassword + ) + } + private func configureDatabase(customWriter: DatabaseWriter? = nil) { // Create the database directory if needed and ensure it's protection level is set before attempting to // create the database KeySpec or the database itself @@ -446,6 +458,16 @@ open class Storage { Log.info("[Storage] Database access resumed.") } + public func checkpoint(_ mode: Database.CheckpointMode) throws { + try dbWriter?.writeWithoutTransaction { db in _ = try db.checkpoint(mode) } + } + + public func closeDatabase() throws { + suspendDatabaseAccess() + isValid = false + dbWriter = nil + } + public func resetAllStorage() { isValid = false Storage.internalHasCreatedValidInstance.mutate { $0 = false } @@ -733,6 +755,58 @@ public extension ValueObservation { // MARK: - Debug Convenience public extension Storage { + static let encKeyFilename: String = "key.enc" + + func testAccess( + databasePath: String, + encryptedKeyPath: String, + encryptedKeyPassword: String + ) throws { + /// First we need to ensure we can decrypt the encrypted key file + do { + var tmpKeySpec: Data = try decryptSecureExportedKey( + path: encryptedKeyPath, + password: encryptedKeyPassword + ) + tmpKeySpec.resetBytes(in: 0.. String { var keySpec: Data = try getOrGenerateDatabaseKeySpec() defer { keySpec.resetBytes(in: 0.. Data { + let encKeyBase64: String = try String(contentsOf: URL(fileURLWithPath: path), encoding: .utf8) + + guard + var passwordData: Data = password.data(using: .utf8), + var encKeyData: Data = Data(base64Encoded: encKeyBase64) + else { throw StorageError.generic } + defer { + // Reset content immediately after use + passwordData.resetBytes(in: 0.. Void)? @@ -42,19 +67,47 @@ public class DirectoryArchiver { // Stream-based directory traversal and compression let enumerator: FileManager.DirectoryEnumerator? = FileManager.default.enumerator( at: sourceUrl, - includingPropertiesForKeys: [.isRegularFileKey] + includingPropertiesForKeys: [.isRegularFileKey, .isDirectoryKey] ) - let fileUrls: [URL] = (enumerator?.allObjects.compactMap { $0 as? URL } ?? []) - .appending(contentsOf: additionalPaths.map { URL(fileURLWithPath: $0) }) + let fileUrls: [URL] = (enumerator?.allObjects + .compactMap { $0 as? URL } + .filter { url -> Bool in + guard !filenamesToExclude.contains(url.lastPathComponent) else { return false } + guard + let resourceValues = try? url.resourceValues( + forKeys: [.isRegularFileKey, .isDirectoryKey] + ) + else { return true } + + return (resourceValues.isRegularFile == true) + }) + .defaulting(to: []) var index: Int = 0 - progressChanged?(index, fileUrls.count, 0, 0) + progressChanged?(index, (fileUrls.count + additionalPaths.count), 0, 0) + + // Include the archiver version so we can validate compatibility when importing + var version: UInt32 = DirectoryArchiver.version + let versionData: [UInt8] = Array(Data(bytes: &version, count: MemoryLayout.size)) + try write(versionData, to: outputStream, blockSize: UInt8.self, password: password) + + // Store general metadata to help with validation and any other non-file related info + var fileCount: UInt32 = UInt32(fileUrls.count) + var additionalFileCount: UInt32 = UInt32(additionalPaths.count) + let metadata: Data = ( + Data(bytes: &fileCount, count: MemoryLayout.size) + + Data(bytes: &additionalFileCount, count: MemoryLayout.size) + ) + try write(Array(metadata), to: outputStream, blockSize: UInt64.self, password: password) + + // Write the main file content try fileUrls.forEach { url in index += 1 try exportFile( sourcePath: sourcePath, fileURL: url, + customRelativePath: nil, outputStream: outputStream, password: password, index: index, @@ -67,10 +120,12 @@ public class DirectoryArchiver { // Add any extra files which we want to include try additionalPaths.forEach { path in index += 1 - // TODO: Need to fix these so that the path replacement with `sourcePath` isn't busted + + let fileUrl: URL = URL(fileURLWithPath: path) try exportFile( sourcePath: sourcePath, - fileURL: URL(fileURLWithPath: path), + fileURL: fileUrl, + customRelativePath: "_extra/\(fileUrl.lastPathComponent)", outputStream: outputStream, password: password, index: index, @@ -81,103 +136,12 @@ public class DirectoryArchiver { } } - private static func exportFile( - sourcePath: String, - fileURL: URL, - outputStream: OutputStream, - password: String?, - index: Int, - totalFiles: Int, - isExtraFile: Bool, - progressChanged: ((Int, Int, UInt64, UInt64) -> Void)? - ) throws { - guard - let values: URLResourceValues = try? fileURL.resourceValues( - forKeys: [.isRegularFileKey, .fileSizeKey] - ), - values.isRegularFile == true, - var fileSize: UInt64 = values.fileSize.map({ UInt64($0) }) - else { - progressChanged?(index, totalFiles, 1, 1) - return - } - - // Relative path preservation - let relativePath: String = fileURL.path.replacingOccurrences( - of: sourcePath, - with: "" - ).trimmingCharacters(in: CharacterSet(charactersIn: "/")) - - // Write path length and path - let pathData: Data = relativePath.data(using: .utf8)! - var pathLength: UInt32 = UInt32(pathData.count) - var isExtraFile: Bool = isExtraFile - - // Encrypt and write metadata (path length + path data) - let metadata: Data = ( - Data(bytes: &pathLength, count: MemoryLayout.size) + - pathData + - Data(bytes: &fileSize, count: MemoryLayout.size) + - Data(bytes: &isExtraFile, count: MemoryLayout.size) - ) - let processedMetadata: [UInt8] - - switch password { - case .none: processedMetadata = Array(metadata) - case .some(let password): - processedMetadata = try encrypt( - buffer: Array(metadata), - password: password - ) - } - - var blockSize: UInt64 = UInt64(processedMetadata.count) - let blockSizeData: [UInt8] = Array(Data(bytes: &blockSize, count: MemoryLayout.size)) - outputStream.write(blockSizeData, maxLength: blockSizeData.count) - outputStream.write(processedMetadata, maxLength: processedMetadata.count) - - // Stream file contents - guard let inputStream: InputStream = InputStream(url: fileURL) else { - progressChanged?(index, totalFiles, 1, 1) - return - } - - inputStream.open() - defer { inputStream.close() } - - var buffer: [UInt8] = [UInt8](repeating: 0, count: 4096) - var currentFileProcessAmount: UInt64 = 0 - while inputStream.hasBytesAvailable { - let bytesRead: Int = inputStream.read(&buffer, maxLength: buffer.count) - currentFileProcessAmount += UInt64(bytesRead) - progressChanged?(index, totalFiles, currentFileProcessAmount, fileSize) - - if bytesRead > 0 { - let processedBytes: [UInt8] - - switch password { - case .none: processedBytes = buffer - case .some(let password): - processedBytes = try encrypt( - buffer: Array(buffer.prefix(bytesRead)), - password: password - ) - } - - var chunkSize: UInt32 = UInt32(processedBytes.count) - let chunkSizeData: [UInt8] = Array(Data(bytes: &chunkSize, count: MemoryLayout.size)) - outputStream.write(chunkSizeData, maxLength: chunkSizeData.count) - outputStream.write(processedBytes, maxLength: processedBytes.count) - } - } - } - public static func unarchiveDirectory( archivePath: String, destinationPath: String, password: String?, - progressChanged: ((UInt64, UInt64) -> Void)? - ) throws -> [String] { + progressChanged: ((Int, Int, UInt64, UInt64) -> Void)? + ) throws -> (paths: [String], additional: [String]) { // Remove any old imported data as we don't want to muddy the new data if FileManager.default.fileExists(atPath: destinationPath) { try? FileManager.default.removeItem(atPath: destinationPath) @@ -200,38 +164,57 @@ public class DirectoryArchiver { inputStream.open() defer { inputStream.close() } - var extraFilePaths: [String] = [] + // First we need to check the version included in the export is compatible with the current one + let (versionData, _, _): ([UInt8], Int, UInt8) = try read(from: inputStream, password: password) + + guard !versionData.isEmpty else { throw ArchiveError.incompatibleVersion } + + var version: UInt32 = 0 + _ = withUnsafeMutableBytes(of: &version) { versionBuffer in + versionData.copyBytes(to: versionBuffer) + } + + // Retrieve and process the general metadata + var metadataOffset = 0 + let (metadataBytes, _, _): ([UInt8], Int, UInt64) = try read(from: inputStream, password: password) + + guard !metadataBytes.isEmpty else { throw ArchiveError.unarchiveFailed } + + // Extract path length and path + let expectedFileCountRange: Range = metadataOffset..<(metadataOffset + MemoryLayout.size) + var expectedFileCount: UInt32 = 0 + _ = withUnsafeMutableBytes(of: &expectedFileCount) { expectedFileCountBuffer in + metadataBytes.copyBytes(to: expectedFileCountBuffer, from: expectedFileCountRange) + } + metadataOffset += MemoryLayout.size + + let expectedAdditionalFileCountRange: Range = metadataOffset..<(metadataOffset + MemoryLayout.size) + var expectedAdditionalFileCount: UInt32 = 0 + _ = withUnsafeMutableBytes(of: &expectedAdditionalFileCount) { expectedAdditionalFileCountBuffer in + metadataBytes.copyBytes(to: expectedAdditionalFileCountBuffer, from: expectedAdditionalFileCountRange) + } + + var filePaths: [String] = [] + var additionalFilePaths: [String] = [] var fileAmountProcessed: UInt64 = 0 - progressChanged?(0, encryptedFileSize) + progressChanged?(0, Int(expectedFileCount + expectedAdditionalFileCount), 0, encryptedFileSize) while inputStream.hasBytesAvailable { - // Read block size - var blockSizeBytes: [UInt8] = [UInt8](repeating: 0, count: MemoryLayout.size) - let bytesRead: Int = inputStream.read(&blockSizeBytes, maxLength: blockSizeBytes.count) - fileAmountProcessed += UInt64(bytesRead) - progressChanged?(fileAmountProcessed, encryptedFileSize) - - switch bytesRead { - case 0: continue // We have finished reading - case blockSizeBytes.count: break // We have started the next block - default: throw ArchiveError.unarchiveFailed // Invalid - } - - var blockSize: UInt64 = 0 - _ = withUnsafeMutableBytes(of: &blockSize) { blockSizeBuffer in - blockSizeBytes.copyBytes(to: blockSizeBuffer, from: ...size) - } + let (metadata, blockSizeBytesRead, encryptedSize): ([UInt8], Int, UInt64) = try read( + from: inputStream, + password: password + ) + fileAmountProcessed += UInt64(blockSizeBytesRead) + progressChanged?( + (filePaths.count + additionalFilePaths.count), + Int(expectedFileCount + expectedAdditionalFileCount), + fileAmountProcessed, + encryptedFileSize + ) - // Read and decrypt metadata - var encryptedMetadata: [UInt8] = [UInt8](repeating: 0, count: Int(blockSize)) - guard inputStream.read(&encryptedMetadata, maxLength: encryptedMetadata.count) == encryptedMetadata.count else { - throw ArchiveError.unarchiveFailed - } + // Stop here if we have finished reading + guard blockSizeBytesRead > 0 else { continue } - let metadata: [UInt8] - switch password { - case .none: metadata = encryptedMetadata - case .some(let password): metadata = try decrypt(buffer: encryptedMetadata, password: password) - } + // Process the metadata var offset = 0 // Extract path length and path @@ -267,8 +250,13 @@ public class DirectoryArchiver { atPath: (fullPath as NSString).deletingLastPathComponent, withIntermediateDirectories: true ) - fileAmountProcessed += UInt64(encryptedMetadata.count) - progressChanged?(fileAmountProcessed, encryptedFileSize) + fileAmountProcessed += encryptedSize + progressChanged?( + (filePaths.count + additionalFilePaths.count), + Int(expectedFileCount + expectedAdditionalFileCount), + fileAmountProcessed, + encryptedFileSize + ) // Read and decrypt file content guard let outputStream: OutputStream = OutputStream(toFileAtPath: fullPath, append: false) else { @@ -279,43 +267,65 @@ public class DirectoryArchiver { var remainingFileSize: Int = Int(fileSize) while remainingFileSize > 0 { - // Read chunk size - var chunkSizeBytes: [UInt8] = [UInt8](repeating: 0, count: MemoryLayout.size) - guard inputStream.read(&chunkSizeBytes, maxLength: chunkSizeBytes.count) == chunkSizeBytes.count else { - throw ArchiveError.unarchiveFailed - } - var chunkSize: UInt32 = 0 - _ = withUnsafeMutableBytes(of: &chunkSize) { chunkSizeBuffer in - chunkSizeBytes.copyBytes(to: chunkSizeBuffer, from: ...size) - } - fileAmountProcessed += UInt64(chunkSizeBytes.count) - progressChanged?(fileAmountProcessed, encryptedFileSize) - - // Read the chunk - var chunkBytes: [UInt8] = [UInt8](repeating: 0, count: Int(chunkSize)) - guard inputStream.read(&chunkBytes, maxLength: chunkBytes.count) == chunkBytes.count else { - throw ArchiveError.unarchiveFailed - } - - let processedChunk: [UInt8] - switch password { - case .none: processedChunk = chunkBytes - case .some(let password): processedChunk = try decrypt(buffer: chunkBytes, password: password) - } + let (chunk, chunkSizeBytesRead, encryptedSize): ([UInt8], Int, UInt32) = try read( + from: inputStream, + password: password + ) - outputStream.write(processedChunk, maxLength: processedChunk.count) - remainingFileSize -= processedChunk.count + // Write to the output + outputStream.write(chunk, maxLength: chunk.count) + remainingFileSize -= chunk.count - fileAmountProcessed += UInt64(chunkBytes.count) - progressChanged?(fileAmountProcessed, encryptedFileSize) + // Update the progress + fileAmountProcessed += UInt64(chunkSizeBytesRead + chunk.count) + progressChanged?( + (filePaths.count + additionalFilePaths.count), + Int(expectedFileCount + expectedAdditionalFileCount), + fileAmountProcessed, + encryptedFileSize + ) } - if isExtraFile { - extraFilePaths.append(fullPath) + // Store the file path info and update the progress + switch isExtraFile { + case false: filePaths.append(fullPath) + case true: additionalFilePaths.append(fullPath) } + progressChanged?( + (filePaths.count + additionalFilePaths.count), + Int(expectedFileCount + expectedAdditionalFileCount), + fileAmountProcessed, + encryptedFileSize + ) + } + + // Validate that the number of files exported matches the number of paths we got back + let testEnumerator: FileManager.DirectoryEnumerator? = FileManager.default.enumerator( + at: URL(fileURLWithPath: destinationPath), + includingPropertiesForKeys: [.isRegularFileKey, .isDirectoryKey] + ) + let tempFileUrls: [URL] = (testEnumerator?.allObjects + .compactMap { $0 as? URL } + .filter { url -> Bool in + guard + let resourceValues = try? url.resourceValues( + forKeys: [.isRegularFileKey, .isDirectoryKey] + ) + else { return true } + + return (resourceValues.isRegularFile == true) + }) + .defaulting(to: []) + + guard tempFileUrls.count == (filePaths.count + additionalFilePaths.count) else { + throw ArchiveError.importedFileCountMismatch } + guard + filePaths.count == expectedFileCount && + additionalFilePaths.count == expectedAdditionalFileCount + else { throw ArchiveError.importedFileCountMetadataMismatch } - return extraFilePaths + return (filePaths, additionalFilePaths) } private static func encrypt(buffer: [UInt8], password: String) throws -> [UInt8] { @@ -374,6 +384,132 @@ public class DirectoryArchiver { let decryptedData: Data = try AES.GCM.open(sealedBox, using: symmetricKey) return [UInt8](decryptedData) } + + private static func write( + _ data: [UInt8], + to outputStream: OutputStream, + blockSize: T.Type, + password: String? + ) throws where T: FixedWidthInteger, T: UnsignedInteger { + let processedBytes: [UInt8] + + switch password { + case .none: processedBytes = data + case .some(let password): + processedBytes = try encrypt( + buffer: data, + password: password + ) + } + + var blockSize: T = T(processedBytes.count) + let blockSizeData: [UInt8] = Array(Data(bytes: &blockSize, count: MemoryLayout.size)) + outputStream.write(blockSizeData, maxLength: blockSizeData.count) + outputStream.write(processedBytes, maxLength: processedBytes.count) + } + + private static func read( + from inputStream: InputStream, + password: String? + ) throws -> (value: [UInt8], blockSizeBytesRead: Int, encryptedSize: T) where T: FixedWidthInteger, T: UnsignedInteger { + var blockSizeBytes: [UInt8] = [UInt8](repeating: 0, count: MemoryLayout.size) + let bytesRead: Int = inputStream.read(&blockSizeBytes, maxLength: blockSizeBytes.count) + + switch bytesRead { + case 0: return ([], bytesRead, 0) // We have finished reading + case blockSizeBytes.count: break // We have started the next block + default: throw ArchiveError.unarchiveFailed // Invalid + } + + var blockSize: T = 0 + _ = withUnsafeMutableBytes(of: &blockSize) { blockSizeBuffer in + blockSizeBytes.copyBytes(to: blockSizeBuffer, from: ...size) + } + + var encryptedResult: [UInt8] = [UInt8](repeating: 0, count: Int(blockSize)) + guard inputStream.read(&encryptedResult, maxLength: encryptedResult.count) == encryptedResult.count else { + throw ArchiveError.unarchiveFailed + } + + let result: [UInt8] + switch password { + case .none: result = encryptedResult + case .some(let password): result = try decrypt(buffer: encryptedResult, password: password) + } + + return (result, bytesRead, blockSize) + } + + private static func exportFile( + sourcePath: String, + fileURL: URL, + customRelativePath: String?, + outputStream: OutputStream, + password: String?, + index: Int, + totalFiles: Int, + isExtraFile: Bool, + progressChanged: ((Int, Int, UInt64, UInt64) -> Void)? + ) throws { + guard + let values: URLResourceValues = try? fileURL.resourceValues( + forKeys: [.isRegularFileKey, .fileSizeKey] + ), + values.isRegularFile == true, + var fileSize: UInt64 = values.fileSize.map({ UInt64($0) }) + else { + progressChanged?(index, totalFiles, 1, 1) + return + } + + // Relative path preservation + let relativePath: String = customRelativePath + .defaulting( + to: fileURL.path + .replacingOccurrences(of: sourcePath, with: "") + .trimmingCharacters(in: CharacterSet(charactersIn: "/")) + ) + + // Write path length and path + let pathData: Data = relativePath.data(using: .utf8)! + var pathLength: UInt32 = UInt32(pathData.count) + var isExtraFile: Bool = isExtraFile + + // Encrypt and write metadata (path length + path data) + let metadata: Data = ( + Data(bytes: &pathLength, count: MemoryLayout.size) + + pathData + + Data(bytes: &fileSize, count: MemoryLayout.size) + + Data(bytes: &isExtraFile, count: MemoryLayout.size) + ) + try write(Array(metadata), to: outputStream, blockSize: UInt64.self, password: password) + + // Stream file contents + guard let inputStream: InputStream = InputStream(url: fileURL) else { + progressChanged?(index, totalFiles, 1, 1) + return + } + + inputStream.open() + defer { inputStream.close() } + + var buffer: [UInt8] = [UInt8](repeating: 0, count: 4096) + var currentFileProcessAmount: UInt64 = 0 + while inputStream.hasBytesAvailable { + let bytesRead: Int = inputStream.read(&buffer, maxLength: buffer.count) + currentFileProcessAmount += UInt64(bytesRead) + progressChanged?(index, totalFiles, currentFileProcessAmount, fileSize) + + if bytesRead > 0 { + try write( + Array(buffer.prefix(bytesRead)), + to: outputStream, + blockSize: UInt32.self, + password: password + ) + } + } + } } fileprivate extension InputStream {