From 17cede4a144b5939554da178b7dd802e9890964d Mon Sep 17 00:00:00 2001
From: Mikunj <mikunj@live.com.au>
Date: Mon, 7 Oct 2019 16:09:44 +1100
Subject: [PATCH] Refactor DeviceLinkMessage. Send a prekeybundle with a
 pairing request.

---
 .../Loki/Settings/DeviceLinkingModal.swift    |  6 +--
 .../src/Loki/Messaging/LKDeviceLinkMessage.m  | 43 +++++++++++--------
 .../Messages/Interactions/TSOutgoingMessage.m |  2 +
 3 files changed, 29 insertions(+), 22 deletions(-)

diff --git a/Signal/src/Loki/Settings/DeviceLinkingModal.swift b/Signal/src/Loki/Settings/DeviceLinkingModal.swift
index 6d1a3eb13..6ccbba09e 100644
--- a/Signal/src/Loki/Settings/DeviceLinkingModal.swift
+++ b/Signal/src/Loki/Settings/DeviceLinkingModal.swift
@@ -171,9 +171,9 @@ final class DeviceLinkingModal : Modal, DeviceLinkingSessionDelegate {
     }
     
     @objc override func cancel() {
-        let session = DeviceLinkingSession.current!
-        session.stopListeningForLinkingRequests()
-        session.markLinkingRequestAsProcessed() // Only relevant in master mode
+        let session = DeviceLinkingSession.current
+        session?.stopListeningForLinkingRequests()
+        session?.markLinkingRequestAsProcessed() // Only relevant in master mode
         delegate?.handleDeviceLinkingModalDismissed() // Only relevant in slave mode
         if let deviceLink = deviceLink {
             OWSPrimaryStorage.shared().dbReadWriteConnection.readWrite { transaction in
diff --git a/SignalServiceKit/src/Loki/Messaging/LKDeviceLinkMessage.m b/SignalServiceKit/src/Loki/Messaging/LKDeviceLinkMessage.m
index 73c671028..252c4bb6c 100644
--- a/SignalServiceKit/src/Loki/Messaging/LKDeviceLinkMessage.m
+++ b/SignalServiceKit/src/Loki/Messaging/LKDeviceLinkMessage.m
@@ -1,5 +1,6 @@
 #import "LKDeviceLinkMessage.h"
 #import "OWSIdentityManager.h"
+#import "OWSPrimaryStorage+Loki.h"
 #import "SignalRecipient.h"
 #import <SignalCoreKit/NSData+OWS.h>
 #import <SignalCoreKit/NSDate+OWS.h>
@@ -30,19 +31,28 @@
 }
 
 #pragma mark Building
-- (nullable NSData *)buildPlainTextData:(SignalRecipient *)recipient
+- (SSKProtoContentBuilder *)contentBuilder:(SignalRecipient *)recipient
 {
-    // Prepare
-    SSKProtoContentBuilder *contentBuilder = [self contentBuilder:recipient];
-    NSError *error;
-    // Data message
-    SSKProtoDataMessage *_Nullable dataMessage = [self buildDataMessage:recipient.recipientId];
-    if (!dataMessage) {
-        OWSFailDebug(@"Failed to build data message.");
-        return nil;
+    SSKProtoContentBuilder *contentBuilder = [super contentBuilder:recipient];
+       
+    // If it's a linking request then we should send a prekey bundle
+    if (self.masterSignature == nil) {
+        PreKeyBundle *bundle = [OWSPrimaryStorage.sharedManager generatePreKeyBundleForContact:recipient.recipientId];
+        SSKProtoPrekeyBundleMessageBuilder *preKeyBuilder = [SSKProtoPrekeyBundleMessage builderFromPreKeyBundle:bundle];
+
+        // Build the prekey bundle message
+        NSError *error;
+        SSKProtoPrekeyBundleMessage *_Nullable message = [preKeyBuilder buildAndReturnError:&error];
+        if (error || !message) {
+            OWSFailDebug(@"Failed to build pre key bundle for %@: %@", recipient.recipientId, error);
+            return nil;
+        } else {
+            [contentBuilder setPrekeyBundleMessage:message];
+        }
     }
-    [contentBuilder setDataMessage:dataMessage];
+    
     // Device link message
+    NSError *error;
     SSKProtoLokiDeviceLinkMessageBuilder *deviceLinkMessageBuilder = [SSKProtoLokiDeviceLinkMessage builder];
     [deviceLinkMessageBuilder setMasterHexEncodedPublicKey:self.masterHexEncodedPublicKey];
     [deviceLinkMessageBuilder setSlaveHexEncodedPublicKey:self.slaveHexEncodedPublicKey];
@@ -52,16 +62,11 @@
     if (error || deviceLinkMessage == nil) {
         OWSFailDebug(@"Failed to build device link message due to error: %@.", error);
         return nil;
+    } else {
+        [contentBuilder setLokiDeviceLinkMessage:deviceLinkMessage];
     }
-    [contentBuilder setLokiDeviceLinkMessage:deviceLinkMessage];
-    // Serialize
-    NSData *_Nullable contentAsData = [contentBuilder buildSerializedDataAndReturnError:&error];
-    if (error || !contentAsData) {
-        OWSFailDebug(@"Failed to build serialized message content due to error: %@.", error);
-        return nil;
-    }
-    // Return
-    return contentAsData;
+    
+    return contentBuilder;
 }
 
 #pragma mark Settings
diff --git a/SignalServiceKit/src/Messages/Interactions/TSOutgoingMessage.m b/SignalServiceKit/src/Messages/Interactions/TSOutgoingMessage.m
index 6620a57b9..883d14874 100644
--- a/SignalServiceKit/src/Messages/Interactions/TSOutgoingMessage.m
+++ b/SignalServiceKit/src/Messages/Interactions/TSOutgoingMessage.m
@@ -1137,6 +1137,8 @@ NSString *NSStringForOutgoingMessageRecipientState(OWSOutgoingMessageRecipientSt
     }
 
     SSKProtoContentBuilder *contentBuilder = [self contentBuilder:recipient];
+    if (contentBuilder == nil) { return nil; }
+    
     [contentBuilder setDataMessage:dataMessage];
     NSData *_Nullable contentData = [contentBuilder buildSerializedDataAndReturnError:&error];
     if (error || !contentData) {