Store devices that we need to restore session with.

Signal/src/ViewControllers/ConversationView/ConversationViewController.m

@@ -418,6 +418,10 @@ typedef enum : NSUInteger {
                                              selector:@selector(handleThreadFriendRequestStatusChangedNotification:)
                                                  name:NSNotification.threadFriendRequestStatusChanged
                                                object:nil];
+    [[NSNotificationCenter defaultCenter] addObserver:self
+                                             selector:@selector(handleThreadSessionRestoreDevicesChangedNotifiaction:)
+                                                 name:NSNotification.threadSessionRestoreDevicesChanged
+                                               object:nil];
 }
 
 - (BOOL)isGroupConversation
@@ -494,6 +498,17 @@ typedef enum : NSUInteger {
     [self resetContentAndLayout];
 }
 
+- (void)handleThreadSessionRestoreDevicesChangedNotifiaction:(NSNotification *)notification
+{
+    // Check thread
+    NSString *threadID = (NSString *)notification.object;
+    if (![threadID isEqualToString:self.thread.uniqueId]) { return; }
+    // Ensure thread instance is up to date
+    [self.thread reload];
+    // Update UI
+    // TODO: Show banner here
+}
+
 - (void)peekSetup
 {
     _peek = YES;

SignalServiceKit/src/Contacts/Threads/TSContactThread.h

@@ -22,6 +22,7 @@ extern NSString *const TSContactThreadPrefix;
 
 // Loki: The current session reset state for this thread
 @property (atomic) TSContactThreadSessionResetState sessionResetState;
+@property (atomic, readonly) NSArray<NSString *> *sessionRestoreDevices;
 
 @property (nonatomic) BOOL hasDismissedOffers;
 
@@ -47,6 +48,12 @@ extern NSString *const TSContactThreadPrefix;
 + (NSString *)conversationColorNameForRecipientId:(NSString *)recipientId
                                       transaction:(YapDatabaseReadTransaction *)transaction;
 
+#pragma mark - Loki Session Restore
+
+- (void)addSessionRestoreDevice:(NSString *)hexEncodedPublicKey transaction:(YapDatabaseReadWriteTransaction *_Nullable)transaction;
+
+- (void)removeAllRessionRestoreDevicesWithTransaction:(YapDatabaseReadWriteTransaction *_Nullable)transaction;
+
 @end
 
 NS_ASSUME_NONNULL_END

SignalServiceKit/src/Contacts/Threads/TSContactThread.m

@@ -8,6 +8,7 @@
 #import "NotificationsProtocol.h"
 #import "OWSIdentityManager.h"
 #import "SSKEnvironment.h"
+#import <SignalServiceKit/SignalServiceKit-Swift.h>
 #import <YapDatabase/YapDatabaseConnection.h>
 #import <YapDatabase/YapDatabaseTransaction.h>
 
@@ -26,6 +27,7 @@ NSString *const TSContactThreadPrefix = @"c";
     
     // No session reset ongoing
     _sessionResetState = TSContactThreadSessionResetStateNone;
+    _sessionRestoreDevices = @[];
 
     return self;
 }
@@ -107,6 +109,34 @@ NSString *const TSContactThreadPrefix = @"c";
     return [self stableColorNameForNewConversationWithString:recipientId];
 }
 
+#pragma mark - Loki Session Restore
+
+- (void)addSessionRestoreDevice:(NSString *)hexEncodedPublicKey transaction:(YapDatabaseReadWriteTransaction *_Nullable)transaction
+{
+    NSMutableSet *set = [[NSMutableSet alloc] initWithArray:_sessionRestoreDevices];
+    [set addObject:hexEncodedPublicKey];
+    [self setSessionRestoreDevices:[set allObjects] transaction:transaction];
+}
+
+- (void)removeAllRessionRestoreDevicesWithTransaction:(YapDatabaseReadWriteTransaction *_Nullable)transaction
+{
+    [self setSessionRestoreDevices:@[] transaction:transaction];
+}
+
+- (void)setSessionRestoreDevices:(NSArray<NSString *> * _Nonnull)sessionRestoreDevices transaction:(YapDatabaseReadWriteTransaction *_Nullable)transaction {
+    _sessionRestoreDevices = sessionRestoreDevices;
+     void (^postNotification)() = ^() {
+         [NSNotificationCenter.defaultCenter postNotificationName:NSNotification.threadSessionRestoreDevicesChanged object:self.uniqueId];
+    };
+    if (transaction == nil) {
+        [self save];
+        [self.dbReadWriteConnection flushTransactionsWithCompletionQueue:dispatch_get_main_queue() completionBlock:^{ postNotification(); }];
+    } else {
+        [self saveWithTransaction:transaction];
+        [transaction.connection flushTransactionsWithCompletionQueue:dispatch_get_main_queue() completionBlock:^{ postNotification(); }];
+    }
+}
+
 @end
 
 NS_ASSUME_NONNULL_END

