Fix deadlock crash when receiving PreKeyBundle message.

It was deadlocking because we had a transaction inside another transaction.
To stop this we can pass in the parent transaction when setting or removing bundles, as well as generating PreKeyBundle from the proto message.

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

@@ -1,6 +1,7 @@
 #import "OWSPrimaryStorage.h"
 #import "PreKeyRecord.h"
 #import "PreKeyBundle.h"
+#import <YapDatabase/YapDatabase.h>
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -49,16 +50,18 @@ NS_ASSUME_NONNULL_BEGIN
  Set the `PreKeyBundle` for the given contact.
 
  @param bundle The pre key bundle.
+ @param transaction A YapDatabaseReadWriteTransaction.
  @param pubKey The hex encoded public key of the contact.
  */
-- (void)setPreKeyBundle:(PreKeyBundle *)bundle forContact:(NSString *)pubKey;
+- (void)setPreKeyBundle:(PreKeyBundle *)bundle forContact:(NSString *)pubKey transaction:(YapDatabaseReadWriteTransaction *)transaction;
 
 /**
  Remove the `PreKeyBundle` for the given contact.
 
  @param pubKey The hex encoded public key of the contact.
+ @param transaction A YapDatabaseReadWriteTransaction.
  */
-- (void)removePreKeyBundleForContact:(NSString *)pubKey;
+- (void)removePreKeyBundleForContact:(NSString *)pubKey transaction:(YapDatabaseReadWriteTransaction *)transaction;
 
 @end
 

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

@@ -91,17 +91,17 @@
 }
 
 - (PreKeyBundle *_Nullable)getPreKeyBundleForContact:(NSString *)pubKey {
-    return [self.dbReadWriteConnection preKeyBundleForKey:pubKey inCollection:LokiPreKeyBundleCollection];
+    return [self.dbReadConnection preKeyBundleForKey:pubKey inCollection:LokiPreKeyBundleCollection];
 }
 
