signalapp/Signal-iOS
signalapp/Signal-iOS
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases 4 Wiki Activity

Use ConcurrentTaskQueue in PreKeyManager

Browse Source
This commit is contained in:
Max Radermacher 2026-01-22 15:09:10 -06:00 committed by GitHub
parent 3b5badc5cf
commit 7ab18d8133
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
12 changed files with 131 additions and 207 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
Signal/AppLaunch/AppDelegate.swift
View File

@ -1337,9 +1337,9 @@ final class AppDelegate: UIResponder, UIApplicationDelegate {
let registeredState = try? tsAccountManager.registeredStateWithMaybeSneakyTransaction()
if registeredState != nil {
DependenciesBridge.shared.db.read { tx in
Task {
// Always check prekeys after app launches, and sometimes check on app activation.
DependenciesBridge.shared.preKeyManager.checkPreKeysIfNecessary(tx: tx)
try? await DependenciesBridge.shared.preKeyManager.checkPreKeysIfNecessary()
}
}

29
Signal/Provisioning/ProvisioningCoordinatorImpl.swift
View File

@ -238,23 +238,10 @@ class ProvisioningCoordinatorImpl: ProvisioningCoordinator {
pni: Pni,
phoneNumber: E164,
) async throws(CompleteProvisioningError) -> CompleteProvisioningStepResult {
let prekeyBundles: RegistrationPreKeyUploadBundles
do {
// This should be the last failable thing we do before making the verification
// request, because if the verification request fails we need to clean up prekey
// state created by this method.
// If we did add new (failable) method calls between this and the verification
// request invocation, we would have to make sure we similarly clean up prekey
// state if there are failures.
prekeyBundles = try await self.preKeyManager
.createPreKeysForProvisioning(
aciIdentityKeyPair: provisionMessage.aciIdentityKeyPair.asECKeyPair,
pniIdentityKeyPair: provisionMessage.pniIdentityKeyPair.asECKeyPair,
)
.value
} catch {
throw .genericError(error)
}
let prekeyBundles = await self.preKeyManager.createPreKeysForProvisioning(
aciIdentityKeyPair: provisionMessage.aciIdentityKeyPair.asECKeyPair,
pniIdentityKeyPair: provisionMessage.pniIdentityKeyPair.asECKeyPair,
)
return try await completeProvisioning_createRegistrationIds(
provisionMessage: provisionMessage,
@ -264,10 +251,10 @@ class ProvisioningCoordinatorImpl: ProvisioningCoordinator {
phoneNumber: phoneNumber,
prekeyBundles: prekeyBundles,
).withUndoOnFailureStep {
try await self.preKeyManager.finalizeRegistrationPreKeys(
await self.preKeyManager.finalizeRegistrationPreKeys(
prekeyBundles,
uploadDidSucceed: false,
).value
)
}
}
@ -434,12 +421,10 @@ class ProvisioningCoordinatorImpl: ProvisioningCoordinator {
authedDevice: AuthedDevice.Explicit,
) async throws(CompleteProvisioningError) -> CompleteProvisioningStepResult {
do {
try await self.preKeyManager
await self.preKeyManager
.finalizeRegistrationPreKeys(prekeyBundles, uploadDidSucceed: true)
.value
try await self.preKeyManager
.rotateOneTimePreKeysForRegistration(auth: authedDevice.authedAccount.chatServiceAuth)
.value
} catch {
throw .genericError(error)
}

22
Signal/Registration/RegistrationCoordinatorImpl.swift
View File

@ -3483,7 +3483,7 @@ public class RegistrationCoordinatorImpl: RegistrationCoordinator {
// But we should still upload one-time prekeys, as that is not part
// of account creation.
do {
try await deps.preKeyManager.rotateOneTimePreKeysForRegistration(auth: accountIdentity.chatServiceAuth).value
try await deps.preKeyManager.rotateOneTimePreKeysForRegistration(auth: accountIdentity.chatServiceAuth)
self.db.write { tx in
self.updatePersistedState(tx) {
// No harm marking both down as done even though
@ -4494,12 +4494,7 @@ public class RegistrationCoordinatorImpl: RegistrationCoordinator {
persistRegistrationMessage(registrationMessage)
}
let prekeyBundles: RegistrationPreKeyUploadBundles
do {
prekeyBundles = try await deps.preKeyManager.createPreKeysForRegistration().value
} catch {
return .showErrorSheet(.genericError)
}
let prekeyBundles = await deps.preKeyManager.createPreKeysForRegistration()
let shouldSkipDeviceTransfer = self.shouldSkipDeviceTransfer()
let signalService = self.deps.signalService
@ -4523,15 +4518,10 @@ public class RegistrationCoordinatorImpl: RegistrationCoordinator {
.genericError,
.deviceTransferPossible: false
}
do {
try await deps.preKeyManager.finalizeRegistrationPreKeys(
prekeyBundles,
uploadDidSucceed: isPrekeyUploadSuccess,
).value
} catch {
// Finalizing is best effort.
Logger.error("Unable to finalize prekeys, ignoring and continuing")
}
await deps.preKeyManager.finalizeRegistrationPreKeys(
prekeyBundles,
uploadDidSucceed: isPrekeyUploadSuccess,
)
return await responseHandler(accountResponse)
}

24
Signal/test/Registration/RegistrationCoordinatorTestShims.swift
View File

@ -168,34 +168,34 @@ public class _RegistrationCoordinator_PreKeyManagerMock: PreKeyManager {
}
public func isAppLockedDueToPreKeyUpdateFailures(tx: DBReadTransaction) -> Bool { fatalError() }
public func checkPreKeysIfNecessary(tx: DBReadTransaction) { fatalError() }
public func checkPreKeysIfNecessary() async throws { fatalError() }
public func rotatePreKeysOnUpgradeIfNecessary(for identity: OWSIdentity) async throws { fatalError() }
public func createPreKeysForProvisioning(aciIdentityKeyPair: ECKeyPair, pniIdentityKeyPair: ECKeyPair) -> Task<RegistrationPreKeyUploadBundles, any Error> { fatalError() }
public func rotateSignedPreKeysIfNeeded() -> Task<Void, any Error> { fatalError() }
public func refreshOneTimePreKeys(forIdentity identity: OWSIdentity, alsoRefreshSignedPreKey shouldRefreshSignedPreKey: Bool) { fatalError() }
public func createPreKeysForProvisioning(aciIdentityKeyPair: ECKeyPair, pniIdentityKeyPair: ECKeyPair) async -> RegistrationPreKeyUploadBundles { fatalError() }
public func rotateSignedPreKeysIfNeeded() async throws { fatalError() }
public func refreshOneTimePreKeys(forIdentity identity: OWSIdentity, alsoRefreshSignedPreKey shouldRefreshSignedPreKey: Bool) async throws { fatalError() }
public typealias CreatePreKeysMock = () -> Task<RegistrationPreKeyUploadBundles, any Error>
public typealias CreatePreKeysMock = () -> Task<RegistrationPreKeyUploadBundles, Never>
private var createPreKeysMocks = [CreatePreKeysMock]()
public func addCreatePreKeysMock(_ mock: @escaping CreatePreKeysMock) { createPreKeysMocks.append(mock) }
public func createPreKeysForRegistration() -> Task<RegistrationPreKeyUploadBundles, any Error> {
public func createPreKeysForRegistration() async -> RegistrationPreKeyUploadBundles {
run.addObservedStep(.createPreKeys)
return createPreKeysMocks.removeFirst()()
return await createPreKeysMocks.removeFirst()().value
}
public typealias FinalizePreKeysMock = (Bool) -> Task<Void, any Error>
public typealias FinalizePreKeysMock = (Bool) -> Task<Void, Never>
private var finalizePreKeysMocks = [FinalizePreKeysMock]()
public func addFinalizePreKeyMock(_ mock: @escaping FinalizePreKeysMock) { finalizePreKeysMocks.append(mock) }
public func finalizeRegistrationPreKeys(_ bundles: RegistrationPreKeyUploadBundles, uploadDidSucceed: Bool) -> Task<Void, any Error> {
public func finalizeRegistrationPreKeys(_ bundles: RegistrationPreKeyUploadBundles, uploadDidSucceed: Bool) async {
run.addObservedStep(.finalizePreKeys)
return finalizePreKeysMocks.removeFirst()(uploadDidSucceed)
await finalizePreKeysMocks.removeFirst()(uploadDidSucceed).value
}
public typealias RotateOneTimePreKeysMock = (ChatServiceAuth) -> Task<Void, any Error>
private var rotateOneTimePreKeysMocks = [RotateOneTimePreKeysMock]()
public func addRotateOneTimePreKeyMock(_ mock: @escaping RotateOneTimePreKeysMock) { rotateOneTimePreKeysMocks.append(mock) }
public func rotateOneTimePreKeysForRegistration(auth: ChatServiceAuth) -> Task<Void, any Error> {
public func rotateOneTimePreKeysForRegistration(auth: ChatServiceAuth) async throws {
run.addObservedStep(.rotateOneTimePreKeys)
return rotateOneTimePreKeysMocks.removeFirst()(auth)
return try await rotateOneTimePreKeysMocks.removeFirst()(auth).value
}
public func setIsChangingNumber(_ isChangingNumber: Bool) {

77
SignalServiceKit/Account/PreKeys/MockPreKeyManager.swift
View File

@ -11,51 +11,47 @@ import LibSignalClient
class MockPreKeyManager: PreKeyManager {
func isAppLockedDueToPreKeyUpdateFailures(tx: SignalServiceKit.DBReadTransaction) -> Bool { false }
func refreshOneTimePreKeysCheckDidSucceed() { }
func checkPreKeysIfNecessary(tx: SignalServiceKit.DBReadTransaction) { }
func checkPreKeysIfNecessary() async throws { }
func rotatePreKeysOnUpgradeIfNecessary(for identity: OWSIdentity) async throws { }
var attemptedRefreshes: [(OWSIdentity, Bool)] = []
func createPreKeysForRegistration() -> Task<RegistrationPreKeyUploadBundles, Error> {
func createPreKeysForRegistration() async -> RegistrationPreKeyUploadBundles {
let identityKeyPair = ECKeyPair.generateKeyPair()
return Task {
.init(
aci: .init(
identity: .aci,
identityKeyPair: identityKeyPair,
signedPreKey: SignedPreKeyStoreImpl.generateSignedPreKey(keyId: PreKeyId.random(), signedBy: identityKeyPair.keyPair.privateKey),
lastResortPreKey: generateLastResortKyberPreKey(signedBy: identityKeyPair.keyPair.privateKey),
),
pni: .init(
identity: .pni,
identityKeyPair: identityKeyPair,
signedPreKey: SignedPreKeyStoreImpl.generateSignedPreKey(keyId: PreKeyId.random(), signedBy: identityKeyPair.keyPair.privateKey),
lastResortPreKey: generateLastResortKyberPreKey(signedBy: identityKeyPair.keyPair.privateKey),
),
)
}
return .init(
aci: .init(
identity: .aci,
identityKeyPair: identityKeyPair,
signedPreKey: SignedPreKeyStoreImpl.generateSignedPreKey(keyId: PreKeyId.random(), signedBy: identityKeyPair.keyPair.privateKey),
lastResortPreKey: generateLastResortKyberPreKey(signedBy: identityKeyPair.keyPair.privateKey),
),
pni: .init(
identity: .pni,
identityKeyPair: identityKeyPair,
signedPreKey: SignedPreKeyStoreImpl.generateSignedPreKey(keyId: PreKeyId.random(), signedBy: identityKeyPair.keyPair.privateKey),
lastResortPreKey: generateLastResortKyberPreKey(signedBy: identityKeyPair.keyPair.privateKey),
),
)
}
func createPreKeysForProvisioning(
aciIdentityKeyPair: ECKeyPair,
pniIdentityKeyPair: ECKeyPair,
) -> Task<RegistrationPreKeyUploadBundles, Error> {
) async -> RegistrationPreKeyUploadBundles {
let identityKeyPair = ECKeyPair.generateKeyPair()
return Task {
.init(
aci: .init(
identity: .aci,
identityKeyPair: identityKeyPair,
signedPreKey: SignedPreKeyStoreImpl.generateSignedPreKey(keyId: PreKeyId.random(), signedBy: identityKeyPair.keyPair.privateKey),
lastResortPreKey: generateLastResortKyberPreKey(signedBy: identityKeyPair.keyPair.privateKey),
),
pni: .init(
identity: .pni,
identityKeyPair: identityKeyPair,
signedPreKey: SignedPreKeyStoreImpl.generateSignedPreKey(keyId: PreKeyId.random(), signedBy: identityKeyPair.keyPair.privateKey),
lastResortPreKey: generateLastResortKyberPreKey(signedBy: identityKeyPair.keyPair.privateKey),
),
)
}
return .init(
aci: .init(
identity: .aci,
identityKeyPair: identityKeyPair,
signedPreKey: SignedPreKeyStoreImpl.generateSignedPreKey(keyId: PreKeyId.random(), signedBy: identityKeyPair.keyPair.privateKey),
lastResortPreKey: generateLastResortKyberPreKey(signedBy: identityKeyPair.keyPair.privateKey),
),
pni: .init(
identity: .pni,
identityKeyPair: identityKeyPair,
signedPreKey: SignedPreKeyStoreImpl.generateSignedPreKey(keyId: PreKeyId.random(), signedBy: identityKeyPair.keyPair.privateKey),
lastResortPreKey: generateLastResortKyberPreKey(signedBy: identityKeyPair.keyPair.privateKey),
),
)
}
var didFinalizeRegistrationPrekeys = false
@ -63,17 +59,16 @@ class MockPreKeyManager: PreKeyManager {
func finalizeRegistrationPreKeys(
_ bundles: RegistrationPreKeyUploadBundles,
uploadDidSucceed: Bool,
) -> Task<Void, Error> {
) async {
didFinalizeRegistrationPrekeys = true
return Task {}
}
func rotateOneTimePreKeysForRegistration(auth: ChatServiceAuth) -> Task<Void, Error> {
return Task {}
func rotateOneTimePreKeysForRegistration(auth: ChatServiceAuth) async throws {
}
func rotateSignedPreKeysIfNeeded() -> Task<Void, Error> { Task {} }
func refreshOneTimePreKeys(forIdentity identity: OWSIdentity, alsoRefreshSignedPreKey shouldRefreshSignedPreKey: Bool) {
func rotateSignedPreKeysIfNeeded() async throws {}
func refreshOneTimePreKeys(forIdentity identity: OWSIdentity, alsoRefreshSignedPreKey shouldRefreshSignedPreKey: Bool) async throws {
attemptedRefreshes.append((identity, shouldRefreshSignedPreKey))
}

54
SignalServiceKit/Account/PreKeys/PreKeyManager.swift
View File

@ -8,59 +8,43 @@ import Foundation
public protocol PreKeyManager {
func isAppLockedDueToPreKeyUpdateFailures(tx: DBReadTransaction) -> Bool
func checkPreKeysIfNecessary(tx: DBReadTransaction)
func checkPreKeysIfNecessary() async throws
func rotatePreKeysOnUpgradeIfNecessary(for identity: OWSIdentity) async throws
/// Creates a new set of prekeys for registration, creating a new identity key if needed
/// (or reusing the existing identity key).
/// These keys are persisted before this method's promise resolves, but best effort
/// should be taken to finalize the keys once they have been accepted by the server.
/// Creates a new set of prekeys for registration, creating a new identity
/// key if needed (or reusing the existing identity key).
///
/// - returns: A task representing the completion of the prekey operation. This task is _not_
/// a child task of the calling context; this call returns once the task has been scheduled, but running
/// the task is handled separately (but can be optionally waited on by the caller).
func createPreKeysForRegistration() -> Task<RegistrationPreKeyUploadBundles, Error>
/// These keys are persisted before this method returns, but best effort
/// should be taken to finalize the keys after the server accepts them.
func createPreKeysForRegistration() async -> RegistrationPreKeyUploadBundles
/// Creates a new set of prekeys for provisioning (linking a new secondary device),
/// using the provided identity keys (which are delivered from the primary during linking).
/// These keys are persisted before this method's promise resolves, but best effort
/// should be taken to finalize the keys once they have been accepted by the server.
/// Use `finalizeRegistrationPreKeys` to finalize once linking is complete.
///
/// - returns: A task representing the completion of the prekey operation. This task is _not_
/// a child task of the calling context; this call returns once the task has been scheduled, but running
/// the task is handled separately (but can be optionally waited on by the caller).
/// Creates a new set of prekeys for provisioning (linking a new secondary
/// device), using the provided identity keys (which are delivered from the
/// primary during linking). These keys are persisted before this method
/// returns, but best effort should be taken to finalize the keys after the
/// server accepts them.
func createPreKeysForProvisioning(
aciIdentityKeyPair: ECKeyPair,
pniIdentityKeyPair: ECKeyPair,
) -> Task<RegistrationPreKeyUploadBundles, Error>
) async -> RegistrationPreKeyUploadBundles
/// Called on a best-effort basis. Consequences of not calling this is that the keys are still
/// persisted (from prior to uploading) but they aren't marked current and accepted.
///
/// - returns: A task representing the completion of the prekey operation. This task is _not_
/// a child task of the calling context; this call returns once the task has been scheduled, but running
/// the task is handled separately (but can be optionally waited on by the caller).
/// Called on a best-effort basis. Consequences of not calling this is that
/// the keys are still persisted (from prior to uploading) but they aren't
/// marked current and accepted.
func finalizeRegistrationPreKeys(
_ bundles: RegistrationPreKeyUploadBundles,
uploadDidSucceed: Bool,
) -> Task<Void, Error>
) async
/// - returns: A task representing the completion of the prekey operation. This task is _not_
/// a child task of the calling context; this call returns once the task has been scheduled, but running
/// the task is handled separately (but can be optionally waited on by the caller).
func rotateOneTimePreKeysForRegistration(auth: ChatServiceAuth) -> Task<Void, Error>
func rotateOneTimePreKeysForRegistration(auth: ChatServiceAuth) async throws
/// - returns: A task representing the completion of the prekey operation. This task is _not_
/// a child task of the calling context; this call returns once the task has been scheduled, but running
/// the task is handled separately (but can be optionally waited on by the caller).
func rotateSignedPreKeysIfNeeded() -> Task<Void, Error>
func rotateSignedPreKeysIfNeeded() async throws
func refreshOneTimePreKeys(
forIdentity identity: OWSIdentity,
alsoRefreshSignedPreKey shouldRefreshSignedPreKey: Bool,
)
) async throws
func setIsChangingNumber(_ isChangingNumber: Bool)
}

2
SignalServiceKit/Account/PreKeys/PreKeyTaskManager.swift
View File

@ -79,7 +79,6 @@ struct PreKeyTaskManager {
/// ALWAYS changes the targeted keys (regardless of current key state)
func createForRegistration() async -> RegistrationPreKeyUploadBundles {
logger.info("Create for registration")
let (aciBundle, pniBundle) = await db.awaitableWrite { tx in
let aciBundle = self.generateKeysForRegistration(identity: .aci, tx: tx)
let pniBundle = self.generateKeysForRegistration(identity: .pni, tx: tx)
@ -98,7 +97,6 @@ struct PreKeyTaskManager {
pniIdentityKeyPair: ECKeyPair,
) async -> RegistrationPreKeyUploadBundles {
logger.info("Create for provisioning")
let (aciBundle, pniBundle) = await db.awaitableWrite { tx in
let aciBundle = self.generateKeysForProvisioning(
identity: .aci,

97
SignalServiceKit/Account/PreKeys/PrekeyManagerImpl.swift
View File

@ -37,7 +37,7 @@ public class PreKeyManagerImpl: PreKeyManager {
/// Some of our pre-key operations depend on the service state, e.g. we need to check our one-time-prekey count
/// before we decide to upload new ones. This potentially entails multiple async operations, all of which should
/// complete before starting any other pre-key operation. That's why they must run in serial.
private static let taskQueue = SerialTaskQueue()
private let taskQueue = ConcurrentTaskQueue(concurrentLimit: 1)
private let db: any DB
private let identityManager: OWSIdentityManager
@ -115,15 +115,13 @@ public class PreKeyManagerImpl: PreKeyManager {
lastOneTimePreKeyCheckTimestamp = Date()
}
public func checkPreKeysIfNecessary(tx: DBReadTransaction) {
checkPreKeys(shouldThrottle: true, tx: tx)
public func checkPreKeysIfNecessary() async throws {
try await checkPreKeys(shouldThrottle: true)
}
fileprivate func checkPreKeys(shouldThrottle: Bool, tx: DBReadTransaction) {
guard
CurrentAppContext().isMainAppAndActive
else {
return
fileprivate func checkPreKeys(shouldThrottle: Bool) async throws {
guard CurrentAppContext().isMainAppAndActive else {
throw OWSGenericError("must be the main app")
}
let shouldCheckOneTimePreKeys = {
@ -145,24 +143,22 @@ public class PreKeyManagerImpl: PreKeyManager {
logger.warn("Skipping PNI pre key check due to change number.")
}
_ = self._checkPreKeys(
try await self._checkPreKeys(
shouldCheckOneTimePreKeys: shouldCheckOneTimePreKeys,
shouldCheckPniPreKeys: !shouldSkipPniPreKeyCheck,
tx: tx,
)
}
private func _checkPreKeys(
shouldCheckOneTimePreKeys: Bool,
shouldCheckPniPreKeys: Bool,
tx: DBReadTransaction,
) -> Task<Void, any Error> {
) async throws {
var targets: PreKeyTargets = [.signedPreKey, .lastResortPqPreKey]
if shouldCheckOneTimePreKeys {
targets.insert(target: .oneTimePreKey)
targets.insert(target: .oneTimePqPreKey)
}
return Self.taskQueue.enqueue { [self, chatConnectionManager, taskManager, targets] in
try await taskQueue.run {
try await chatConnectionManager.waitForIdentifiedConnectionToOpen()
try Task.checkCancellation()
try await taskManager.refresh(identity: .aci, targets: targets, auth: .implicit())
@ -177,66 +173,48 @@ public class PreKeyManagerImpl: PreKeyManager {
}
}
public func createPreKeysForRegistration() -> Task<RegistrationPreKeyUploadBundles, Error> {
public func createPreKeysForRegistration() async -> RegistrationPreKeyUploadBundles {
logger.info("Create registration prekeys")
/// Note that we do not report a `refreshOneTimePreKeysCheckDidSucceed`
/// because this operation does not generate one time prekeys, so we
/// shouldn't mark the routine refresh as having been "checked".
return Self.taskQueue.enqueueCancellingPrevious { [taskManager] in
try Task.checkCancellation()
return await taskManager.createForRegistration()
}
return await taskManager.createForRegistration()
}
public func createPreKeysForProvisioning(
aciIdentityKeyPair: ECKeyPair,
pniIdentityKeyPair: ECKeyPair,
) -> Task<RegistrationPreKeyUploadBundles, Error> {
) async -> RegistrationPreKeyUploadBundles {
logger.info("Create provisioning prekeys")
/// Note that we do not report a `refreshOneTimePreKeysCheckDidSucceed`
/// because this operation does not generate one time prekeys, so we
/// shouldn't mark the routine refresh as having been "checked".
return Self.taskQueue.enqueueCancellingPrevious { [taskManager] in
try Task.checkCancellation()
return await taskManager.createForProvisioning(
aciIdentityKeyPair: aciIdentityKeyPair,
pniIdentityKeyPair: pniIdentityKeyPair,
)
}
return await taskManager.createForProvisioning(
aciIdentityKeyPair: aciIdentityKeyPair,
pniIdentityKeyPair: pniIdentityKeyPair,
)
}
public func finalizeRegistrationPreKeys(
_ bundles: RegistrationPreKeyUploadBundles,
uploadDidSucceed: Bool,
) -> Task<Void, Error> {
) async {
logger.info("Finalize registration prekeys")
return Self.taskQueue.enqueue { [taskManager] in
try Task.checkCancellation()
await taskManager.persistAfterRegistration(
bundles: bundles,
uploadDidSucceed: uploadDidSucceed,
)
}
await taskManager.persistAfterRegistration(
bundles: bundles,
uploadDidSucceed: uploadDidSucceed,
)
}
public func rotateOneTimePreKeysForRegistration(auth: ChatServiceAuth) -> Task<Void, Error> {
public func rotateOneTimePreKeysForRegistration(auth: ChatServiceAuth) async throws {
logger.info("Rotate one-time prekeys for registration")
return Self.taskQueue.enqueue { [weak self, taskManager] in
return try await taskQueue.run {
try Task.checkCancellation()
try await taskManager.createOneTimePreKeys(identity: .aci, auth: auth)
try Task.checkCancellation()
try await taskManager.createOneTimePreKeys(identity: .pni, auth: auth)
self?.refreshOneTimePreKeysCheckDidSucceed()
self.refreshOneTimePreKeysCheckDidSucceed()
}
}
public func rotateSignedPreKeysIfNeeded() -> Task<Void, Error> {
public func rotateSignedPreKeysIfNeeded() async throws {
logger.info("Rotating signed prekeys if needed")
return db.read { tx in
return _checkPreKeys(shouldCheckOneTimePreKeys: false, shouldCheckPniPreKeys: true, tx: tx)
}
try await _checkPreKeys(shouldCheckOneTimePreKeys: false, shouldCheckPniPreKeys: true)
}
/// Refresh one-time pre-keys for the given identity, and optionally refresh
@ -244,18 +222,6 @@ public class PreKeyManagerImpl: PreKeyManager {
public func refreshOneTimePreKeys(
forIdentity identity: OWSIdentity,
alsoRefreshSignedPreKey shouldRefreshSignedPreKey: Bool,
) {
Task {
try? await self._refreshOneTimePreKeys(
forIdentity: identity,
alsoRefreshSignedPreKey: shouldRefreshSignedPreKey,
)
}
}
private func _refreshOneTimePreKeys(
forIdentity identity: OWSIdentity,
alsoRefreshSignedPreKey shouldRefreshSignedPreKey: Bool,
) async throws {
logger.info("[\(identity)] Force refresh onetime prekeys (also refresh signed pre key? \(shouldRefreshSignedPreKey))")
/// Note that we do not report a `refreshOneTimePreKeysCheckDidSucceed`
@ -265,11 +231,11 @@ public class PreKeyManagerImpl: PreKeyManager {
var targets: PreKeyTargets = [.oneTimePreKey, .oneTimePqPreKey]
if shouldRefreshSignedPreKey {
targets.insert(.signedPreKey)
targets.insert(target: .lastResortPqPreKey)
targets.insert(.lastResortPqPreKey)
}
try await waitUntilNotChangingNumberIfNeeded(targets: targets)
let task = Self.taskQueue.enqueue { [taskManager, targets] in
try await taskQueue.run {
try Task.checkCancellation()
try await taskManager.refresh(
identity: identity,
@ -278,7 +244,6 @@ public class PreKeyManagerImpl: PreKeyManager {
auth: .implicit(),
)
}
try await task.value
}
/// If we don't have a PNI identity key, we should not run PNI operations.
@ -310,7 +275,7 @@ public class PreKeyManagerImpl: PreKeyManager {
}
do {
try await chatConnectionManager.waitForIdentifiedConnectionToOpen()
try await _refreshOneTimePreKeys(forIdentity: identity, alsoRefreshSignedPreKey: true)
try await refreshOneTimePreKeys(forIdentity: identity, alsoRefreshSignedPreKey: true)
} catch {
logger.warn("Couldn't rotate pre keys: \(error)")
throw error
@ -366,8 +331,8 @@ public class PreKeyManagerImpl: PreKeyManager {
#if TESTABLE_BUILD
public extension PreKeyManagerImpl {
func checkPreKeysImmediately(tx: DBReadTransaction) {
checkPreKeys(shouldThrottle: false, tx: tx)
func checkPreKeysImmediately() async throws {
try await checkPreKeys(shouldThrottle: false)
}
}

13
SignalServiceKit/ChangePhoneNumber/ChangePhoneNumberPniManager.swift
View File

@ -205,13 +205,12 @@ class ChangePhoneNumberPniManagerImpl: ChangePhoneNumberPniManager {
// Followup tasks
tx.addSyncCompletion { [preKeyManager] in
// Since we rotated the identity key, we need new one-time pre-keys.
// However, no need to update the signed pre-key, which we also just
// rotated.
preKeyManager.refreshOneTimePreKeys(
forIdentity: .pni,
alsoRefreshSignedPreKey: refreshSignedPreKey,
)
Task {
try? await preKeyManager.refreshOneTimePreKeys(
forIdentity: .pni,
alsoRefreshSignedPreKey: refreshSignedPreKey,
)
}
}
}
}

6
SignalServiceKit/Messages/IncomingPniChangeNumberProcessor.swift
View File

@ -87,7 +87,11 @@ public class IncomingPniChangeNumberProcessorImpl: IncomingPniChangeNumberProces
// We need to refresh our one-time pre-keys, and should also refresh
// our signed pre-key so we use the one generated on the primary for as
// little time as possible.
preKeyManager.refreshOneTimePreKeys(forIdentity: .pni, alsoRefreshSignedPreKey: true)
tx.addSyncCompletion { [preKeyManager] in
Task {
try? await preKeyManager.refreshOneTimePreKeys(forIdentity: .pni, alsoRefreshSignedPreKey: true)
}
}
}
private func deserializeIncomingPniChangePhoneNumber(

2
SignalServiceKit/Messages/MessageSender.swift
View File

@ -442,7 +442,7 @@ public class MessageSender {
// If this succeeds, or if we hit an error, allow another attempt.
self.pendingPreKeyRotation.set(nil)
}
try await self.preKeyManager.rotateSignedPreKeysIfNeeded().value
try await self.preKeyManager.rotateSignedPreKeysIfNeeded()
}
}
return nil

8
SignalServiceKit/Messages/OWSMessageDecrypter.swift
View File

@ -419,7 +419,9 @@ public class OWSMessageDecrypter {
sendReactiveProfileKeyIfNecessary(to: sourceAci, tx: transaction)
case .preKey:
if DependenciesBridge.shared.tsAccountManager.registrationState(tx: transaction).isRegistered {
DependenciesBridge.shared.preKeyManager.checkPreKeysIfNecessary(tx: transaction)
Task {
try? await DependenciesBridge.shared.preKeyManager.checkPreKeysIfNecessary()
}
}
let message = try PreKeySignalMessage(bytes: encryptedData)
plaintext = try signalDecryptPreKey(
@ -607,7 +609,9 @@ public class OWSMessageDecrypter {
decryptResult.messageType == .prekey,
DependenciesBridge.shared.tsAccountManager.registrationState(tx: transaction).isRegistered
{
DependenciesBridge.shared.preKeyManager.checkPreKeysIfNecessary(tx: transaction)
Task {
try? await DependenciesBridge.shared.preKeyManager.checkPreKeysIfNecessary()
}
}
let envelopeBuilder = validatedEnvelope.envelope.asBuilder()