From cef72e5a5a9cdfeb529e8828d145b1078cca498c Mon Sep 17 00:00:00 2001 From: Max Radermacher Date: Mon, 8 Jun 2026 14:17:53 -0500 Subject: [PATCH] Wait for sync message after storage service error --- .../Account/AccountKeyStore.swift | 18 ++++++ .../SecureValueRecovery2Impl.swift | 25 +++++--- .../StorageServiceManager.swift | 57 +++++++++++-------- 3 files changed, 70 insertions(+), 30 deletions(-) diff --git a/SignalServiceKit/Account/AccountKeyStore.swift b/SignalServiceKit/Account/AccountKeyStore.swift index 6b1753fd19..94b339b3f8 100644 --- a/SignalServiceKit/Account/AccountKeyStore.swift +++ b/SignalServiceKit/Account/AccountKeyStore.swift @@ -20,6 +20,7 @@ public class AccountKeyStore { private let aepKvStore: KeyValueStore private let mrbkKvStore: NewKeyValueStore private let masterKeyKvStore: NewKeyValueStore + private let syncStore: NewKeyValueStore private let backupSettingsStore: BackupSettingsStore @@ -30,6 +31,7 @@ public class AccountKeyStore { self.masterKeyKvStore = NewKeyValueStore(collection: "kOWSKeyBackupService_Keys") self.mrbkKvStore = NewKeyValueStore(collection: "MediaRootBackupKey") self.aepKvStore = KeyValueStore(collection: "AccountEntropyPool") + self.syncStore = NewKeyValueStore(collection: "AccountKey.Sync") self.backupSettingsStore = backupSettingsStore } @@ -137,4 +139,20 @@ public class AccountKeyStore { aepKvStore.setString(accountEntropyPool.rawString, key: Keys.aepKeyName, transaction: tx) } + + // MARK: - + + private static let isWaitingForKeysSyncKey = "isWaitingForKeysSync" + + func isWaitingForKeysSyncMessage(tx: DBReadTransaction) -> Bool { + return syncStore.fetchValue(Bool.self, forKey: Self.isWaitingForKeysSyncKey, tx: tx) == true + } + + func setWaitingForKeysSyncMessage(_ isWaitingForKeysSyncMessage: Bool, tx: DBWriteTransaction) { + if isWaitingForKeysSyncMessage { + syncStore.writeValue(true, forKey: Self.isWaitingForKeysSyncKey, tx: tx) + } else { + syncStore.removeValue(forKey: Self.isWaitingForKeysSyncKey, tx: tx) + } + } } diff --git a/SignalServiceKit/SecureValueRecovery/SecureValueRecovery2Impl.swift b/SignalServiceKit/SecureValueRecovery/SecureValueRecovery2Impl.swift index 45ef33c894..25efeb1cb4 100644 --- a/SignalServiceKit/SecureValueRecovery/SecureValueRecovery2Impl.swift +++ b/SignalServiceKit/SecureValueRecovery/SecureValueRecovery2Impl.swift @@ -102,6 +102,7 @@ public class SecureValueRecovery2Impl: SecureValueRecovery { Logger.info("") accountKeyStore.setMediaRootBackupKey(provisioningMessage.mrbk, tx: tx) accountKeyStore.setAccountEntropyPool(provisioningMessage.aep, tx: tx) + accountKeyStore.setWaitingForKeysSyncMessage(false, tx: tx) } public func storeKeys( @@ -121,14 +122,24 @@ public class SecureValueRecovery2Impl: SecureValueRecovery { guard let newAep else { throw SVR.KeysError.missingAep } - let oldAep = accountKeyStore.getAccountEntropyPool(tx: tx) - if newAep != oldAep { + var shouldRestore = false + if newAep != accountKeyStore.getAccountEntropyPool(tx: tx) { accountKeyStore.setAccountEntropyPool(newAep, tx: tx) - // Trigger a re-fetch of the storage manifest if our keys have changed - storageServiceManager.restoreOrCreateManifestIfNecessary( - authedDevice: authedDevice, - masterKeySource: .implicit, - ) + shouldRestore = true + } + if accountKeyStore.isWaitingForKeysSyncMessage(tx: tx) { + accountKeyStore.setWaitingForKeysSyncMessage(false, tx: tx) + shouldRestore = true + } + if shouldRestore { + // Trigger a re-fetch of the storage manifest if our keys have changed or + // if we've gotten a key that we requested. + tx.addSyncCompletion { [storageServiceManager] in + storageServiceManager.restoreOrCreateManifestIfNecessary( + authedDevice: authedDevice, + masterKeySource: .implicit, + ) + } } } diff --git a/SignalServiceKit/StorageService/StorageServiceManager.swift b/SignalServiceKit/StorageService/StorageServiceManager.swift index 189caaf535..1fc5fda810 100644 --- a/SignalServiceKit/StorageService/StorageServiceManager.swift +++ b/SignalServiceKit/StorageService/StorageServiceManager.swift @@ -775,7 +775,9 @@ class StorageServiceOperation { } private func _run() async throws { + let accountKeyStore = DependenciesBridge.shared.accountKeyStore let databaseStorage = SSKEnvironment.shared.databaseStorageRef + let tsAccountManager = DependenciesBridge.shared.tsAccountManager let (currentStateIfRotatingManifest, masterKey) = databaseStorage.read { tx in let state: State? @@ -786,12 +788,17 @@ class StorageServiceOperation { state = nil } - let masterKey: MasterKey? + var masterKey: MasterKey? switch masterKeySource { case .explicit(let keyData): masterKey = keyData case .implicit: - masterKey = DependenciesBridge.shared.accountKeyStore.getMasterKey(tx: tx) + masterKey = accountKeyStore.getMasterKey(tx: tx) + } + + if !isPrimaryDevice, accountKeyStore.isWaitingForKeysSyncMessage(tx: tx) { + // We hit a failure and are waiting for a "new" AEP/MasterKey. + masterKey = nil } return (state, masterKey) @@ -800,13 +807,9 @@ class StorageServiceOperation { guard let masterKey else { if !isPrimaryDevice, - DependenciesBridge.shared.tsAccountManager.registrationStateWithMaybeSneakyTransaction.isRegistered + tsAccountManager.registrationStateWithMaybeSneakyTransaction.isRegistered { - // This is a linked device, and keys are missing. There's nothing that can be done - // until we receive new keys, so send a key sync message and return early. - await databaseStorage.awaitableWrite { tx in - SSKEnvironment.shared.syncManagerRef.sendKeysSyncRequestMessage(transaction: tx) - } + await sendKeysSyncRequestMessageIfNeeded() } else { // We're either not registered, or a primary. Either way, // we don't have keys, or a means to get them, so do nothing. @@ -847,6 +850,25 @@ class StorageServiceOperation { } } + private func sendKeysSyncRequestMessageIfNeeded() async { + owsPrecondition(!isPrimaryDevice) + let accountKeyStore = DependenciesBridge.shared.accountKeyStore + let databaseStorage = SSKEnvironment.shared.databaseStorageRef + let syncManager = SSKEnvironment.shared.syncManagerRef + + await databaseStorage.awaitableWrite { tx in + if accountKeyStore.isWaitingForKeysSyncMessage(tx: tx) { + // We've already requested keys; if the request got lost, we can rely on + // the periodic keys sync message. + return + } + // This is a linked device, and keys are missing. There's nothing that can + // be done until we receive new keys, so send a key sync message. + syncManager.sendKeysSyncRequestMessage(transaction: tx) + accountKeyStore.setWaitingForKeysSyncMessage(true, tx: tx) + } + } + // MARK: - Mark Pending Changes fileprivate static func recordPendingMutations(_ pendingMutations: PendingMutations) -> (() async -> Void) { @@ -1212,14 +1234,8 @@ class StorageServiceOperation { case StorageService.StorageError.manifestDecryptionFailed(_) where !isPrimaryDevice: // If this is a linked device, give up and request the latest storage // service key from the primary device. - Logger.warn("Manifest decryption failed on linked device, clearing storage service keys.") - - await SSKEnvironment.shared.databaseStorageRef.awaitableWrite { transaction in - // Clear out the key, it's no longer valid. This will prevent us - // from trying to backup again until the sync response is received. - DependenciesBridge.shared.accountKeyStore.setMasterKey(nil, tx: transaction) - SSKEnvironment.shared.syncManagerRef.sendKeysSyncRequestMessage(transaction: transaction) - } + Logger.warn("Manifest decryption failed on linked device; waiting for keys sync") + await sendKeysSyncRequestMessageIfNeeded() default: break } @@ -1762,16 +1778,11 @@ class StorageServiceOperation { return } - Logger.warn("Item decryption failed, clearing storage service keys.") + Logger.warn("Item decryption failed; waiting for keys sync") // If this is a linked device, give up and request the latest storage // service key from the primary device. - await SSKEnvironment.shared.databaseStorageRef.awaitableWrite { transaction in - // Clear out the key, it's no longer valid. This will prevent us - // from trying to backup again until the sync response is received. - DependenciesBridge.shared.accountKeyStore.setMasterKey(nil, tx: transaction) - SSKEnvironment.shared.syncManagerRef.sendKeysSyncRequestMessage(transaction: transaction) - } + await sendKeysSyncRequestMessageIfNeeded() } else if case .itemProtoDeserializationFailed = storageError, self.isPrimaryDevice