Extract SPK rotation and CreatePreKey operations

This commit is contained in:
Michael Kirk 2018-09-05 10:50:01 -06:00
parent 286d3c8ce9
commit 3df0e72eda
10 changed files with 202 additions and 112 deletions

View File

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

View File

@ -26,4 +26,8 @@ public class AccountManager: NSObject {
public func setPreKeys(identityKey: IdentityKey, signedPreKeyRecord: SignedPreKeyRecord, preKeyRecords: [PreKeyRecord]) -> Promise<Void> {
return serviceSocket.registerPreKeys(identityKey: identityKey, signedPreKeyRecord: signedPreKeyRecord, preKeyRecords: preKeyRecords)
}
public func setSignedPreKey(_ signedPreKey: SignedPreKeyRecord) -> Promise<Void> {
return serviceSocket.setCurrentSignedPreKey(signedPreKey)
}
}

View File

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

View File

@ -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: {

View File

@ -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;

View File

@ -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();
});
}
});
}

View File

@ -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);

View File

@ -8,6 +8,7 @@ import PromiseKit
protocol ServiceSocket {
func getAvailablePreKeys() -> Promise<Int>
func registerPreKeys(identityKey: IdentityKey, signedPreKeyRecord: SignedPreKeyRecord, preKeyRecords: [PreKeyRecord]) -> Promise<Void>
func setCurrentSignedPreKey(_ signedPreKey: SignedPreKeyRecord) -> Promise<Void>
}
class ServiceRestSocket: ServiceSocket {
@ -69,4 +70,21 @@ class ServiceRestSocket: ServiceSocket {
})
return promise
}
public func setCurrentSignedPreKey(_ signedPreKey: SignedPreKeyRecord) -> Promise<Void> {
let (promise, fulfill, reject) = Promise<Void>.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
}
}

View File

@ -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

View File

@ -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