From de346cb645dfa4a19e6fe0cfeddb862d30827e3e Mon Sep 17 00:00:00 2001
From: Niels Andriesse <andriesseniels@gmail.com>
Date: Thu, 27 Jun 2019 11:04:56 +1000
Subject: [PATCH] Fix pre key bundle usage

---
 .../src/Account/CreatePreKeysOperation.swift       | 14 +++++++-------
 .../src/Account/PreKeyRefreshOperation.swift       |  9 ++++-----
 .../src/Account/RotateSignedKeyOperation.swift     |  4 ++--
 .../src/Loki/Crypto/OWSPrimaryStorage+Loki.m       | 11 ++++++++++-
 4 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/SignalServiceKit/src/Account/CreatePreKeysOperation.swift b/SignalServiceKit/src/Account/CreatePreKeysOperation.swift
index 1bb9cf2d9..ea1b091e6 100644
--- a/SignalServiceKit/src/Account/CreatePreKeysOperation.swift
+++ b/SignalServiceKit/src/Account/CreatePreKeysOperation.swift
@@ -23,21 +23,21 @@ public class CreatePreKeysOperation: OWSOperation {
     public override func run() {
         Logger.debug("")
 
-        if self.identityKeyManager.identityKeyPair() == nil {
-            self.identityKeyManager.generateNewIdentityKey()
+        if identityKeyManager.identityKeyPair() == nil {
+            identityKeyManager.generateNewIdentityKey()
         }
         
         // Loki: We don't generate PreKeyRecords here.
         // This is because we need the records to be linked to a contact since we don't have a central server.
         // It is done automatically when we generate a PreKeyBundle to send to a contact (`generatePreKeyBundleForContact:`).
         // You can use `getOrCreatePreKeyForContact:` to generate one if needed.
-        let signedPreKeyRecord = self.primaryStorage.generateRandomSignedRecord()
+        let signedPreKeyRecord = primaryStorage.generateRandomSignedRecord()
         signedPreKeyRecord.markAsAcceptedByService()
-        self.primaryStorage.storeSignedPreKey(signedPreKeyRecord.id, signedPreKeyRecord: signedPreKeyRecord)
-        self.primaryStorage.setCurrentSignedPrekeyId(signedPreKeyRecord.id)
+        primaryStorage.storeSignedPreKey(signedPreKeyRecord.id, signedPreKeyRecord: signedPreKeyRecord)
+        primaryStorage.setCurrentSignedPrekeyId(signedPreKeyRecord.id)
 
-        print("[Loki] Create pre keys operation done.")
-        self.reportSuccess()
+        print("[Loki] Pre keys created successfully.")
+        reportSuccess()
         
         /* Loki: Original code
          * ================
diff --git a/SignalServiceKit/src/Account/PreKeyRefreshOperation.swift b/SignalServiceKit/src/Account/PreKeyRefreshOperation.swift
index 75bcd7d16..f6eeb9c5c 100644
--- a/SignalServiceKit/src/Account/PreKeyRefreshOperation.swift
+++ b/SignalServiceKit/src/Account/PreKeyRefreshOperation.swift
@@ -36,12 +36,11 @@ public class RefreshPreKeysOperation: OWSOperation {
             return
         }
         
-        // Loki: Doing this on the global queue because they do it at the bottom
+        // Loki: Doing this on the global queue to match Signal
         DispatchQueue.global().async {
             guard self.primaryStorage.currentSignedPrekeyId() == nil else {
-                Logger.debug("Already have a signed prekey set")
-                self.reportSuccess()
-                return
+                print("[Loki] Using existing signed pre key.")
+                return self.reportSuccess()
             }
             
             let signedPreKeyRecord = self.primaryStorage.generateRandomSignedRecord()
@@ -52,7 +51,7 @@ public class RefreshPreKeysOperation: OWSOperation {
             TSPreKeyManager.clearPreKeyUpdateFailureCount()
             TSPreKeyManager.clearSignedPreKeyRecords()
             
-            print("[Loki] Pre key refresh operation done.")
+            print("[Loki] Pre keys refreshed successfully.")
             self.reportSuccess()
         }
         
diff --git a/SignalServiceKit/src/Account/RotateSignedKeyOperation.swift b/SignalServiceKit/src/Account/RotateSignedKeyOperation.swift
index e09d2f693..c99cb3733 100644
--- a/SignalServiceKit/src/Account/RotateSignedKeyOperation.swift
+++ b/SignalServiceKit/src/Account/RotateSignedKeyOperation.swift
@@ -27,7 +27,7 @@ public class RotateSignedPreKeyOperation: OWSOperation {
             return
         }
         
-        // Loki: Doing this on the global queue because they do it at the bottom
+        // Loki: Doing this on the global queue to match Signal
         DispatchQueue.global().async {
             let signedPreKeyRecord = self.primaryStorage.generateRandomSignedRecord()
             signedPreKeyRecord.markAsAcceptedByService()
@@ -37,7 +37,7 @@ public class RotateSignedPreKeyOperation: OWSOperation {
             TSPreKeyManager.clearPreKeyUpdateFailureCount()
             TSPreKeyManager.clearSignedPreKeyRecords()
             
-            print("[Loki] Rotate signed pre key operation done.")
+            print("[Loki] Pre keys rotated successfully.")
             self.reportSuccess()
         }
 
diff --git a/SignalServiceKit/src/Loki/Crypto/OWSPrimaryStorage+Loki.m b/SignalServiceKit/src/Loki/Crypto/OWSPrimaryStorage+Loki.m
index 8d07532d0..1e72f3830 100644
--- a/SignalServiceKit/src/Loki/Crypto/OWSPrimaryStorage+Loki.m
+++ b/SignalServiceKit/src/Loki/Crypto/OWSPrimaryStorage+Loki.m
@@ -89,7 +89,16 @@
     
     ECKeyPair *_Nullable keyPair = self.identityManager.identityKeyPair;
     OWSAssertDebug(keyPair);
-    
+
+    // Refresh the signed pre key if needed
+    if (self.currentSignedPreKey == nil) {
+        SignedPreKeyRecord *signedPreKeyRecord = [self generateRandomSignedRecord];
+        [signedPreKeyRecord markAsAcceptedByService];
+        [self storeSignedPreKey:signedPreKeyRecord.Id signedPreKeyRecord:signedPreKeyRecord];
+        [self setCurrentSignedPrekeyId:signedPreKeyRecord.Id];
+        NSLog(@"[Loki] Pre keys refreshed successfully.");
+    }
+
     SignedPreKeyRecord *_Nullable signedPreKey = self.currentSignedPreKey;
     if (!signedPreKey) {
         OWSFailDebug(@"Signed prekey is null");