diff --git a/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.h b/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.h index 83986879d..ea2efd40b 100644 --- a/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.h +++ b/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.h @@ -1,5 +1,6 @@ #import "OWSPrimaryStorage.h" #import "PreKeyRecord.h" +#import "PreKeyBundle.h" NS_ASSUME_NONNULL_BEGIN @@ -24,6 +25,34 @@ NS_ASSUME_NONNULL_BEGIN */ - (PreKeyRecord *)getPreKeyForContact:(NSString *)pubKey; +# pragma mark - PreKeyBundle + +/** + Generate a PreKeyBundle for the given contact. + This doesn't store the pre key bundle, and you shouldn't store this bundle. + It's used for generating bundles to send to other users. + + @param pubKey The hex encoded public key of the contact. + @return A pre key bundle for the contact. + */ +- (PreKeyBundle *)generatePreKeyBundleForContact:(NSString *)pubKey; + +/** + Get the PreKeyBundle associated with the given contact. + + @param pubKey The hex encoded public key of the contact. + @return The pre key bundle or nil if it doesn't exist. + */ +- (PreKeyBundle *_Nullable)getPreKeyBundleForContact:(NSString *)pubKey; + +/** + Set the PreKeyBundle for the given contact. + + @param bundle The pre key bundle. + @param pubKey The hex encoded public key of the contact. + */ +- (void)setPreKeyBundle:(PreKeyBundle *)bundle forContact:(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 a517d90e6..8e7aa84a9 100644 --- a/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.m +++ b/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.m @@ -1,8 +1,13 @@ #import "OWSPrimaryStorage+Loki.h" #import "OWSPrimaryStorage+PreKeyStore.h" +#import "OWSPrimaryStorage+SignedPreKeyStore.h" +#import "OWSDevice.h" +#import "OWSIdentityManager.h" +#import "TSAccountManager.h" #import "YapDatabaseConnection+OWS.h" #define LokiPreKeyContactCollection @"LokiPreKeyContactCollection" +#define LokiPreKeyBundleCollection @"LokiPreKeyBundleCollection" @implementation OWSPrimaryStorage (Loki) @@ -45,4 +50,37 @@ return record; } +# pragma mark - PreKeyBundle + +- (PreKeyBundle *)generatePreKeyBundleForContact:(NSString *)pubKey { + ECKeyPair *_Nullable myKeyPair = [[OWSIdentityManager sharedManager] identityKeyPair]; + OWSAssertDebug(myKeyPair); + + SignedPreKeyRecord *signedPreKey = [self currentSignedPreKey]; + PreKeyRecord *preKey = [self getPreKeyForContact:pubKey]; + + uint32_t registrationId = [[TSAccountManager sharedInstance] getOrGenerateRegistrationId]; + + PreKeyBundle *bundle = [[PreKeyBundle alloc] initWithRegistrationId:registrationId + deviceId:OWSDevicePrimaryDeviceId + preKeyId:preKey.Id + preKeyPublic:preKey.keyPair.publicKey + signedPreKeyPublic:signedPreKey.keyPair.publicKey + signedPreKeyId:signedPreKey.Id + signedPreKeySignature:signedPreKey.signature + identityKey:myKeyPair.publicKey]; + return bundle; +} + +- (PreKeyBundle *_Nullable)getPreKeyBundleForContact:(NSString *)pubKey { + PreKeyBundle *bundle = [self.dbReadWriteConnection preKeyBundleForKey:pubKey inCollection:LokiPreKeyBundleCollection]; + return bundle; +} + +- (void)setPreKeyBundle:(PreKeyBundle *)bundle forContact:(NSString *)pubKey { + [self.dbReadWriteConnection setObject:bundle + forKey:pubKey + inCollection:LokiPreKeyBundleCollection]; +} + @end diff --git a/SignalServiceKit/src/Storage/YapDatabaseConnection+OWS.h b/SignalServiceKit/src/Storage/YapDatabaseConnection+OWS.h index 1b9e3afab..96b27b277 100644 --- a/SignalServiceKit/src/Storage/YapDatabaseConnection+OWS.h +++ b/SignalServiceKit/src/Storage/YapDatabaseConnection+OWS.h @@ -7,6 +7,7 @@ @class ECKeyPair; @class PreKeyRecord; @class SignedPreKeyRecord; +@class PreKeyBundle; NS_ASSUME_NONNULL_BEGIN @@ -23,6 +24,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; - (NSUInteger)numberOfKeysInCollection:(NSString *)collection; diff --git a/SignalServiceKit/src/Storage/YapDatabaseConnection+OWS.m b/SignalServiceKit/src/Storage/YapDatabaseConnection+OWS.m index f21372f80..428444ca7 100644 --- a/SignalServiceKit/src/Storage/YapDatabaseConnection+OWS.m +++ b/SignalServiceKit/src/Storage/YapDatabaseConnection+OWS.m @@ -4,6 +4,7 @@ #import "YapDatabaseConnection+OWS.h" #import +#import #import #import #import @@ -78,6 +79,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 { return [self objectForKey:key inCollection:collection ofExpectedType:[SignedPreKeyRecord class]];