From 619597cd6102ab97baf031d5af3034eae4a57495 Mon Sep 17 00:00:00 2001
From: Michael Kirk <michael@signal.org>
Date: Wed, 5 Sep 2018 17:27:26 -0600
Subject: [PATCH] ensure operations run to completion on the PreKey operation
 queue

---
 .../src/Account/TSPreKeyManager.m             | 54 +++++++------------
 SignalServiceKit/src/Util/OWSOperation.m      | 15 ++++--
 2 files changed, 30 insertions(+), 39 deletions(-)

diff --git a/SignalServiceKit/src/Account/TSPreKeyManager.m b/SignalServiceKit/src/Account/TSPreKeyManager.m
index 182ad8491..4a2cb009c 100644
--- a/SignalServiceKit/src/Account/TSPreKeyManager.m
+++ b/SignalServiceKit/src/Account/TSPreKeyManager.m
@@ -67,17 +67,6 @@ static const NSUInteger kMaxPrekeyUpdateFailureCount = 5;
     [primaryStorage clearPrekeyUpdateFailureCount];
 }
 
-// We should never dispatch sync to this queue.
-+ (dispatch_queue_t)prekeyQueue
-{
-    static dispatch_once_t onceToken;
-    static dispatch_queue_t queue;
-    dispatch_once(&onceToken, ^{
-        queue = dispatch_queue_create("org.whispersystems.signal.prekeyQueue", NULL);
-    });
-    return queue;
-}
-
 + (NSOperationQueue *)operationQueue
 {
     static dispatch_once_t onceToken;
@@ -100,41 +89,36 @@ static const NSUInteger kMaxPrekeyUpdateFailureCount = 5;
         return;
     }
 
-    NSBlockOperation *checkIfPreKeyRefreshNecessaryOperation = [NSBlockOperation blockOperationWithBlock:^{
+    SSKRefreshPreKeysOperation *refreshOperation = [SSKRefreshPreKeysOperation new];
+    NSBlockOperation *checkIfRefreshNecessaryOperation = [NSBlockOperation blockOperationWithBlock:^{
         BOOL shouldCheck = (lastPreKeyCheckTimestamp == nil
-            || fabs([lastPreKeyCheckTimestamp timeIntervalSinceNow]) >= kPreKeyCheckFrequencySeconds);
+                            || fabs([lastPreKeyCheckTimestamp timeIntervalSinceNow]) >= kPreKeyCheckFrequencySeconds);
         if (!shouldCheck) {
-            return;
+            [refreshOperation cancel];
         }
-
-        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:^{
+    [refreshOperation addDependency:checkIfRefreshNecessaryOperation];
+    
+    SSKRotateSignedPreKeyOperation *rotationOperation = [SSKRotateSignedPreKeyOperation new];
+    NSBlockOperation *checkIfRotationNecessaryOperation = [NSBlockOperation blockOperationWithBlock:^{
         OWSPrimaryStorage *primaryStorage = [OWSPrimaryStorage sharedManager];
         SignedPreKeyRecord *_Nullable signedPreKey = [primaryStorage currentSignedPreKey];
-
+        
         BOOL shouldCheck
-            = !signedPreKey || fabs(signedPreKey.generatedAt.timeIntervalSinceNow) >= kSignedPreKeyRotationTime;
+        = !signedPreKey || fabs(signedPreKey.generatedAt.timeIntervalSinceNow) >= kSignedPreKeyRotationTime;
         if (!shouldCheck) {
-            return;
+            [rotationOperation cancel];
         }
-
-        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];
     }];
+    [rotationOperation addDependency:checkIfRotationNecessaryOperation];
+
+    // Order matters here - if we rotated *before* refreshing, we'd risk uploading
+    // two SPK's in a row since RefreshPreKeysOperation can also upload a new SPK.
+    [checkIfRotationNecessaryOperation addDependency:refreshOperation];
 
-    // 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];
+    NSArray<NSOperation *> *operations =
+        @[ checkIfRefreshNecessaryOperation, refreshOperation, checkIfRotationNecessaryOperation, rotationOperation ];
+    [self.operationQueue addOperations:operations waitUntilFinished:NO];
 }
 
 + (void)createPreKeysWithSuccess:(void (^)(void))successHandler failure:(void (^)(NSError *error))failureHandler
diff --git a/SignalServiceKit/src/Util/OWSOperation.m b/SignalServiceKit/src/Util/OWSOperation.m
index 4ebd055a4..ae16d8b59 100644
--- a/SignalServiceKit/src/Util/OWSOperation.m
+++ b/SignalServiceKit/src/Util/OWSOperation.m
@@ -50,10 +50,12 @@ NSString *const OWSOperationKeyIsFinished = @"isFinished";
 {
     for (NSOperation *dependency in self.dependencies) {
         if (![dependency isKindOfClass:[OWSOperation class]]) {
-            NSString *errorDescription =
-                [NSString stringWithFormat:@"%@ unknown dependency: %@", self.logTag, dependency.class];
-
-            return OWSErrorMakeAssertionError(errorDescription);
+            // If you want an operation to have a cascading failure, it must
+            // subclass OWSOperation.
+            //
+            // Native operations like NSOperation or NSBlockOperation don't
+            // don't support this feature.
+            continue;
         }
 
         OWSOperation *dependentOperation = (OWSOperation *)dependency;
@@ -107,6 +109,11 @@ NSString *const OWSOperationKeyIsFinished = @"isFinished";
         return;
     }
 
+    if (self.isCancelled) {
+        [self reportCancelled];
+        return;
+    }
+    
     [self run];
 }