From f8de85ac42c44b582a6212353de4ab5b2b7768b6 Mon Sep 17 00:00:00 2001 From: Mikunj Date: Tue, 14 May 2019 12:33:48 +1000 Subject: [PATCH] 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. --- .../Loki/Extensions/OWSPrimaryStorage+Loki.h | 7 ++++-- .../Loki/Extensions/OWSPrimaryStorage+Loki.m | 14 +++++------ ...rotoPrekeyBundleMessage+PreKeyBundle.swift | 25 +++++++++++-------- .../src/Messages/OWSMessageManager.m | 4 +-- .../src/Messages/OWSMessageSender.m | 4 +-- .../src/Storage/YapDatabaseTransaction+OWS.h | 2 ++ .../src/Storage/YapDatabaseTransaction+OWS.m | 6 +++++ 7 files changed, 38 insertions(+), 24 deletions(-) diff --git a/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.h b/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.h index 0afb932bf..5cc04ab62 100644 --- a/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.h +++ b/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.h @@ -1,6 +1,7 @@ #import "OWSPrimaryStorage.h" #import "PreKeyRecord.h" #import "PreKeyBundle.h" +#import 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 diff --git a/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.m b/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.m index 30e09d50d..a120e43ce 100644 --- a/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.m +++ b/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 { - [self.dbReadWriteConnection setObject:bundle - forKey:pubKey - inCollection:LokiPreKeyBundleCollection]; +- (void)setPreKeyBundle:(PreKeyBundle *)bundle forContact:(NSString *)pubKey transaction:(YapDatabaseReadWriteTransaction *)transaction { + [transaction setObject:bundle + forKey:pubKey + inCollection:LokiPreKeyBundleCollection]; } -- (void)removePreKeyBundleForContact:(NSString *)pubKey { - [self.dbReadWriteConnection removeObjectForKey:pubKey inCollection:LokiPreKeyBundleCollection]; +- (void)removePreKeyBundleForContact:(NSString *)pubKey transaction:(YapDatabaseReadWriteTransaction *)transaction { + [transaction removeObjectForKey:pubKey inCollection:LokiPreKeyBundleCollection]; } @end diff --git a/SignalServiceKit/src/Loki/Extensions/SSKProtoPrekeyBundleMessage+PreKeyBundle.swift b/SignalServiceKit/src/Loki/Extensions/SSKProtoPrekeyBundleMessage+PreKeyBundle.swift index 623846f8c..67789154e 100644 --- a/SignalServiceKit/src/Loki/Extensions/SSKProtoPrekeyBundleMessage+PreKeyBundle.swift +++ b/SignalServiceKit/src/Loki/Extensions/SSKProtoPrekeyBundleMessage+PreKeyBundle.swift @@ -1,16 +1,7 @@ - @objc public extension SSKProtoPrekeyBundleMessage { - @objc public var preKeyBundle: PreKeyBundle? { - let registrationId = TSAccountManager.sharedInstance().getOrGenerateRegistrationId() - return PreKeyBundle(registrationId: Int32(registrationId), - deviceId: Int32(deviceID), - preKeyId: Int32(prekeyID), - preKeyPublic: prekey, - signedPreKeyPublic: signedKey, - signedPreKeyId: Int32(signedKeyID), - signedPreKeySignature: signature, - identityKey: identityKey) + private var accountManager: TSAccountManager { + return TSAccountManager.sharedInstance() } @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) + } } diff --git a/SignalServiceKit/src/Messages/OWSMessageManager.m b/SignalServiceKit/src/Messages/OWSMessageManager.m index eceedafe3..f43e6b740 100644 --- a/SignalServiceKit/src/Messages/OWSMessageManager.m +++ b/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) { diff --git a/SignalServiceKit/src/Messages/OWSMessageSender.m b/SignalServiceKit/src/Messages/OWSMessageSender.m index 051701e95..092605e21 100644 --- a/SignalServiceKit/src/Messages/OWSMessageSender.m +++ b/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; } diff --git a/SignalServiceKit/src/Storage/YapDatabaseTransaction+OWS.h b/SignalServiceKit/src/Storage/YapDatabaseTransaction+OWS.h index b533e367f..fb9f6524c 100644 --- a/SignalServiceKit/src/Storage/YapDatabaseTransaction+OWS.h +++ b/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 diff --git a/SignalServiceKit/src/Storage/YapDatabaseTransaction+OWS.m b/SignalServiceKit/src/Storage/YapDatabaseTransaction+OWS.m index 255663255..0fbd9f64e 100644 --- a/SignalServiceKit/src/Storage/YapDatabaseTransaction+OWS.m +++ b/SignalServiceKit/src/Storage/YapDatabaseTransaction+OWS.m @@ -4,6 +4,7 @@ #import "YapDatabaseTransaction+OWS.h" #import +#import #import #import @@ -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);