Loki session reset (#14)

* Added session reset.

* Hooked up session reset internals to UI.

* Send empty message when we have received an end session message.

* Verify incoming PreKeyWhisperMessage.

* Fix indentations in SessionReset.md

Pods

@@ -1 +1 @@
-Subproject commit 7f42f93c7df8127331d26d0109170a0524f67f7b
+Subproject commit 870d1b5be23fd8fb5d68af6c20e36b3ed5dcde0f

Signal/src/Jobs/SessionResetJob.swift

@@ -108,6 +108,9 @@ public class SessionResetOperation: OWSOperation, DurableOperation {
     override public func run() {
         assert(self.durableOperationDelegate != nil)
 
+        /* Loki Original Code
+         * We don't want to delete session. Ref: SignalServiceKit/Loki/Docs/SessionReset.md
+         * ==================
         if firstAttempt {
             self.dbConnection.readWrite { transaction in
                 Logger.info("deleting sessions for recipient: \(self.recipientId)")
@@ -115,6 +118,7 @@ public class SessionResetOperation: OWSOperation, DurableOperation {
             }
             firstAttempt = false
         }
+        */
 
         let endSessionMessage = EndSessionMessage(timestamp: NSDate.ows_millisecondTimeStamp(), in: self.contactThread)
 
@@ -128,10 +132,25 @@ public class SessionResetOperation: OWSOperation, DurableOperation {
                 // Otherwise if we send another message before them, they wont have the session to decrypt it.
                 self.primaryStorage.archiveAllSessions(forContact: self.recipientId, protocolContext: transaction)
 
+                /* Loki original code
+                 * ==================
                 let message = TSInfoMessage(timestamp: NSDate.ows_millisecondTimeStamp(),
                                             in: self.contactThread,
                                             messageType: TSInfoMessageType.typeSessionDidEnd)
                 message.save(with: transaction)
+                */
+                
+                if (self.contactThread.sessionResetState != .requestReceived) {
+                    let message = TSInfoMessage(timestamp: NSDate.ows_millisecondTimeStamp(),
+                                                in: self.contactThread,
+                                                messageType: .typeLokiSessionResetProgress)
+                    message.save(with: transaction)
+                    
+                    /// Loki: We have initiated a session reset
+                    Logger.debug("[Loki Session Reset] Session reset has been initiated")
+                    self.contactThread.sessionResetState = .initiated
+                    self.contactThread.save(with: transaction)
+                }
             }
             self.reportSuccess()
         }.catch { error in

Signal/src/ViewControllers/ConversationView/Cells/OWSSystemMessageCell.m

@@ -273,6 +273,8 @@ typedef void (^SystemMessageActionBlock)(void);
             case TSInfoMessageAddGroupToProfileWhitelistOffer:
             case TSInfoMessageTypeGroupUpdate:
             case TSInfoMessageTypeGroupQuit:
+            case TSInfoMessageTypeLokiSessionResetProgress:
+            case TSInfoMessageTypeLokiSessionResetDone:
                 return nil;
             case TSInfoMessageTypeDisappearingMessagesUpdate: {
                 BOOL areDisappearingMessagesEnabled = YES;
@@ -459,6 +461,8 @@ typedef void (^SystemMessageActionBlock)(void);
     switch (message.messageType) {
         case TSInfoMessageUserNotRegistered:
         case TSInfoMessageTypeSessionDidEnd:
+        case TSInfoMessageTypeLokiSessionResetProgress:
+        case TSInfoMessageTypeLokiSessionResetDone:
             return nil;
         case TSInfoMessageTypeUnsupportedMessage:
             // Unused.

Signal/translations/en.lproj/Localizable.strings

@@ -2576,6 +2576,8 @@
 "Decline" = "Decline";
 "Pending Friend Request..." = "Pending Friend Request...";
 "New Message" = "New Message";
+"Secure session reset in progress" = "Secure session reset in progress";
+"Secure session reset done" = "Secure session reset done";
 "Session" = "Session";
 "You've declined %@'s friend request" = "You've declined %@'s friend request";
 "You've accepted %@'s friend request" = "You've accepted %@'s friend request";

SignalServiceKit/src/Contacts/Threads/TSContactThread.h

@@ -6,10 +6,23 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+// Loki: Session reset state
+typedef NS_ENUM(NSInteger, TSContactThreadSessionResetState) {
+    // No ongoing session reset
+    TSContactThreadSessionResetStateNone,
+    // We initiated session reset
+    TSContactThreadSessionResetStateInitiated,
+    // We received the session reset
+    TSContactThreadSessionResetStateRequestReceived,
+};
+
 extern NSString *const TSContactThreadPrefix;
 
 @interface TSContactThread : TSThread
 
+// Loki: The current session reset state with this thread
+@property (atomic) TSContactThreadSessionResetState sessionResetState;
+
 @property (nonatomic) BOOL hasDismissedOffers;
 
 + (instancetype)getOrCreateThreadWithContactId:(NSString *)contactId NS_SWIFT_NAME(getOrCreateThread(contactId:));

SignalServiceKit/src/Contacts/Threads/TSContactThread.m

@@ -23,6 +23,9 @@ NSString *const TSContactThreadPrefix = @"c";
     OWSAssertDebug(contactId.length > 0);
 
     self = [super initWithUniqueId:uniqueIdentifier];
+    
+    // No session reset ongoing
+    _sessionResetState = TSContactThreadSessionResetStateNone;
 
     return self;
 }

SignalServiceKit/src/Loki/Docs/SessionReset.md

@@ -0,0 +1,56 @@
+# Loki Session Reset
+
+## Signal
+Since Signal uses a centralised server, creating sessions is easy as the prekeys can be easily fetched.
+
+The process is as follows:
+
+1. `A` deletes all their sessions and sends `End Session` to `B`
+    - `A` contacts the server and creates a new session
+2. `B` Gets this message and deletes all sessions.
+3. `B` Sends a message with a newly created session
+    - `B` contacted server and established this
+4. `A` and `B` now have the same sessions so they can delete any archived ones.
+
+## Loki
+Loki doesn't have a centralised server and thus we need to change the process above with something similar. 
+
+We have to introduce a session reset state `sessionState` which can take the following states:
+- `none`: No session reset is in progress
+- `initiated`: We have initiated the session reset
+- `received`: We have received a session reset from the other user
+
+The new process is as follows:
+
+1. `A` Sends `End Session` with a `PreKeyBundle` and archives its own session.
+    - `sessionState = initiated`
+    - The session is archived as we could get a message from `B` using the archived session, so we still want to be able to decrypt that.
+    - We can show `Session reset in progress`
+2. `B` Gets this message and saves the `PreKeyBundle` and archives its own sessions.
+    - `sessionState = received`
+    - `B` sends an empty message, which will trigger a new session to be created.
+    - `B` deletes the `PreKeyBundle` once session is created.
+    - We can show `Session reset in progress`
+3. `A` and `B` both do the routine below when receiving messages.
+
+### Upon receiving message (Only applies to PreKey and Cipher messages)
+
+- Store the current active session `PS`
+- Decrypt the message
+    - Decrypting a message can cause the active session to change
+- If `sessionState == none` then it means that we haven't started session reset and we can abort.
+- Get the current session `CS`
+- If `PS` is `nil` then abort as we didn't have a session before.
+- If `CS != PS` then sessions were changed.
+    - If `sessionState == received` then it means that the sender used an old session to contact us. We need to wait for them to use the new one.
+        - Archive `CS` and set the session to `PS`
+    - If `sessionState == initiated` then it means that the sender acknowledged our session reset and sent a message with a new session
+        - Delete all session except `CS`
+        - `sessionState = none`
+        - Send an empty message to confirm session adoption
+        - We can show `Session reset done`
+- If `CS == PS` then sessions were the same.
+    - If `sessionState == received` then it means that the new session we created is the one the sender used for sending message. We have successfully adopted the new session.
+    - Delete all sessions except `PS`
+    - `sessionState = none`
+    - We can show `Session reset done`

SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.h

@@ -17,6 +17,15 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (BOOL)hasPreKeyForContact:(NSString *)pubKey;
 
+/**
+ Get the `PreKeyRecord` associated with the given contact.
+ 
+ @param pubKey The hex encoded public key of the contact.
+ @param transaction A `YapDatabaseReadTransaction`.
+ @return The record associated with the contact or nil if it didn't exist.
+ */
+- (PreKeyRecord *_Nullable)getPreKeyForContact:(NSString *)pubKey transaction:(YapDatabaseReadTransaction *)transaction;
+
 /**
  Get the `PreKeyRecord` associated with the given contact.
  If the record doesn't exist then this will create a new one.

SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.m

@@ -1,13 +1,16 @@
 #import "OWSPrimaryStorage+Loki.h"
 #import "OWSPrimaryStorage+PreKeyStore.h"
 #import "OWSPrimaryStorage+SignedPreKeyStore.h"
+#import "OWSPrimaryStorage+keyFromIntLong.h"
 #import "OWSDevice.h"
 #import "OWSIdentityManager.h"
 #import "TSAccountManager.h"
 #import "TSPreKeyManager.h"
 #import "YapDatabaseConnection+OWS.h"
+#import "YapDatabaseTransaction+OWS.h"
 #import <AxolotlKit/NSData+keyVersionByte.h>
 
+#define OWSPrimaryStoragePreKeyStoreCollection @"TSStorageManagerPreKeyStoreCollection"
 #define LokiPreKeyContactCollection @"LokiPreKeyContactCollection"
 #define LokiPreKeyBundleCollection @"LokiPreKeyBundleCollection"
 
@@ -30,6 +33,20 @@
     return preKeyId > 0;
 }
 
+- (PreKeyRecord *_Nullable)getPreKeyForContact:(NSString *)pubKey transaction:(YapDatabaseReadTransaction *)transaction {
+    OWSAssertDebug(pubKey.length > 0);
+    int preKeyId = [transaction intForKey:pubKey inCollection:LokiPreKeyContactCollection];
+    
+    // If we don't have an id then return nil
+    if (preKeyId <= 0) {
+        return nil;
+    }
+    
+    /// thows_loadPreKey doesn't allow us to pass transaction ;(
+    return [transaction preKeyRecordForKey:[self keyFromInt:preKeyId]
+                              inCollection:OWSPrimaryStoragePreKeyStoreCollection];
+}
+
 - (PreKeyRecord *)getOrCreatePreKeyForContact:(NSString *)pubKey {
     OWSAssertDebug(pubKey.length > 0);
     int preKeyId = [self.dbReadWriteConnection intForKey:pubKey inCollection:LokiPreKeyContactCollection];

SignalServiceKit/src/Loki/Extensions/SessionCipher+Loki.h

@@ -0,0 +1,25 @@
+/// Loki: Refer to Docs/SessionReset.md for explanations
+
+#import "SessionCipher.h"
+
+NS_ASSUME_NONNULL_BEGIN
+
+extern NSString *const kNSNotificationName_SessionAdopted;
+extern NSString *const kNSNotificationKey_ContactPubKey;
+
+@interface SessionCipher (Loki)
+
+/**
+ Decrypt the given `CipherMessage`.
+ This function is a wrapper around `throws_decrypt:protocolContext:` and adds on the custom loki session handling ontop.
+ Refer to SignalServiceKit/Loki/Docs/SessionReset.md for overview on how it works.
+
+ @param whisperMessage The cipher message.
+ @param protocolContext The protocol context (YapDatabaseReadWriteTransaction)
+ @return The decrypted data.
+ */
+- (NSData *)throws_lokiDecrypt:(id<CipherMessage>)whisperMessage protocolContext:(nullable id)protocolContext NS_SWIFT_UNAVAILABLE("throws objc exceptions");
+
+@end
+
+NS_ASSUME_NONNULL_END

SignalServiceKit/src/Loki/Extensions/SessionCipher+Loki.m

@@ -0,0 +1,181 @@
+/// Loki: Refer to Docs/SessionReset.md for explanations
+
+#import "SessionCipher+Loki.h"
+#import "NSNotificationCenter+OWS.h"
+#import "PreKeyWhisperMessage.h"
+#import "OWSPrimaryStorage+Loki.h"
+#import "TSContactThread.h"
+#import <YapDatabase/YapDatabase.h>
+
+NSString *const kNSNotificationName_SessionAdopted = @"kNSNotificationName_SessionAdopted";
+NSString *const kNSNotificationKey_ContactPubKey = @"kNSNotificationKey_ContactPubKey";
+
+@interface SessionCipher ()
+
+@property (nonatomic, readonly) NSString *recipientId;
+@property (nonatomic, readonly) int deviceId;
+
+@property (nonatomic, readonly) id<SessionStore> sessionStore;
+@property (nonatomic, readonly) id<PreKeyStore> prekeyStore;
+
+@end
+
+@implementation SessionCipher (Loki)
+
+- (NSData *)throws_lokiDecrypt:(id<CipherMessage>)whisperMessage protocolContext:(nullable id)protocolContext
+{
+    // Our state before we decrypt the message
+    SessionState *_Nullable state = [self getCurrentState:protocolContext];
+    
+    // While decrypting our state may change internally
+    NSData *plainText = [self throws_decrypt:whisperMessage protocolContext:protocolContext];
+    
+    // Loki: Verify incoming friend request messages
+    if (!state) {
+        [self throws_verifyFriendRequestAcceptPreKeyForMessage:whisperMessage protocolContext:protocolContext];
+    }
+   
+    // Loki: Handle any session resets
+    [self handleSessionReset:whisperMessage previousState:state protocolContext:protocolContext];
+    
+    return plainText;
+}
+
+/// Get the current session state
+- (SessionState *_Nullable)getCurrentState:(nullable id)protocolContext {
+    SessionRecord *record = [self.sessionStore loadSession:self.recipientId deviceId:self.deviceId protocolContext:protocolContext];
+    SessionState *state = record.sessionState;
+    
+    // Check if session is initialized
+    if (!state.hasSenderChain) {
+        return nil;
+    }
+    
+    return state;
+}
+
+/// Handle any loki session reset stuff
+- (void)handleSessionReset:(id<CipherMessage>)whisperMessage
+             previousState:(SessionState *_Nullable)previousState
+           protocolContext:(nullable id)protocolContext
+{
+    // Don't bother doing anything if we didn't have a session before
+    if (!previousState) {
+        return;
+    }
+    
+    OWSAssertDebug([protocolContext isKindOfClass:[YapDatabaseReadWriteTransaction class]]);
+    YapDatabaseReadWriteTransaction *transaction = protocolContext;
+    
+    // Get the thread
+    TSContactThread *thread = [TSContactThread getThreadWithContactId:self.recipientId transaction:transaction];
+    if (!thread) {
+        return;
+    }
+    
+    // Bail early if no session reset is in progress
+    if (thread.sessionResetState == TSContactThreadSessionResetStateNone) {
+        return;
+    }
+    
+    BOOL sessionResetReceived = thread.sessionResetState == TSContactThreadSessionResetStateRequestReceived;
+    SessionState *_Nullable currentState = [self getCurrentState:protocolContext];
+    
+    // Check if our previous state and our current state differ
+    if (!currentState || ![currentState.aliceBaseKey isEqualToData:previousState.aliceBaseKey]) {
+        
+        if (sessionResetReceived) {
+            // The other user used an old session to contact us.
+            // Wait for them to use a new one
+            [self restoreSession:previousState protocolContext:protocolContext];
+        } else {
+            // Our session reset went through successfully
+            // We had initiated a session reset and got a different session back from the user
+            [self deleteAllSessionsExcept:currentState protocolContext:protocolContext];
+            [self notifySessionAdopted];
+        }
+        
+    } else if (sessionResetReceived) {
+        // Our session reset went through successfully
+        // We got a message with the same session from the other user
+        [self deleteAllSessionsExcept:previousState protocolContext:protocolContext];
+        [self notifySessionAdopted];
+    }
+}
+
+/// Send a notification about a new session being adopted
+- (void)notifySessionAdopted
+{
+    [[NSNotificationCenter defaultCenter]
+     postNotificationNameAsync:kNSNotificationName_SessionAdopted
+     object:nil
+     userInfo:@{
+                kNSNotificationKey_ContactPubKey : self.recipientId,
+                }];
+}
+
+/// Delete all other sessions except the given one
+- (void)deleteAllSessionsExcept:(SessionState *)state protocolContext:(nullable id)protocolContext
+{
+    SessionRecord *record = [self.sessionStore loadSession:self.recipientId deviceId:self.deviceId protocolContext:protocolContext];
+    [record removePreviousSessionStates];
+    [record setState:state];
+    
+    [self.sessionStore storeSession:self.recipientId
+                           deviceId:self.deviceId
+                            session:record
+                    protocolContext:protocolContext];
+}
+
+/// Set the given session as the active one while archiving the old one
+- (void)restoreSession:(SessionState *)state protocolContext:(nullable id)protocolContext
+{
+    SessionRecord *record = [self.sessionStore loadSession:self.recipientId deviceId:self.deviceId protocolContext:protocolContext];
+    
+    // Remove the state from previous session states
+    [record.previousSessionStates enumerateObjectsWithOptions:NSEnumerationReverse usingBlock:^(SessionState *obj, NSUInteger idx, BOOL *stop) {
+        if ([state.aliceBaseKey isEqualToData:obj.aliceBaseKey]) {
+            [record.previousSessionStates removeObjectAtIndex:idx];
+            *stop = true;
+        }
+    }];
+    
+    // Promote it so the previous state gets archived
+    [record promoteState:state];
+    
+    [self.sessionStore storeSession:self.recipientId
+                           deviceId:self.deviceId
+                            session:record
+                    protocolContext:protocolContext];
+}
+
+/// Check that we have matching prekeys in the case of a `PreKeyWhisperMessage`
+/// This is so that we don't trigger a false friend request accept on unknown contacts
+- (void)throws_verifyFriendRequestAcceptPreKeyForMessage:(id<CipherMessage>)whisperMessage protocolContext:(nullable id)protocolContext {
+    OWSAssertDebug([protocolContext isKindOfClass:[YapDatabaseReadTransaction class]]);
+    YapDatabaseReadTransaction *transaction = protocolContext;
+    
+    /// We only want to look at `PreKeyWhisperMessage`
+    if (![whisperMessage isKindOfClass:[PreKeyWhisperMessage class]]) {
+        return;
+    }
+
+    /// We need the primary storage to access contact prekeys
+    if (![self.prekeyStore isKindOfClass:[OWSPrimaryStorage class]]) {
+        return;
+    }
+    
+    PreKeyWhisperMessage *preKeyMessage = whisperMessage;
+    OWSPrimaryStorage *primaryStorage = self.prekeyStore;
+    
+    PreKeyRecord *_Nullable storedPreKey = [primaryStorage getPreKeyForContact:self.recipientId transaction:transaction];
+    if(!storedPreKey) {
+        OWSRaiseException(@"LokiInvalidPreKey", @"Received a friend request from a pubkey for which no prekey bundle was created");
+    }
+    
+    if (storedPreKey.Id != preKeyMessage.prekeyID) {
+        OWSRaiseException(@"LokiPreKeyIdsDontMatch", @"Received a preKeyWhisperMessage (friend request accept) from an unknown source");
+    }
+}
+
+@end

SignalServiceKit/src/Messages/Interactions/TSInfoMessage.h

@@ -21,6 +21,8 @@ typedef NS_ENUM(NSInteger, TSInfoMessageType) {
     TSInfoMessageVerificationStateChange,
     TSInfoMessageAddUserToProfileWhitelistOffer,
     TSInfoMessageAddGroupToProfileWhitelistOffer,
+    TSInfoMessageTypeLokiSessionResetProgress,
+    TSInfoMessageTypeLokiSessionResetDone,
 };
 
 + (instancetype)userNotRegisteredMessageInThread:(TSThread *)thread recipientId:(NSString *)recipientId;

SignalServiceKit/src/Messages/Interactions/TSInfoMessage.m

@@ -117,6 +117,10 @@ NSUInteger TSInfoMessageSchemaVersion = 1;
 - (NSString *)previewTextWithTransaction:(YapDatabaseReadTransaction *)transaction
 {
     switch (_messageType) {
+        case TSInfoMessageTypeLokiSessionResetProgress:
+            return NSLocalizedString(@"Secure session reset in progress", nil);
+        case TSInfoMessageTypeLokiSessionResetDone:
+            return NSLocalizedString(@"Secure session reset done", nil);
         case TSInfoMessageTypeSessionDidEnd:
             return NSLocalizedString(@"SECURE_SESSION_RESET", nil);
         case TSInfoMessageTypeUnsupportedMessage:

SignalServiceKit/src/Messages/OWSMessageDecrypter.m

@@ -15,6 +15,7 @@
 #import "OWSPrimaryStorage+SessionStore.h"
 #import "OWSPrimaryStorage+SignedPreKeyStore.h"
 #import "OWSPrimaryStorage.h"
+#import "SessionCipher+Loki.h"
 #import "SSKEnvironment.h"
 #import "SignalRecipient.h"
 #import "TSAccountManager.h"
@@ -434,8 +435,14 @@ NSError *EnsureDecryptError(NSError *_Nullable error, NSString *fallbackErrorDes
                                                                            deviceId:deviceId];
 
                 // plaintextData may be nil for some envelope types.
+                NSData *_Nullable plaintextData =
+                    [[cipher throws_lokiDecrypt:cipherMessage protocolContext:transaction] removePadding];
+                
+                /* Loki original code
+                 * =================
                 NSData *_Nullable plaintextData =
                     [[cipher throws_decrypt:cipherMessage protocolContext:transaction] removePadding];
+                 */
                 OWSMessageDecryptResult *result = [OWSMessageDecryptResult resultWithEnvelopeData:envelopeData
                                                                                     plaintextData:plaintextData
                                                                                            source:envelope.source

SignalServiceKit/src/Messages/OWSMessageManager.m

@@ -33,6 +33,7 @@
 #import "OWSSyncGroupsMessage.h"
 #import "OWSSyncGroupsRequestMessage.h"
 #import "ProfileManagerProtocol.h"
+#import "SessionCipher+Loki.h"
 #import "SSKEnvironment.h"
 #import "TSAccountManager.h"
 #import "TSAttachment.h"
@@ -85,12 +86,22 @@ NS_ASSUME_NONNULL_BEGIN
     _primaryStorage = primaryStorage;
     _dbConnection = primaryStorage.newDatabaseConnection;
     _incomingMessageFinder = [[OWSIncomingMessageFinder alloc] initWithPrimaryStorage:primaryStorage];
+    
+    /// Loki: Add observation for new session
+    [[NSNotificationCenter defaultCenter] addObserver:self
+                                             selector:@selector(onNewSessionAdopted:)
+                                                 name:kNSNotificationName_SessionAdopted
+                                               object:nil];
 
     OWSSingletonAssert();
 
     return self;
 }
 
+- (void)dealloc {
+    [[NSNotificationCenter defaultCenter] removeObserver:self];
+}
+
 #pragma mark - Dependencies
 
 - (id<OWSCallMessageHandler>)callMessageHandler
@@ -992,11 +1003,34 @@ NS_ASSUME_NONNULL_BEGIN
     TSContactThread *thread = [TSContactThread getOrCreateThreadWithContactId:envelope.source transaction:transaction];
 
     // MJK TODO - safe to remove senderTimestamp
+    [[[TSInfoMessage alloc] initWithTimestamp:[NSDate ows_millisecondTimeStamp]
+                                     inThread:thread
+                                  messageType:TSInfoMessageTypeLokiSessionResetProgress] saveWithTransaction:transaction];
+    /* Loki original code
+     * ==================
     [[[TSInfoMessage alloc] initWithTimestamp:[NSDate ows_millisecondTimeStamp]
                                      inThread:thread
                                   messageType:TSInfoMessageTypeSessionDidEnd] saveWithTransaction:transaction];
+     */
 
+    /// Loki: Archive all our sessions
+    /// Ref: SignalServiceKit/Loki/Docs/SessionReset.md
+    [self.primaryStorage archiveAllSessionsForContact:envelope.source protocolContext:transaction];
+    
+    /// Loki: Set our session reset state
+    thread.sessionResetState = TSContactThreadSessionResetStateRequestReceived;
+    [thread saveWithTransaction:transaction];
+    
+    /// Loki: Send an empty message to trigger the session reset code for both parties
+    TSOutgoingMessage *emptyMessage = [TSOutgoingMessage createEmptyOutgoingMessageInThread:thread];
+    [self.messageSenderJobQueue addMessage:emptyMessage transaction:transaction];
+
+    OWSLogDebug(@"[Loki Session Reset] Session reset has been received from %@", envelope.source);
+    
+    /* Loki Original Code
+     * ===================
     [self.primaryStorage deleteAllSessionsForContact:envelope.source protocolContext:transaction];
+     */
 }
 
 - (void)handleExpirationTimerUpdateMessageWithEnvelope:(SSKProtoEnvelope *)envelope
@@ -1623,6 +1657,39 @@ NS_ASSUME_NONNULL_BEGIN
     }
 }
 
+# pragma mark - Loki Session
+
+- (void)onNewSessionAdopted:(NSNotification *)notification {
+    NSString *pubKey = notification.userInfo[kNSNotificationKey_ContactPubKey];
+    if (pubKey.length == 0) {
+        return;
+    }
+    
+    [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
+        
+        TSContactThread *_Nullable thread = [TSContactThread getThreadWithContactId:pubKey transaction:transaction];
+        if (!thread) {
+            OWSLogDebug(@"[Loki Session Reset] New session was adopted but we failed to get the thread for %@", pubKey);
+            return;
+        }
+        
+        // If we were the ones to initiate the reset then we need to send back an empty message
+        if (thread.sessionResetState == TSContactThreadSessionResetStateInitiated) {
+            TSOutgoingMessage *emptyMessage = [TSOutgoingMessage createEmptyOutgoingMessageInThread:thread];
+            [self.messageSenderJobQueue addMessage:emptyMessage transaction:transaction];
+        }
+        
+        // Show session reset done message
+        [[[TSInfoMessage alloc] initWithTimestamp:[NSDate ows_millisecondTimeStamp]
+                                         inThread:thread
+                                      messageType:TSInfoMessageTypeLokiSessionResetDone] saveWithTransaction:transaction];
+        
+        /// Loki: Set our session reset state to none
+        thread.sessionResetState = TSContactThreadSessionResetStateNone;
+        [thread saveWithTransaction:transaction];
+    }];
+}
+
 @end
 
 NS_ASSUME_NONNULL_END