SignalServiceKit/src/Loki/Messaging/Notification+Loki.swift

@@ -6,6 +6,7 @@ public extension Notification.Name {
     public static let messageFriendRequestStatusChanged = Notification.Name("messageFriendRequestStatusChanged")
     public static let threadDeleted = Notification.Name("threadDeleted")
     public static let dataNukeRequested = Notification.Name("dataNukeRequested")
+    public static let threadSessionRestoreDevicesChanged = Notification.Name("threadSessionRestoreDevicesChanged")
 }
 
 @objc public extension NSNotification {
@@ -15,4 +16,5 @@ public extension Notification.Name {
     @objc public static let messageFriendRequestStatusChanged = Notification.Name.messageFriendRequestStatusChanged.rawValue as NSString
     @objc public static let threadDeleted = Notification.Name.threadDeleted.rawValue as NSString
     @objc public static let dataNukeRequested = Notification.Name.dataNukeRequested.rawValue as NSString
+    @objc public static let threadSessionRestoreDevicesChanged = Notification.Name.threadSessionRestoreDevicesChanged.rawValue as NSString
 }

SignalServiceKit/src/Messages/OWSMessageDecrypter.m

@@ -678,16 +678,33 @@ NSError *EnsureDecryptError(NSError *_Nullable error, NSString *fallbackErrorDes
         OWSAssertDebug(errorMessage);
         if (errorMessage != nil) {
             [errorMessage saveWithTransaction:transaction];
+            [self handleSessionRestoreForErrorMessage:errorMessage envelope:envelope transaction:transaction];
             [self notifyUserForErrorMessage:errorMessage envelope:envelope transaction:transaction];
         }
     }];
 }
 
+- (void)handleSessionRestoreForErrorMessage:(TSErrorMessage *)errorMessage
+                                   envelope:(SSKProtoEnvelope *)envelope
+                                   transaction:(YapDatabaseReadWriteTransaction *)transaction
+{
+    NSString *masterHexEncodedPublicKey = [LKDatabaseUtilities getMasterHexEncodedPublicKeyFor:envelope.source in:transaction];
+    TSThread *contactThread = [TSContactThread getOrCreateThreadWithContactId:masterHexEncodedPublicKey transaction:transaction];
+       
+   // Trigger a session restore prompt if we get specific errors
+   if (errorMessage.errorType == TSErrorMessageNoSession ||
+       errorMessage.errorType == TSErrorMessageInvalidMessage ||
+       errorMessage.errorType == TSErrorMessageInvalidKeyException) {
+       [((TSContactThread *) contactThread) addSessionRestoreDevice:masterHexEncodedPublicKey transaction:transaction];
+   }
+}
+
 - (void)notifyUserForErrorMessage:(TSErrorMessage *)errorMessage
                          envelope:(SSKProtoEnvelope *)envelope
                       transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
-    TSThread *contactThread = [TSContactThread getOrCreateThreadWithContactId:envelope.source transaction:transaction];
+    NSString *masterHexEncodedPublicKey = [LKDatabaseUtilities getMasterHexEncodedPublicKeyFor:envelope.source in:transaction];
+    TSThread *contactThread = [TSContactThread getOrCreateThreadWithContactId:masterHexEncodedPublicKey transaction:transaction];
     [SSKEnvironment.shared.notificationsManager notifyUserForErrorMessage:errorMessage
                                                                    thread:contactThread
                                                               transaction:transaction];

SignalServiceKit/src/Messages/OWSMessageSender.m

@@ -1316,6 +1316,9 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException";
                                 // Update the thread
                                 [message.thread saveFriendRequestStatus:LKThreadFriendRequestStatusRequestSent withTransaction:transaction];
                                 [message.thread removeOldOutgoingFriendRequestMessagesIfNeededWithTransaction:transaction];
+                                if ([message.thread isKindOfClass:[TSContactThread class]]) {
+                                    [((TSContactThread *) message.thread) removeAllRessionRestoreDevicesWithTransaction:transaction];
+                                }
                                 // Update the message
                                 [message saveFriendRequestStatus:LKMessageFriendRequestStatusPending withTransaction:transaction];
                                 NSTimeInterval expirationInterval = 72 * kHourInterval;