Cleaned up and finished the export/import logic

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
+                            self?.performImport()
-                                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)
                         }
                     }
                 )
@@ -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 {

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)
-    }
-}

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..<tmpKeySpec.count)
+        }
+        catch { return }
+        
+        /// Then configure the database using the key
+        var config = Configuration()
+
+        /// Load in the SQLCipher keys
+        config.prepareDatabase { [weak self] db in
+            var keySpec: Data = try self?.decryptSecureExportedKey(
+                path: encryptedKeyPath,
+                password: encryptedKeyPassword
+            ) ?? { throw StorageError.invalidKeySpec }()
+            defer { keySpec.resetBytes(in: 0..<keySpec.count) } // Reset content immediately after use
+            
+            // Use a raw key spec, where the 96 hexadecimal digits are provided
+            // (i.e. 64 hex for the 256 bit key, followed by 32 hex for the 128 bit salt)
+            // using explicit BLOB syntax, e.g.:
+            //
+            // x'98483C6EB40B6C31A448C22A66DED3B5E5E8D5119CAC8327B655C8B5C483648101010101010101010101010101010101'
+            keySpec = try (keySpec.toHexString().data(using: .utf8) ?? { throw StorageError.invalidKeySpec }())
+            keySpec.insert(contentsOf: [120, 39], at: 0)    // "x'" prefix
+            keySpec.append(39)                              // "'" suffix
+            
+            try db.usePassphrase(keySpec)
+            
+            // According to the SQLCipher docs iOS needs the 'cipher_plaintext_header_size' value set to at least
+            // 32 as iOS extends special privileges to the database and needs this header to be in plaintext
+            // to determine the file type
+            //
+            // For more info see: https://www.zetetic.net/sqlcipher/sqlcipher-api/#cipher_plaintext_header_size
+            try db.execute(sql: "PRAGMA cipher_plaintext_header_size = 32")
+        }
+        
+        // Create the DatabasePool to allow us to connect to the database and mark the storage as valid
+        dbWriter = try DatabasePool(path: databasePath, configuration: config)
+        isValid = true
+    }
+    
     func secureExportKey(password: String) throws -> String {
         var keySpec: Data = try getOrGenerateDatabaseKeySpec()
         defer { keySpec.resetBytes(in: 0..<keySpec.count) } // Reset content immediately after use
@@ -753,14 +827,42 @@ public extension Storage {
         let key: SymmetricKey = SymmetricKey(data: Data(hash.makeIterator()))
         let sealedBox: ChaChaPoly.SealedBox = try ChaChaPoly.seal(keySpec, using: key, nonce: nonce, authenticating: Data())
         let keyInfoPath: String = (Singleton.hasAppContext ?
-            "\(Singleton.appContext.temporaryDirectory)/key.enc" :
+            "\(Singleton.appContext.temporaryDirectory)/\(Storage.encKeyFilename)" :
-            "\(NSTemporaryDirectory())key.enc"
+            "\(NSTemporaryDirectory())\(Storage.encKeyFilename)"
         )
         let encryptedKeyBase64: String = sealedBox.combined.base64EncodedString()
         try encryptedKeyBase64.write(toFile: keyInfoPath, atomically: true, encoding: .utf8)
         
         return keyInfoPath
     }
+    
+    func replaceDatabaseKey(path: String, password: String) throws {
+        var keySpec: Data = try decryptSecureExportedKey(path: path, password: password)
+        defer { keySpec.resetBytes(in: 0..<keySpec.count) } // Reset content immediately after use
+        
+        try Singleton.keychain.set(data: keySpec, forKey: .dbCipherKeySpec)
+    }
+    
+    fileprivate func decryptSecureExportedKey(path: String, password: String) throws -> 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..<passwordData.count)
+            encKeyData.resetBytes(in: 0..<encKeyData.count)
+        }
+        
+        let hash: SHA256.Digest = SHA256.hash(data: passwordData)
+        let key: SymmetricKey = SymmetricKey(data: Data(hash.makeIterator()))
+        
+        let sealedBox: ChaChaPoly.SealedBox = try ChaChaPoly.SealedBox(combined: encKeyData)
+        
+        return try ChaChaPoly.open(sealedBox, using: key, authenticating: Data())
+    }
 }
 
 // MARK: - CallInfo