-- (void)setPreKeyBundle:(PreKeyBundle *)bundle forContact:(NSString *)pubKey {
+- (void)setPreKeyBundle:(PreKeyBundle *)bundle forContact:(NSString *)pubKey transaction:(YapDatabaseReadWriteTransaction *)transaction {
-    [self.dbReadWriteConnection setObject:bundle
+    [transaction setObject:bundle
-                                   forKey:pubKey
+                    forKey:pubKey
-                             inCollection:LokiPreKeyBundleCollection];
+              inCollection:LokiPreKeyBundleCollection];
 }
 
-- (void)removePreKeyBundleForContact:(NSString *)pubKey {
+- (void)removePreKeyBundleForContact:(NSString *)pubKey transaction:(YapDatabaseReadWriteTransaction *)transaction {
-    [self.dbReadWriteConnection removeObjectForKey:pubKey inCollection:LokiPreKeyBundleCollection];
+    [transaction removeObjectForKey:pubKey inCollection:LokiPreKeyBundleCollection];
 }
 
 @end

SignalServiceKit/src/Loki/Extensions/SSKProtoPrekeyBundleMessage+PreKeyBundle.swift

@@ -1,16 +1,7 @@
-
 @objc public extension SSKProtoPrekeyBundleMessage {
     
-    @objc public var preKeyBundle: PreKeyBundle? {
+    private var accountManager: TSAccountManager {
-        let registrationId = TSAccountManager.sharedInstance().getOrGenerateRegistrationId()
+        return TSAccountManager.sharedInstance()
-        return PreKeyBundle(registrationId: Int32(registrationId),
-                            deviceId: Int32(deviceID),
-                            preKeyId: Int32(prekeyID),
-                            preKeyPublic: prekey,
-                            signedPreKeyPublic: signedKey,
-                            signedPreKeyId: Int32(signedKeyID),
-                            signedPreKeySignature: signature,
-                            identityKey: identityKey)
     }
     
     @objc public class func builder(fromPreKeyBundle preKeyBundle: PreKeyBundle) -> SSKProtoPrekeyBundleMessageBuilder {
@@ -25,4 +16,16 @@
         
         return builder
     }
+    
+    @objc public func preKeyBundle(withTransaction transaction: YapDatabaseReadWriteTransaction) -> PreKeyBundle? {
+        let registrationId = accountManager.getOrGenerateRegistrationId(transaction)
+        return PreKeyBundle(registrationId: Int32(registrationId),
+                            deviceId: Int32(deviceID),
+                            preKeyId: Int32(prekeyID),
+                            preKeyPublic: prekey,
+                            signedPreKeyPublic: signedKey,
+                            signedPreKeyId: Int32(signedKeyID),
+                            signedPreKeySignature: signature,
+                            identityKey: identityKey)
+    }
 }

SignalServiceKit/src/Messages/OWSMessageManager.m

@@ -416,11 +416,11 @@ NS_ASSUME_NONNULL_BEGIN
         // Loki: Handle PreKeyBundle message
         if (contentProto.prekeyBundleMessage) {
             OWSLogInfo(@"Received a prekey bundle message from: %@", envelope.source);
-            PreKeyBundle *_Nullable bundle = contentProto.prekeyBundleMessage.preKeyBundle;
+            PreKeyBundle *_Nullable bundle = [contentProto.prekeyBundleMessage preKeyBundleWithTransaction:transaction];
             if (!bundle) {
                 OWSFailDebug(@"Failed to create PreKeyBundle from message");
             }
-            [[OWSPrimaryStorage sharedManager] setPreKeyBundle:bundle forContact:envelope.source];
+            [self.primaryStorage setPreKeyBundle:bundle forContact:envelope.source transaction:transaction];
         }
 
         if (contentProto.syncMessage) {

SignalServiceKit/src/Messages/OWSMessageSender.m

@@ -1616,7 +1616,7 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException";
         OWSRaiseException(NoSessionForTransientMessageException, @"No session for transient message.");
     }
     
-    PreKeyBundle *_Nullable bundle = [[OWSPrimaryStorage sharedManager] getPreKeyBundleForContact:recipientId];
+    PreKeyBundle *_Nullable bundle = [storage getPreKeyBundleForContact:recipientId];
     __block NSException *_Nullable exception;
 
     /** Loki: Original code
@@ -1667,7 +1667,7 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException";
                 [builder throws_processPrekeyBundle:bundle protocolContext:transaction];
                 
                 // Loki: Discard the prekey bundle here since the session is initiated
-                [[OWSPrimaryStorage sharedManager] removePreKeyBundleForContact:recipientId];
+                [storage removePreKeyBundleForContact:recipientId transaction:transaction];
             } @catch (NSException *caughtException) {
                 exception = caughtException;
             }

SignalServiceKit/src/Storage/YapDatabaseTransaction+OWS.h

@@ -6,6 +6,7 @@
 
 @class ECKeyPair;
 @class PreKeyRecord;
+@class PreKeyBundle;
 @class SignedPreKeyRecord;
 
 NS_ASSUME_NONNULL_BEGIN
@@ -20,6 +21,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (nullable NSData *)dataForKey:(NSString *)key inCollection:(NSString *)collection;
 - (nullable ECKeyPair *)keyPairForKey:(NSString *)key inCollection:(NSString *)collection;
 - (nullable PreKeyRecord *)preKeyRecordForKey:(NSString *)key inCollection:(NSString *)collection;
+- (nullable PreKeyBundle *)preKeyBundleForKey:(NSString *)key inCollection:(NSString *)collection;
 - (nullable SignedPreKeyRecord *)signedPreKeyRecordForKey:(NSString *)key inCollection:(NSString *)collection;
 
 @end

SignalServiceKit/src/Storage/YapDatabaseTransaction+OWS.m

@@ -4,6 +4,7 @@
 
 #import "YapDatabaseTransaction+OWS.h"
 #import <AxolotlKit/PreKeyRecord.h>
+#import <AxolotlKit/PreKeyBundle.h>
 #import <AxolotlKit/SignedPrekeyRecord.h>
 #import <Curve25519Kit/Curve25519.h>
 
@@ -69,6 +70,11 @@ NS_ASSUME_NONNULL_BEGIN
     return [self objectForKey:key inCollection:collection ofExpectedType:[PreKeyRecord class]];
 }
 
+- (nullable PreKeyBundle *)preKeyBundleForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    return [self objectForKey:key inCollection:collection ofExpectedType:[PreKeyBundle class]];
+}
+
 - (nullable SignedPreKeyRecord *)signedPreKeyRecordForKey:(NSString *)key inCollection:(NSString *)collection
 {
     OWSAssertDebug(key.length > 0);