Fix crash after using device transfer

Author: max-signal

Signal/AppLaunch/AppDelegate.swift

@@ -178,6 +178,13 @@ final class AppDelegate: UIResponder, UIApplicationDelegate {
 
         MessageFetchBGRefreshTask.register()
 
+        let keychainStorage = KeychainStorageImpl(isUsingProductionService: TSConstants.isUsingProductionService)
+        let deviceTransferService = DeviceTransferService(keychainStorage: keychainStorage)
+
+        AppEnvironment.setSharedEnvironment(AppEnvironment(
+            deviceTransferService: deviceTransferService
+        ))
+
         // This *must* happen before we try and access or verify the database,
         // since we may be in a state where the database has been partially
         // restored from transfer (e.g. the key was replaced, but the database
@@ -186,10 +193,9 @@ final class AppDelegate: UIResponder, UIApplicationDelegate {
             title: "Slow device transfer service launch",
             logIfLongerThan: 0.01,
             logInProduction: true,
-            block: { DeviceTransferService.shared.launchCleanup() }
+            block: { deviceTransferService.launchCleanup() }
         )
 
-        let keychainStorage = KeychainStorageImpl(isUsingProductionService: TSConstants.isUsingProductionService)
         let databaseStorage: SDSDatabaseStorage
         do {
             databaseStorage = try SDSDatabaseStorage(

Signal/AppLaunch/AppEnvironment.swift

@@ -7,28 +7,21 @@ import SignalServiceKit
 
 public class AppEnvironment: NSObject {
 
-    private static var _shared: AppEnvironment = AppEnvironment()
+    private static var _shared: AppEnvironment?
 
-    @objc
-    public class var shared: AppEnvironment {
-        get {
-            return _shared
-        }
-        set {
-            guard CurrentAppContext().isRunningTests else {
-                owsFailDebug("Can only switch environments in tests.")
-                return
-            }
-
-            _shared = newValue
-        }
+    static func setSharedEnvironment(_ appEnvironment: AppEnvironment) {
+        owsAssert(self._shared == nil)
+        self._shared = appEnvironment
     }
 
-    public var pushRegistrationManagerRef: PushRegistrationManager
+    @objc
+    public class var shared: AppEnvironment { _shared! }
+
+    let pushRegistrationManagerRef: PushRegistrationManager
 
     var callService: CallService!
 
-    let deviceTransferServiceRef = DeviceTransferService()
+    let deviceTransferServiceRef: DeviceTransferService
 
     let avatarHistorManagerRef = AvatarHistoryManager()
 
@@ -42,7 +35,8 @@ public class AppEnvironment: NSObject {
     private(set) var badgeManager: BadgeManager!
     private var usernameValidationObserverRef: UsernameValidationObserver?
 
-    private override init() {
+    init(deviceTransferService: DeviceTransferService) {
+        self.deviceTransferServiceRef = deviceTransferService
         self.pushRegistrationManagerRef = PushRegistrationManager()
 
         super.init()

Signal/DeviceTransfer/DeviceTransferService+Restore.swift

@@ -145,7 +145,7 @@ extension DeviceTransferService {
         }
 
         do {
-            try databaseStorage.keyFetcher.store(data: database.key)
+            try GRDBKeyFetcher(keychainStorage: keychainStorage).store(data: database.key)
         } catch {
             owsFailDebug("failed to restore database key")
             return false
@@ -547,7 +547,7 @@ extension DeviceTransferService {
             throw OWSAssertionError("No manifest database available")
         }
 
-        try databaseStorage.keyFetcher.store(data: database.key)
+        try GRDBKeyFetcher(keychainStorage: keychainStorage).store(data: database.key)
         GRDBDatabaseStorageAdapter.promoteTransferDirectoryToPrimary()
     }
 

Signal/DeviceTransfer/DeviceTransferService.swift

@@ -123,7 +123,11 @@ class DeviceTransferService: NSObject {
 
     // MARK: -
 
-    override init() {
+    let keychainStorage: any KeychainStorage
+
+    init(keychainStorage: any KeychainStorage) {
+        self.keychainStorage = keychainStorage
+
         super.init()
 
         SwiftSingletons.register(self)

SignalNSE/NSEEnvironment.swift

@@ -131,6 +131,7 @@ class NSEEnvironment: Dependencies {
             databaseFileUrl: SDSDatabaseStorage.grdbDatabaseFileUrl,
             keychainStorage: keychainStorage
         )
+        databaseStorage.grdbStorage.setUpDatabasePathKVO()
 
         didFinishDatabaseSetup = true
 

SignalServiceKit/Storage/Database/GRDBDatabaseStorageAdapter.swift

@@ -87,7 +87,7 @@ public class GRDBDatabaseStorageAdapter: NSObject {
             return primaryFolderName
         }
         if let primaryFolderName = CurrentAppContext().appUserDefaults().string(forKey: DirectoryMode.primaryFolderNameKey) {
-            // Make sure its also written to the file.
+            // Make sure it's also written to the file.
             OWSFileSystem.ensureDirectoryExists(fileUrl.deletingLastPathComponent().path)
             try? primaryFolderName.write(toFile: fileUrl.pathExtension, atomically: true, encoding: .utf8)
             return primaryFolderName
@@ -97,7 +97,7 @@ public class GRDBDatabaseStorageAdapter: NSObject {
 
     fileprivate static func writeStoredPrimaryFolderName(_ newPrimaryFolderName: String) {
         CurrentAppContext().appUserDefaults().set(newPrimaryFolderName, forKey: DirectoryMode.primaryFolderNameKey)
-        // Make sure its also written to the file.
+        // Make sure it's also written to the file.
         let fileUrl = storedPrimaryFolderNameFileUrl()
         OWSFileSystem.ensureDirectoryExists(fileUrl.deletingLastPathComponent().path)
         try? newPrimaryFolderName.write(toFile: fileUrl.pathExtension, atomically: true, encoding: .utf8)
@@ -139,9 +139,6 @@ public class GRDBDatabaseStorageAdapter: NSObject {
         } catch {
             throw error
         }
-
-        super.init()
-        setUpDatabasePathKVO()
     }
 
     deinit {
@@ -207,7 +204,7 @@ public class GRDBDatabaseStorageAdapter: NSObject {
 
     private var darwinToken: Int32?
 
-    func setUpDatabasePathKVO() {
+    public func setUpDatabasePathKVO() {
         darwinToken = DarwinNotificationCenter.addObserver(
             for: .primaryDBFolderNameDidChange,
             queue: .main,
@@ -351,7 +348,7 @@ extension GRDBDatabaseStorageAdapter {
     }
 
     public static func promoteTransferDirectoryToPrimary() {
-        owsAssert(CurrentAppContext().isMainApp, "Only the main app can't swap databases")
+        owsAssert(CurrentAppContext().isMainApp, "Only the main app can swap databases")
 
         // Ordering matters here. We should be able to crash and recover without issue
         // A prior run may have already performed the swap but crashed, so we should not expect a transfer folder
@@ -863,7 +860,11 @@ public struct GRDBKeyFetcher {
         public static let kSQLCipherKeySpecLength: Int32 = 48
     }
 
-    let keychainStorage: any KeychainStorage
+    private let keychainStorage: any KeychainStorage
+
+    public init(keychainStorage: any KeychainStorage) {
+        self.keychainStorage = keychainStorage
+    }
 
     func fetchString() throws -> String {
         // Use a raw key spec, where the 96 hexadecimal digits are provided

SignalShareExtension/ShareViewController.swift

@@ -79,6 +79,7 @@ public class ShareViewController: UIViewController, ShareViewDelegate, SAEFailed
             self.showNotRegisteredView()
             return
         }
+        databaseStorage.grdbStorage.setUpDatabasePathKVO()
 
         // We shouldn't set up our environment until after we've consulted isReadyForAppExtensions.
         let databaseContinuation = AppSetup().start(