From ea7205ce7681537df58d79addead060f0c5d024d Mon Sep 17 00:00:00 2001 From: Mikunj Date: Mon, 13 May 2019 11:40:33 +1000 Subject: [PATCH] Discard PreKeyBundle once session has been built. --- .../src/Loki/Extensions/OWSPrimaryStorage+Loki.h | 7 +++++++ .../src/Loki/Extensions/OWSPrimaryStorage+Loki.m | 6 ++++++ SignalServiceKit/src/Messages/OWSMessageSender.m | 3 +++ 3 files changed, 16 insertions(+) diff --git a/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.h b/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.h index ea2efd40b..0ccc9a7ff 100644 --- a/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.h +++ b/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.h @@ -53,6 +53,13 @@ NS_ASSUME_NONNULL_BEGIN */ - (void)setPreKeyBundle:(PreKeyBundle *)bundle forContact:(NSString *)pubKey; +/** + Remove the PreKeyBundle for the given contact. + + @param pubKey The hex encoded public key of the contact. + */ +- (void)removePreKeyBundleForContact:(NSString *)pubKey; + @end NS_ASSUME_NONNULL_END diff --git a/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.m b/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.m index ad27c72c7..a74f873c0 100644 --- a/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.m +++ b/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.m @@ -91,4 +91,10 @@ inCollection:LokiPreKeyBundleCollection]; } +- (void)removePreKeyBundleForContact:(NSString *)pubKey +{ + [self.dbReadWriteConnection removeObjectForKey:pubKey + inCollection:LokiPreKeyBundleCollection]; +} + @end diff --git a/SignalServiceKit/src/Messages/OWSMessageSender.m b/SignalServiceKit/src/Messages/OWSMessageSender.m index 4247bf716..051701e95 100644 --- a/SignalServiceKit/src/Messages/OWSMessageSender.m +++ b/SignalServiceKit/src/Messages/OWSMessageSender.m @@ -1665,6 +1665,9 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException"; [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) { @try { [builder throws_processPrekeyBundle:bundle protocolContext:transaction]; + + // Loki: Discard the prekey bundle here since the session is initiated + [[OWSPrimaryStorage sharedManager] removePreKeyBundleForContact:recipientId]; } @catch (NSException *caughtException) { exception = caughtException; }