Add protocol context to protocol kit.

SignalServiceKit/src/Messages/OWSBatchMessageProcessor.m

@@ -12,6 +12,7 @@
 #import "OWSSignalServiceProtos.pb.h"
 #import "OWSStorage.h"
 #import "TSDatabaseView.h"
+#import "TSStorageManager+SessionStore.h"
 #import "TSStorageManager.h"
 #import "TSYapDatabaseObject.h"
 #import "Threading.h"
@@ -363,7 +364,9 @@ NSString *const OWSMessageContentJobFinderExtensionGroup = @"OWSMessageContentJo
 {
     AssertOnDispatchQueue(self.serialQueue);
 
-    [self.dbReadWriteConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
+    // Note that we use the protocolStoreDBConnection to process incoming messages,
+    // to ensure session & identity store consistency.
+    [TSStorageManager.protocolStoreDBConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
         for (OWSMessageContentJob *job in jobs) {
             [self.messagesManager processEnvelope:job.envelopeProto
                                     plaintextData:job.plaintextData

SignalServiceKit/src/Messages/OWSMessageUtils.h

@@ -0,0 +1,29 @@
+//
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
+//
+
+#import "OWSMessageHandler.h"
+
+NS_ASSUME_NONNULL_BEGIN
+
+@class OWSSignalServiceProtosEnvelope;
+@class TSThread;
+@class YapDatabaseReadWriteTransaction;
+
+@interface OWSMessageManager : OWSMessageHandler
+
+- (instancetype)init NS_UNAVAILABLE;
++ (instancetype)sharedManager;
+
+// processEnvelope: can be called from any thread.
+- (void)processEnvelope:(OWSSignalServiceProtosEnvelope *)envelope
+          plaintextData:(NSData *_Nullable)plaintextData
+            transaction:(YapDatabaseReadWriteTransaction *)transaction;
+
+- (NSUInteger)unreadMessagesCount;
+- (NSUInteger)unreadMessagesCountExcept:(TSThread *)thread;
+- (NSUInteger)unreadMessagesInThread:(TSThread *)thread;
+
+@end
+
+NS_ASSUME_NONNULL_END

SignalServiceKit/src/Storage/AxolotlStore/TSStorageManager+SessionStore.h

@@ -9,6 +9,8 @@ NS_ASSUME_NONNULL_BEGIN
 
 @interface TSStorageManager (SessionStore) <SessionStore>
 
++ (YapDatabaseConnection *)protocolStoreDBConnection;
+
 - (void)archiveAllSessionsForContact:(NSString *)contactIdentifier protocolContext:(nullable id)protocolContext;
 
 #pragma mark - debug

SignalServiceKit/src/Storage/AxolotlStore/TSStorageManager+SessionStore.m

@@ -196,20 +196,16 @@ NSString *const kSessionStoreDBConnectionKey = @"kSessionStoreDBConnectionKey";
 
 - (void)resetSessionStore
 {
-    AssertIsOnSessionStoreQueue();
-
     DDLogWarn(@"%@ resetting session store", self.logTag);
-    [self.sessionDBConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *_Nonnull transaction) {
+    [self.protocolStoreDBConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *_Nonnull transaction) {
         [transaction removeAllObjectsInCollection:TSStorageManagerSessionStoreCollection];
     }];
 }
 
 - (void)printAllSessions
 {
-    AssertIsOnSessionStoreQueue();
-
     NSString *tag = @"[TSStorageManager (SessionStore)]";
-    [self.sessionDBConnection readWithBlock:^(YapDatabaseReadTransaction *_Nonnull transaction) {
+    [self.protocolStoreDBConnection readWithBlock:^(YapDatabaseReadTransaction *_Nonnull transaction) {
         DDLogDebug(@"%@ All Sessions:", tag);
         [transaction
             enumerateKeysAndObjectsInCollection:TSStorageManagerSessionStoreCollection
@@ -258,18 +254,14 @@ NSString *const kSessionStoreDBConnectionKey = @"kSessionStoreDBConnectionKey";
 
 - (void)snapshotSessionStore
 {
-    AssertIsOnSessionStoreQueue();
-
-    [self.sessionDBConnection snapshotCollection:TSStorageManagerSessionStoreCollection
-                                snapshotFilePath:self.snapshotFilePath];
+    [self.protocolStoreDBConnection snapshotCollection:TSStorageManagerSessionStoreCollection
+                                      snapshotFilePath:self.snapshotFilePath];
 }
 
 - (void)restoreSessionStore
 {
-    AssertIsOnSessionStoreQueue();
-
-    [self.sessionDBConnection restoreSnapshotOfCollection:TSStorageManagerSessionStoreCollection
-                                         snapshotFilePath:self.snapshotFilePath];
+    [self.protocolStoreDBConnection restoreSnapshotOfCollection:TSStorageManagerSessionStoreCollection
+                                               snapshotFilePath:self.snapshotFilePath];
 }
 #endif