SessionUtilitiesKit/Utilities/DirectoryArchiver.swift

@@ -1,16 +1,40 @@
 // Copyright © 2024 Rangeproof Pty Ltd. All rights reserved.
+//
+// stringlint:disable
 
 import Foundation
 import CryptoKit
 import Compression
 
-enum ArchiveError: Error {
+public enum ArchiveError: Error, CustomStringConvertible {
     case invalidSourcePath
     case archiveFailed
     case unarchiveFailed
+    case decryptionFailed
+    case incompatibleVersion
+    case unableToFindDatabaseKey
+    case importedFileCountMismatch
+    case importedFileCountMetadataMismatch
+    
+    public var description: String {
+        switch self {
+            case .invalidSourcePath: "Invalid source path provided."
+            case .archiveFailed: "Failed to archive."
+            case .unarchiveFailed: "Failed to unarchive."
+            case .decryptionFailed: "Decryption failed."
+            case .incompatibleVersion: "This exported bundle is not compatible with this version of Session."
+            case .unableToFindDatabaseKey: "Unable to find database key."
+            case .importedFileCountMismatch: "The number of files imported doesn't match the number of files written to disk."
+            case .importedFileCountMetadataMismatch: "The number of files imported doesn't match the number of files reported."
+        }
+    }
 }
 
 public class DirectoryArchiver {
+    /// This value is here in case we need to change the structure of the exported data in the future, this would allow us to have
+    /// some form of backwards compatibility if desired
+    private static let version: UInt32 = 1
+    
     /// Archive an entire directory
     /// - Parameters:
     ///   - sourcePath: Full path to the directory to compress
@@ -20,6 +44,7 @@ public class DirectoryArchiver {
     public static func archiveDirectory(
         sourcePath: String,
         destinationPath: String,
+        filenamesToExclude: [String] = [],
         additionalPaths: [String] = [],
         password: String?,
         progressChanged: ((Int, Int, UInt64, UInt64) -> 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 } ?? [])
+        let fileUrls: [URL] = (enumerator?.allObjects
-            .appending(contentsOf: additionalPaths.map { URL(fileURLWithPath: $0) })
+            .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<UInt32>.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<UInt32>.size) +
+            Data(bytes: &additionalFileCount, count: MemoryLayout<UInt32>.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<UInt32>.size) +
-            pathData +
-            Data(bytes: &fileSize, count: MemoryLayout<UInt64>.size) +
-            Data(bytes: &isExtraFile, count: MemoryLayout<Bool>.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<UInt64>.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<UInt32>.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)?
+        progressChanged: ((Int, Int, UInt64, UInt64) -> Void)?
-    ) throws -> [String] {
+    ) 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<Int> = metadataOffset..<(metadataOffset + MemoryLayout<UInt32>.size)
+        var expectedFileCount: UInt32 = 0
+        _ = withUnsafeMutableBytes(of: &expectedFileCount) { expectedFileCountBuffer in
+            metadataBytes.copyBytes(to: expectedFileCountBuffer, from: expectedFileCountRange)
+        }
+        metadataOffset += MemoryLayout<UInt32>.size
+        
+        let expectedAdditionalFileCountRange: Range<Int> = metadataOffset..<(metadataOffset + MemoryLayout<UInt32>.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
+            let (metadata, blockSizeBytesRead, encryptedSize): ([UInt8], Int, UInt64) = try read(
-            var blockSizeBytes: [UInt8] = [UInt8](repeating: 0, count: MemoryLayout<UInt64>.size)
+                from: inputStream,
-            let bytesRead: Int = inputStream.read(&blockSizeBytes, maxLength: blockSizeBytes.count)
+                password: password
-            fileAmountProcessed += UInt64(bytesRead)
+            )
-            progressChanged?(fileAmountProcessed, encryptedFileSize)
+            fileAmountProcessed += UInt64(blockSizeBytesRead)
-            
+            progressChanged?(
-            switch bytesRead {
+                (filePaths.count + additionalFilePaths.count),
-                case 0: continue                            // We have finished reading
+                Int(expectedFileCount + expectedAdditionalFileCount),
-                case blockSizeBytes.count: break            // We have started the next block
+                fileAmountProcessed,
-                default: throw ArchiveError.unarchiveFailed // Invalid
+                encryptedFileSize
-            }
+            )
-            
-            var blockSize: UInt64 = 0
-            _ = withUnsafeMutableBytes(of: &blockSize) { blockSizeBuffer in
-                blockSizeBytes.copyBytes(to: blockSizeBuffer, from: ..<MemoryLayout<UInt64>.size)
-            }
             
-            // Read and decrypt metadata
+            // Stop here if we have finished reading
-            var encryptedMetadata: [UInt8] = [UInt8](repeating: 0, count: Int(blockSize))
+            guard blockSizeBytesRead > 0 else { continue }
-            guard inputStream.read(&encryptedMetadata, maxLength: encryptedMetadata.count) == encryptedMetadata.count else {
-                throw ArchiveError.unarchiveFailed
-            }
             
-            let metadata: [UInt8]
+            // Process the metadata
-            switch password {
-                case .none: metadata = encryptedMetadata
-                case .some(let password): metadata = try decrypt(buffer: encryptedMetadata, password: password)
-            }
             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)
+            fileAmountProcessed += encryptedSize
-            progressChanged?(fileAmountProcessed, encryptedFileSize)
+            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
+                let (chunk, chunkSizeBytesRead, encryptedSize): ([UInt8], Int, UInt32) = try read(
-                var chunkSizeBytes: [UInt8] = [UInt8](repeating: 0, count: MemoryLayout<UInt32>.size)
+                    from: inputStream,
-                guard inputStream.read(&chunkSizeBytes, maxLength: chunkSizeBytes.count) == chunkSizeBytes.count else {
+                    password: password
-                    throw ArchiveError.unarchiveFailed
+                )
-                }
-                var chunkSize: UInt32 = 0
-                _ = withUnsafeMutableBytes(of: &chunkSize) { chunkSizeBuffer in
-                    chunkSizeBytes.copyBytes(to: chunkSizeBuffer, from: ..<MemoryLayout<UInt32>.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)
-                }
                 
-                outputStream.write(processedChunk, maxLength: processedChunk.count)
+                // Write to the output
-                remainingFileSize -= processedChunk.count
+                outputStream.write(chunk, maxLength: chunk.count)
+                remainingFileSize -= chunk.count
                 
-                fileAmountProcessed += UInt64(chunkBytes.count)
+                // Update the progress
-                progressChanged?(fileAmountProcessed, encryptedFileSize)
+                fileAmountProcessed += UInt64(chunkSizeBytesRead + chunk.count)
+                progressChanged?(
+                    (filePaths.count + additionalFilePaths.count),
+                    Int(expectedFileCount + expectedAdditionalFileCount),
+                    fileAmountProcessed,
+                    encryptedFileSize
+                )
             }
             
-            if isExtraFile {
+            // Store the file path info and update the progress
-                extraFilePaths.append(fullPath)
+            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<T>(
+        _ 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<T>.size))
+        outputStream.write(blockSizeData, maxLength: blockSizeData.count)
+        outputStream.write(processedBytes, maxLength: processedBytes.count)
+    }
+    
+    private static func read<T>(
+        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<T>.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: ..<MemoryLayout<T>.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<UInt32>.size) +
+            pathData +
+            Data(bytes: &fileSize, count: MemoryLayout<UInt64>.size) +
+            Data(bytes: &isExtraFile, count: MemoryLayout<Bool>.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 {