Add generation and storage of PreKeyBundle.

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

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

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;

SignalServiceKit/src/Storage/YapDatabaseConnection+OWS.m

@@ -4,6 +4,7 @@
 
 #import "YapDatabaseConnection+OWS.h"
 #import <AxolotlKit/PreKeyRecord.h>
+#import <AxolotlKit/PreKeyBundle.h>
 #import <AxolotlKit/SignedPrekeyRecord.h>
 #import <Curve25519Kit/Curve25519.h>
 #import <YapDatabase/YapDatabaseTransaction.h>
@@ -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]];