From 080b24b282a7dccb67092f74ed3595a4faf0fe8e Mon Sep 17 00:00:00 2001
From: Mikunj <mikunj@live.com.au>
Date: Mon, 17 Feb 2020 12:39:22 +1100
Subject: [PATCH] Handle session adoption correctly

---
 .../src/Loki/Database/LokiSessionReset.swift  | 26 +++++++++++++++
 .../src/Messages/OWSMessageManager.m          | 33 -------------------
 2 files changed, 26 insertions(+), 33 deletions(-)

diff --git a/SignalServiceKit/src/Loki/Database/LokiSessionReset.swift b/SignalServiceKit/src/Loki/Database/LokiSessionReset.swift
index edb6df1c3..53b3e4cf1 100644
--- a/SignalServiceKit/src/Loki/Database/LokiSessionReset.swift
+++ b/SignalServiceKit/src/Loki/Database/LokiSessionReset.swift
@@ -41,4 +41,30 @@ public class LokiSessionReset: NSObject, SessionResetProtocol {
         guard let thread = TSContactThread.getWithContactId(recipientId, transaction: transaction) else { return .none }
         return thread.sessionResetStatus
     }
+
+    public func onNewSessionAdopted(for recipientId: String, protocolContext: Any?) {
+        guard let transaction = protocolContext as? YapDatabaseReadWriteTransaction else {
+            Logger.warn("[Loki] Cannot handle new session adoption because invalid transaction was passed")
+            return
+        }
+
+        guard recipientId.count > 0 else { return }
+        guard let thread = TSContactThread.getWithContactId(recipientId, transaction: transaction) else {
+            Logger.debug("[Loki] A new session was adopted but the thread couldn't be found for \(recipientId)")
+            return
+        }
+
+        // If the current user initiated the reset then send back an empty message to acknowledge the completion of the session reset
+        if (thread.sessionResetStatus == .initiated) {
+            let emptyMessage = EphemeralMessage(in: thread)
+            SSKEnvironment.shared.messageSender.sendPromise(message: emptyMessage).retainUntilComplete()
+        }
+
+        // Show session reset done message
+        TSInfoMessage(timestamp: NSDate.ows_millisecondTimeStamp(), in: thread, messageType: .typeLokiSessionResetDone).save(with: transaction)
+
+        // Clear the session reset status
+        thread.sessionResetStatus = .none
+        thread.save(with: transaction)
+    }
 }
diff --git a/SignalServiceKit/src/Messages/OWSMessageManager.m b/SignalServiceKit/src/Messages/OWSMessageManager.m
index 8ae1e0b94..1fb3edd8c 100644
--- a/SignalServiceKit/src/Messages/OWSMessageManager.m
+++ b/SignalServiceKit/src/Messages/OWSMessageManager.m
@@ -93,9 +93,6 @@ NS_ASSUME_NONNULL_BEGIN
     _primaryStorage = primaryStorage;
     _dbConnection = primaryStorage.newDatabaseConnection;
     _incomingMessageFinder = [[OWSIncomingMessageFinder alloc] initWithPrimaryStorage:primaryStorage];
-    
-    // Loki: Observe session changes
-    [NSNotificationCenter.defaultCenter addObserver:self selector:@selector(handleNewSessionAdopted:) name:LKSessionCipher.sessionAdoptedNotification object:nil];
 
     OWSSingletonAssert();
 
@@ -2019,36 +2016,6 @@ NS_ASSUME_NONNULL_BEGIN
     }
 }
 
-# pragma mark - Loki Session Handling
-
-- (void)handleNewSessionAdopted:(NSNotification *)notification {
-    NSString *hexEncodedPublicKey = notification.userInfo[LKSessionCipher.contactPubKeyField];
-    if (hexEncodedPublicKey.length == 0) { return; }
-    
-    [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
-        TSContactThread *thread = [TSContactThread getThreadWithContactId:hexEncodedPublicKey transaction:transaction];
-        if (thread == nil) {
-            NSLog(@"[Loki] A new session was adopted but the thread couldn't be found for: %@.", hexEncodedPublicKey);
-            return;
-        }
-
-        // If the current user initiated the reset then send back an empty message to acknowledge the completion of the session reset
-        if (thread.sessionResetStatus == LKSessionResetStatusInitiated) {
-            LKEphemeralMessage *emptyMessage = [[LKEphemeralMessage alloc] initInThread:thread];
-            [self.messageSenderJobQueue addMessage:emptyMessage transaction:transaction];
-        }
-        
-        // Show session reset done message
-        [[[TSInfoMessage alloc] initWithTimestamp:NSDate.ows_millisecondTimeStamp
-                                         inThread:thread
-                                      messageType:TSInfoMessageTypeLokiSessionResetDone] saveWithTransaction:transaction];
-        
-        // Clear the session reset state
-        thread.sessionResetStatus = LKSessionResetStatusNone;
-        [thread saveWithTransaction:transaction];
-    }];
-}
-
 @end
 
 NS_ASSUME_NONNULL_END