From 3df0e72eda41bcb17ae6254541eac0fc50a5fde9 Mon Sep 17 00:00:00 2001 From: Michael Kirk Date: Wed, 5 Sep 2018 10:50:01 -0600 Subject: [PATCH] Extract SPK rotation and CreatePreKey operations --- .../OWS106EnsureProfileComplete.swift | 11 +- .../src/Account/AccountManager.swift | 4 + .../src/Account/PreKeyRefreshOperation.swift | 82 +++++++++ .../src/Account/TSAccountManager.m | 4 +- .../src/Account/TSPreKeyManager.h | 6 +- .../src/Account/TSPreKeyManager.m | 164 +++++++----------- .../src/Messages/OWSMessageSender.m | 5 +- .../src/Network/ServiceSocket.swift | 18 ++ .../OWSPrimaryStorage+SignedPreKeyStore.h | 1 + .../OWSPrimaryStorage+SignedPreKeyStore.m | 19 ++ 10 files changed, 202 insertions(+), 112 deletions(-) diff --git a/SignalMessaging/environment/migrations/OWS106EnsureProfileComplete.swift b/SignalMessaging/environment/migrations/OWS106EnsureProfileComplete.swift index cec609271..32f21407b 100644 --- a/SignalMessaging/environment/migrations/OWS106EnsureProfileComplete.swift +++ b/SignalMessaging/environment/migrations/OWS106EnsureProfileComplete.swift @@ -101,13 +101,12 @@ public class OWS106EnsureProfileComplete: OWSDatabaseMigration { case SignalServiceProfile.ValidationError.invalidIdentityKey(let description): Logger.warn("detected incomplete profile for \(localRecipientId) error: \(description)") // This is the error condition we're looking for. Update prekeys to properly set the identity key, completing registration. - TSPreKeyManager.registerPreKeys(with: .signedAndOneTime, - success: { - Logger.info("successfully uploaded pre-keys. Profile should be fixed.") - fulfill(()) + TSPreKeyManager.createPreKeys(success: { + Logger.info("successfully uploaded pre-keys. Profile should be fixed.") + fulfill(()) }, - failure: { _ in - reject(OWSErrorWithCodeDescription(.signalServiceFailure, "\(self.logTag) Unknown error")) + failure: { _ in + reject(OWSErrorWithCodeDescription(.signalServiceFailure, "\(self.logTag) Unknown error")) }) default: reject(error) diff --git a/SignalServiceKit/src/Account/AccountManager.swift b/SignalServiceKit/src/Account/AccountManager.swift index bdcf59e43..e4cd34e35 100644 --- a/SignalServiceKit/src/Account/AccountManager.swift +++ b/SignalServiceKit/src/Account/AccountManager.swift @@ -26,4 +26,8 @@ public class AccountManager: NSObject { public func setPreKeys(identityKey: IdentityKey, signedPreKeyRecord: SignedPreKeyRecord, preKeyRecords: [PreKeyRecord]) -> Promise { return serviceSocket.registerPreKeys(identityKey: identityKey, signedPreKeyRecord: signedPreKeyRecord, preKeyRecords: preKeyRecords) } + + public func setSignedPreKey(_ signedPreKey: SignedPreKeyRecord) -> Promise { + return serviceSocket.setCurrentSignedPreKey(signedPreKey) + } } diff --git a/SignalServiceKit/src/Account/PreKeyRefreshOperation.swift b/SignalServiceKit/src/Account/PreKeyRefreshOperation.swift index 7c9f0b8b2..d4def95cb 100644 --- a/SignalServiceKit/src/Account/PreKeyRefreshOperation.swift +++ b/SignalServiceKit/src/Account/PreKeyRefreshOperation.swift @@ -9,6 +9,87 @@ import PromiseKit // whenever ~2/3 of them have been consumed. let kEphemeralPreKeysMinimumCount: UInt = 35 +@objc(SSKCreatePreKeysOperation) +public class CreatePreKeysOperation: OWSOperation { + private var accountManager: AccountManager { + return AccountManager.shared + } + private var primaryStorage: OWSPrimaryStorage { + return OWSPrimaryStorage.shared() + } + + private var identityKeyManager: OWSIdentityManager { + return OWSIdentityManager.shared() + } + + public override func run() { + Logger.debug("") + + if self.identityKeyManager.identityKeyPair() == nil { + self.identityKeyManager.generateNewIdentityKey() + } + let identityKey: Data = self.identityKeyManager.identityKeyPair()!.publicKey + let signedPreKeyRecord: SignedPreKeyRecord = self.primaryStorage.generateRandomSignedRecord() + let preKeyRecords: [PreKeyRecord] = self.primaryStorage.generatePreKeyRecords() + + firstly { + return self.accountManager.setPreKeys(identityKey: identityKey, signedPreKeyRecord: signedPreKeyRecord, preKeyRecords: preKeyRecords) + }.then { () -> Void in + signedPreKeyRecord.markAsAcceptedByService() + self.primaryStorage.storeSignedPreKey(signedPreKeyRecord.id, signedPreKeyRecord: signedPreKeyRecord) + self.primaryStorage.setCurrentSignedPrekeyId(signedPreKeyRecord.id) + self.primaryStorage.storePreKeyRecords(preKeyRecords) + }.then { () -> Void in + Logger.debug("done") + self.reportSuccess() + }.catch { error in + self.reportError(error) + }.retainUntilComplete() + } +} + +@objc(SSKRotateSignedPreKeyOperation) +public class RotateSignedPreKeyOperation: OWSOperation { + private var tsAccountManager: TSAccountManager { + return TSAccountManager.sharedInstance() + } + + private var accountManager: AccountManager { + return AccountManager.shared + } + + private var primaryStorage: OWSPrimaryStorage { + return OWSPrimaryStorage.shared() + } + + public override func run() { + Logger.debug("") + + guard tsAccountManager.isRegistered() else { + Logger.debug("skipping - not registered") + return + } + + let signedPreKeyRecord: SignedPreKeyRecord = self.primaryStorage.generateRandomSignedRecord() + + firstly { + return self.accountManager.setSignedPreKey(signedPreKeyRecord) + }.then(on: DispatchQueue.global()) { () -> Void in + Logger.info("Successfully uploaded signed PreKey") + signedPreKeyRecord.markAsAcceptedByService() + self.primaryStorage.storeSignedPreKey(signedPreKeyRecord.id, signedPreKeyRecord: signedPreKeyRecord) + self.primaryStorage.setCurrentSignedPrekeyId(signedPreKeyRecord.id) + + TSPreKeyManager.clearSignedPreKeyRecords() + }.then { () -> Void in + Logger.debug("done") + self.reportSuccess() + }.catch { error in + self.reportError(error) + }.retainUntilComplete() + } +} + @objc(SSKRefreshPreKeysOperation) public class RefreshPreKeysOperation: OWSOperation { @@ -19,6 +100,7 @@ public class RefreshPreKeysOperation: OWSOperation { private var accountManager: AccountManager { return AccountManager.shared } + private var primaryStorage: OWSPrimaryStorage { return OWSPrimaryStorage.shared() } diff --git a/SignalServiceKit/src/Account/TSAccountManager.m b/SignalServiceKit/src/Account/TSAccountManager.m index bb835ed0a..d5dbca936 100644 --- a/SignalServiceKit/src/Account/TSAccountManager.m +++ b/SignalServiceKit/src/Account/TSAccountManager.m @@ -371,9 +371,7 @@ NSString *const TSAccountManager_ServerSignalingKey = @"TSStorageServerSignaling case 204: { OWSLogInfo(@"Verification code accepted."); [self storeServerAuthToken:authToken signalingKey:signalingKey]; - [TSPreKeyManager registerPreKeysWithMode:RefreshPreKeysMode_SignedAndOneTime - success:successBlock - failure:failureBlock]; + [TSPreKeyManager createPreKeysWithSuccess:successBlock failure:failureBlock]; break; } default: { diff --git a/SignalServiceKit/src/Account/TSPreKeyManager.h b/SignalServiceKit/src/Account/TSPreKeyManager.h index da6e496f9..19964f430 100644 --- a/SignalServiceKit/src/Account/TSPreKeyManager.h +++ b/SignalServiceKit/src/Account/TSPreKeyManager.h @@ -21,9 +21,9 @@ typedef NS_ENUM(NSInteger, RefreshPreKeysMode) { + (BOOL)isAppLockedDueToPreKeyUpdateFailures; -+ (void)registerPreKeysWithMode:(RefreshPreKeysMode)mode - success:(void (^)(void))successHandler - failure:(void (^)(NSError *error))failureHandler; ++ (void)rotateSignedPreKeyWithSuccess:(void (^)(void))successHandler failure:(void (^)(NSError *error))failureHandler; + ++ (void)createPreKeysWithSuccess:(void (^)(void))successHandler failure:(void (^)(NSError *error))failureHandler; + (void)checkPreKeys; diff --git a/SignalServiceKit/src/Account/TSPreKeyManager.m b/SignalServiceKit/src/Account/TSPreKeyManager.m index 7f142f9c4..182ad8491 100644 --- a/SignalServiceKit/src/Account/TSPreKeyManager.m +++ b/SignalServiceKit/src/Account/TSPreKeyManager.m @@ -96,114 +96,82 @@ static const NSUInteger kMaxPrekeyUpdateFailureCount = 5; } OWSAssertDebug(CurrentAppContext().isMainAppAndActive); - // Update the prekey check timestamp. - dispatch_async(TSPreKeyManager.prekeyQueue, ^{ + if (!TSAccountManager.isRegistered) { + return; + } + + NSBlockOperation *checkIfPreKeyRefreshNecessaryOperation = [NSBlockOperation blockOperationWithBlock:^{ BOOL shouldCheck = (lastPreKeyCheckTimestamp == nil || fabs([lastPreKeyCheckTimestamp timeIntervalSinceNow]) >= kPreKeyCheckFrequencySeconds); - if (shouldCheck) { - // Optimistically mark the prekeys as checked. This - // de-bounces prekey checks. - // - // If the check or key registration fails, the prekeys - // will be marked as _NOT_ checked. - // - // Note: [TSPreKeyManager checkPreKeys] will also - // optimistically mark them as checked. This - // redundancy is fine and precludes a race - // condition. - lastPreKeyCheckTimestamp = [NSDate date]; + if (!shouldCheck) { + return; + } - if ([TSAccountManager isRegistered]) { - dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{ - [TSPreKeyManager checkPreKeys]; - }); - } + SSKRefreshPreKeysOperation *refreshOperation = [SSKRefreshPreKeysOperation new]; + // Run inline instead of enqueuing to ensure we don't have multiple "check" operations + // that in turn enqueue multiple "refresh" operations. + [refreshOperation run]; + }]; + + NSBlockOperation *checkIfSignedRotationNecessaryOperation = [NSBlockOperation blockOperationWithBlock:^{ + OWSPrimaryStorage *primaryStorage = [OWSPrimaryStorage sharedManager]; + SignedPreKeyRecord *_Nullable signedPreKey = [primaryStorage currentSignedPreKey]; + + BOOL shouldCheck + = !signedPreKey || fabs(signedPreKey.generatedAt.timeIntervalSinceNow) >= kSignedPreKeyRotationTime; + if (!shouldCheck) { + return; + } + + SSKRotateSignedPreKeyOperation *rotationOperation = [SSKRotateSignedPreKeyOperation new]; + // Run inline instead of enqueuing to ensure we don't have multiple "check" operations + // that in turn enqueue multiple "refresh" operations. + [rotationOperation run]; + }]; + + // Order matters here. + // We add the PreKeyRefresh first - if it *does* upload new keys, it'll upload a new SignedPreKey + // In that case our SPK rotation will be a no-op. + [self.operationQueue + addOperations:@[ checkIfPreKeyRefreshNecessaryOperation, checkIfSignedRotationNecessaryOperation ] + waitUntilFinished:NO]; +} + ++ (void)createPreKeysWithSuccess:(void (^)(void))successHandler failure:(void (^)(NSError *error))failureHandler +{ + dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{ + SSKCreatePreKeysOperation *operation = [SSKCreatePreKeysOperation new]; + [self.operationQueue addOperations:@[ operation ] waitUntilFinished:YES]; + + NSError *_Nullable error = operation.failingError; + if (error) { + dispatch_async(dispatch_get_main_queue(), ^{ + failureHandler(error); + }); + } else { + dispatch_async(dispatch_get_main_queue(), ^{ + successHandler(); + }); } }); } -+ (void)registerPreKeysWithMode:(RefreshPreKeysMode)mode - success:(void (^)(void))successHandler - failure:(void (^)(NSError *error))failureHandler ++ (void)rotateSignedPreKeyWithSuccess:(void (^)(void))successHandler failure:(void (^)(NSError *error))failureHandler { - // We use prekeyQueue to serialize this logic and ensure that only - // one thread is "registering" or "clearing" prekeys at a time. - dispatch_async(TSPreKeyManager.prekeyQueue, ^{ - // Mark the prekeys as checked every time we try to register prekeys. - lastPreKeyCheckTimestamp = [NSDate date]; - - RefreshPreKeysMode modeCopy = mode; - OWSPrimaryStorage *primaryStorage = [OWSPrimaryStorage sharedManager]; - ECKeyPair *identityKeyPair = [[OWSIdentityManager sharedManager] identityKeyPair]; - - if (!identityKeyPair) { - [[OWSIdentityManager sharedManager] generateNewIdentityKey]; - identityKeyPair = [[OWSIdentityManager sharedManager] identityKeyPair]; - - // Switch modes if necessary. - modeCopy = RefreshPreKeysMode_SignedAndOneTime; - } - - SignedPreKeyRecord *signedPreKey = [primaryStorage generateRandomSignedRecord]; - // Store the new signed key immediately, before it is sent to the - // service to prevent race conditions and other edge cases. - [primaryStorage storeSignedPreKey:signedPreKey.Id signedPreKeyRecord:signedPreKey]; - - NSArray *preKeys = nil; - TSRequest *request; - NSString *description; - if (modeCopy == RefreshPreKeysMode_SignedAndOneTime) { - description = @"signed and one-time prekeys"; - preKeys = [primaryStorage generatePreKeyRecords]; - // Store the new one-time keys immediately, before they are sent to the - // service to prevent race conditions and other edge cases. - [primaryStorage storePreKeyRecords:preKeys]; - - request = [OWSRequestFactory registerPrekeysRequestWithPrekeyArray:preKeys - identityKey:identityKeyPair.publicKey - signedPreKey:signedPreKey]; - } else { - description = @"just signed prekey"; - request = [OWSRequestFactory registerSignedPrekeyRequestWithSignedPreKeyRecord:signedPreKey]; - } - - [[TSNetworkManager sharedManager] makeRequest:request - success:^(NSURLSessionDataTask *task, id responseObject) { - OWSLogInfo(@"Successfully registered %@.", description); - - // Mark signed prekey as accepted by service. - [signedPreKey markAsAcceptedByService]; - [primaryStorage storeSignedPreKey:signedPreKey.Id signedPreKeyRecord:signedPreKey]; - - // On success, update the "current" signed prekey state. - [primaryStorage setCurrentSignedPrekeyId:signedPreKey.Id]; - - successHandler(); - - [TSPreKeyManager clearPreKeyUpdateFailureCount]; - } - failure:^(NSURLSessionDataTask *task, NSError *error) { - if (!IsNSErrorNetworkFailure(error)) { - if (modeCopy == RefreshPreKeysMode_SignedAndOneTime) { - OWSProdError([OWSAnalyticsEvents errorPrekeysUpdateFailedSignedAndOnetime]); - } else { - OWSProdError([OWSAnalyticsEvents errorPrekeysUpdateFailedJustSigned]); - } - } + dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{ + SSKRotateSignedPreKeyOperation *operation = [SSKRotateSignedPreKeyOperation new]; + [self.operationQueue addOperations:@[ operation ] waitUntilFinished:YES]; + NSError *_Nullable error = operation.failingError; + if (error) { + dispatch_async(dispatch_get_main_queue(), ^{ failureHandler(error); - - NSInteger statusCode = 0; - if ([task.response isKindOfClass:[NSHTTPURLResponse class]]) { - NSHTTPURLResponse *httpResponse = (NSHTTPURLResponse *)task.response; - statusCode = httpResponse.statusCode; - } - if (statusCode >= 400 && statusCode <= 599) { - // Only treat 4xx and 5xx errors from the service as failures. - // Ignore network failures, for example. - [TSPreKeyManager incrementPreKeyUpdateFailureCount]; - } - }]; + }); + } else { + dispatch_async(dispatch_get_main_queue(), ^{ + successHandler(); + }); + } }); } diff --git a/SignalServiceKit/src/Messages/OWSMessageSender.m b/SignalServiceKit/src/Messages/OWSMessageSender.m index 8267d5ce0..caabda65f 100644 --- a/SignalServiceKit/src/Messages/OWSMessageSender.m +++ b/SignalServiceKit/src/Messages/OWSMessageSender.m @@ -732,14 +732,15 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException"; // // Only try to update the signed prekey; updating it is sufficient to // re-enable message sending. - [TSPreKeyManager registerPreKeysWithMode:RefreshPreKeysMode_SignedOnly - success:^{ + [TSPreKeyManager + rotateSignedPreKeyWithSuccess:^{ OWSLogInfo(@"New prekeys registered with server."); } failure:^(NSError *error) { OWSLogWarn(@"Failed to update prekeys with the server: %@", error); }]; + // MJK TODO move this into the success/failure handlers, otherwise we'll just fail again. NSError *error = OWSErrorMakeMessageSendDisabledDueToPreKeyUpdateFailuresError(); [error setIsRetryable:YES]; return failureHandler(error); diff --git a/SignalServiceKit/src/Network/ServiceSocket.swift b/SignalServiceKit/src/Network/ServiceSocket.swift index c00798842..a503d9ae0 100644 --- a/SignalServiceKit/src/Network/ServiceSocket.swift +++ b/SignalServiceKit/src/Network/ServiceSocket.swift @@ -8,6 +8,7 @@ import PromiseKit protocol ServiceSocket { func getAvailablePreKeys() -> Promise func registerPreKeys(identityKey: IdentityKey, signedPreKeyRecord: SignedPreKeyRecord, preKeyRecords: [PreKeyRecord]) -> Promise + func setCurrentSignedPreKey(_ signedPreKey: SignedPreKeyRecord) -> Promise } class ServiceRestSocket: ServiceSocket { @@ -69,4 +70,21 @@ class ServiceRestSocket: ServiceSocket { }) return promise } + + public func setCurrentSignedPreKey(_ signedPreKey: SignedPreKeyRecord) -> Promise { + let (promise, fulfill, reject) = Promise.pending() + let request = OWSRequestFactory.registerSignedPrekeyRequest(with: signedPreKey) + + networkManager.makeRequest(request, + success: { (_, _) in + Logger.debug("success") + fulfill(()) + + }, + failure: { (_, error) in + Logger.debug("error: \(error)") + reject(error) + }) + return promise + } } diff --git a/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+SignedPreKeyStore.h b/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+SignedPreKeyStore.h index 687d6f6cf..752581f57 100644 --- a/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+SignedPreKeyStore.h +++ b/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+SignedPreKeyStore.h @@ -19,6 +19,7 @@ extern NSString *const OWSPrimaryStorageSignedPreKeyStoreCollection; // Returns nil if no current signed prekey id is found. - (nullable NSNumber *)currentSignedPrekeyId; - (void)setCurrentSignedPrekeyId:(int)value; +- (nullable SignedPreKeyRecord *)currentSignedPreKey; #pragma mark - Prekey update failures diff --git a/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+SignedPreKeyStore.m b/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+SignedPreKeyStore.m index 84cfbbfc0..8d608b483 100644 --- a/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+SignedPreKeyStore.m +++ b/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+SignedPreKeyStore.m @@ -105,6 +105,25 @@ NSString *const OWSPrimaryStorageKeyPrekeyCurrentSignedPrekeyId = @"currentSigne inCollection:OWSPrimaryStorageSignedPreKeyMetadataCollection]; } +- (nullable SignedPreKeyRecord *)currentSignedPreKey +{ + __block SignedPreKeyRecord *_Nullable currentRecord; + + [self.dbReadConnection readWithBlock:^(YapDatabaseReadTransaction *_Nonnull transaction) { + NSNumber *preKeyId = [transaction objectForKey:OWSPrimaryStorageKeyPrekeyCurrentSignedPrekeyId + inCollection:OWSPrimaryStorageSignedPreKeyMetadataCollection]; + + if (preKeyId == nil) { + return; + } + + currentRecord = + [transaction objectForKey:preKeyId.stringValue inCollection:OWSPrimaryStorageSignedPreKeyStoreCollection]; + }]; + + return currentRecord; +} + #pragma mark - Prekey update failures - (int)prekeyUpdateFailureCount