From 8fff98d94eb6fc0259f06a6ec6a02851cfc39f11 Mon Sep 17 00:00:00 2001
From: Mikunj <mikunj@live.com.au>
Date: Mon, 13 May 2019 12:19:05 +1000
Subject: [PATCH] Save prekey bundle from content message.

---
 .../SSKProtoPrekeyBundleMessage+PreKeyBundle.swift  | 13 ++++++++++++-
 SignalServiceKit/src/Messages/OWSMessageManager.m   | 11 +++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/SignalServiceKit/src/Loki/Extensions/SSKProtoPrekeyBundleMessage+PreKeyBundle.swift b/SignalServiceKit/src/Loki/Extensions/SSKProtoPrekeyBundleMessage+PreKeyBundle.swift
index 38cc82c5b..9bd4da1f6 100644
--- a/SignalServiceKit/src/Loki/Extensions/SSKProtoPrekeyBundleMessage+PreKeyBundle.swift
+++ b/SignalServiceKit/src/Loki/Extensions/SSKProtoPrekeyBundleMessage+PreKeyBundle.swift
@@ -1,5 +1,17 @@
 public extension SSKProtoPrekeyBundleMessage {
     
+    @objc public var preKeyBundle: PreKeyBundle? {
+        let registrationId = TSAccountManager.sharedInstance().getOrGenerateRegistrationId()
+        return PreKeyBundle(registrationId: Int32(registrationId),
+                            deviceId: Int32(self.deviceID),
+                            preKeyId: Int32(self.prekeyID),
+                            preKeyPublic: self.prekey,
+                            signedPreKeyPublic: self.signedKey,
+                            signedPreKeyId: Int32(self.signedKeyID),
+                            signedPreKeySignature: self.signature,
+                            identityKey: self.identityKey)
+    }
+    
     @objc public class func builder(fromPreKeyBundle preKeyBundle: PreKeyBundle) -> SSKProtoPrekeyBundleMessageBuilder {
         let builder = self.builder()
         
@@ -12,5 +24,4 @@ public extension SSKProtoPrekeyBundleMessage {
         
         return builder
     }
-    
 }
diff --git a/SignalServiceKit/src/Messages/OWSMessageManager.m b/SignalServiceKit/src/Messages/OWSMessageManager.m
index 8c0d7223e..eceedafe3 100644
--- a/SignalServiceKit/src/Messages/OWSMessageManager.m
+++ b/SignalServiceKit/src/Messages/OWSMessageManager.m
@@ -26,6 +26,7 @@
 #import "OWSMessageUtils.h"
 #import "OWSOutgoingReceiptManager.h"
 #import "OWSPrimaryStorage+SessionStore.h"
+#import "OWSPrimaryStorage+Loki.h"
 #import "OWSPrimaryStorage.h"
 #import "OWSReadReceiptManager.h"
 #import "OWSRecordTranscriptJob.h"
@@ -411,6 +412,16 @@ NS_ASSUME_NONNULL_BEGIN
             return;
         }
         OWSLogInfo(@"handling content: <Content: %@>", [self descriptionForContent:contentProto]);
+        
+        // Loki: Handle PreKeyBundle message
+        if (contentProto.prekeyBundleMessage) {
+            OWSLogInfo(@"Received a prekey bundle message from: %@", envelope.source);
+            PreKeyBundle *_Nullable bundle = contentProto.prekeyBundleMessage.preKeyBundle;
+            if (!bundle) {
+                OWSFailDebug(@"Failed to create PreKeyBundle from message");
+            }
+            [[OWSPrimaryStorage sharedManager] setPreKeyBundle:bundle forContact:envelope.source];
+        }
 
         if (contentProto.syncMessage) {
             [self throws_handleIncomingEnvelope:envelope