Merge branch 'charlesmchen/unregisterVsSAE'

SignalServiceKit/src/Account/TSAccountManager.m

@@ -3,6 +3,7 @@
 //
 
 #import "TSAccountManager.h"
+#import "AppContext.h"
 #import "NSData+Base64.h"
 #import "NSData+hexString.h"
 #import "NSNotificationCenter+OWS.h"
@@ -58,9 +59,21 @@ NSString *const TSAccountManager_ServerSignalingKey = @"TSStorageServerSignaling
 
     OWSSingletonAssert();
 
+    if (!CurrentAppContext().isMainApp) {
+        [[NSNotificationCenter defaultCenter] addObserver:self
+                                                 selector:@selector(yapDatabaseModifiedExternally:)
+                                                     name:YapDatabaseModifiedExternallyNotification
+                                                   object:nil];
+    }
+
     return self;
 }
 
+- (void)dealloc
+{
+    [[NSNotificationCenter defaultCenter] removeObserver:self];
+}
+
 + (instancetype)sharedInstance
 {
     static dispatch_once_t onceToken;
@@ -454,6 +467,21 @@ NSString *const TSAccountManager_ServerSignalingKey = @"TSStorageServerSignaling
         }];
 }
 
+- (void)yapDatabaseModifiedExternally:(NSNotification *)notification
+{
+    OWSAssertIsOnMainThread();
+
+    DDLogVerbose(@"%@ %s", self.logTag, __PRETTY_FUNCTION__);
+
+    // Any database write by the main app might reflect a deregistration,
+    // so clear the cached "is registered" state.  This will significantly
+    // erode the value of this cache in the SAE.
+    @synchronized(self)
+    {
+        _isRegistered = NO;
+    }
+}
+
 @end
 
 NS_ASSUME_NONNULL_END

SignalShareExtension/ShareViewController.swift

@@ -1,5 +1,5 @@
 //
-//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
 //
 
 import UIKit
@@ -12,6 +12,12 @@ import PromiseKit
 @objc
 public class ShareViewController: UINavigationController, ShareViewDelegate, SAEFailedViewDelegate {
 
+    enum ShareViewControllerError: Error {
+        case assertionError(description: String)
+        case unsupportedMedia
+        case notRegistered()
+    }
+
     private var hasInitialRootViewController = false
     private var isReadyForAppExtensions = false
 
@@ -109,6 +115,10 @@ public class ShareViewController: UINavigationController, ShareViewDelegate, SAE
                                                selector: #selector(registrationStateDidChange),
                                                name: .RegistrationStateDidChange,
                                                object: nil)
+        NotificationCenter.default.addObserver(self,
+                                               selector: #selector(owsApplicationWillEnterForeground),
+                                               name: .OWSApplicationWillEnterForeground,
+                                               object: nil)
 
         Logger.info("\(self.logTag) application: didFinishLaunchingWithOptions completed.")
 
@@ -342,6 +352,29 @@ public class ShareViewController: UINavigationController, ShareViewDelegate, SAE
         Logger.flush()
     }
 
+    @objc
+    func owsApplicationWillEnterForeground() throws {
+        AssertIsOnMainThread()
+
+        Logger.debug("\(self.logTag) \(#function)")
+
+        // If a user unregisters in the main app, the SAE should shut down
+        // immediately.
+        guard !TSAccountManager.isRegistered() else {
+            // If user is registered, do nothing.
+            return
+        }
+        guard let firstViewController = shareViewNavigationController.viewControllers.first else {
+            // If no view has been presented yet, do nothing.
+            return
+        }
+        if let _ = firstViewController as? SAEFailedViewController {
+            // If root view is an error view, do nothing.
+            return
+        }
+        throw ShareViewControllerError.notRegistered()
+    }
+
     // MARK: ShareViewDelegate, SAEFailedViewDelegate
 
     public func shareViewWasCompleted() {
@@ -408,12 +441,6 @@ public class ShareViewController: UINavigationController, ShareViewDelegate, SAE
         }.retainUntilComplete()
     }
 
-    enum ShareViewControllerError: Error {
-        case assertionError(description: String)
-        case unsupportedMedia
-
-    }
-
     private func buildAttachment() -> Promise<SignalAttachment> {
         guard let inputItem: NSExtensionItem = self.extensionContext?.inputItems.first as? NSExtensionItem else {
             let error = ShareViewControllerError.assertionError(description: "no input item")