From 70bd9350b8e7669ebd4cb809ed6d79b77f755a9f Mon Sep 17 00:00:00 2001
From: Mikunj <mikunj@live.com.au>
Date: Tue, 18 Feb 2020 12:28:49 +1100
Subject: [PATCH] Ensure generated PreKeyBundle is valid

---
 .../jobs/MultiDeviceGroupUpdateJob.java       |  5 +++--
 .../securesms/loki/MultiDeviceUtilities.kt    |  2 +-
 .../messaging/LokiPreKeyBundleDatabase.kt     | 22 ++++++++++++++++++-
 3 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/src/org/thoughtcrime/securesms/jobs/MultiDeviceGroupUpdateJob.java b/src/org/thoughtcrime/securesms/jobs/MultiDeviceGroupUpdateJob.java
index 9ef69669ea..48afa32221 100644
--- a/src/org/thoughtcrime/securesms/jobs/MultiDeviceGroupUpdateJob.java
+++ b/src/org/thoughtcrime/securesms/jobs/MultiDeviceGroupUpdateJob.java
@@ -125,7 +125,9 @@ public class MultiDeviceGroupUpdateJob extends BaseJob implements InjectableType
 
   @Override
   public boolean onShouldRetry(@NonNull Exception exception) {
-    return exception instanceof PushNetworkException;
+    // Loki - Disabled because we have our own retrying
+    // if (exception instanceof PushNetworkException) return true;
+    return false;
   }
 
   @Override
@@ -143,7 +145,6 @@ public class MultiDeviceGroupUpdateJob extends BaseJob implements InjectableType
                                                                               .withLength(contactsFile.length())
                                                                               .build();
 
-    // TODO: Message ID
     messageSender.sendMessage(0, SignalServiceSyncMessage.forGroups(attachmentStream),
                               UnidentifiedAccessUtil.getAccessForSync(context));
   }
diff --git a/src/org/thoughtcrime/securesms/loki/MultiDeviceUtilities.kt b/src/org/thoughtcrime/securesms/loki/MultiDeviceUtilities.kt
index a3e08cc8a5..116f2a8bd2 100644
--- a/src/org/thoughtcrime/securesms/loki/MultiDeviceUtilities.kt
+++ b/src/org/thoughtcrime/securesms/loki/MultiDeviceUtilities.kt
@@ -128,7 +128,7 @@ fun sendDeviceLinkMessage(context: Context, hexEncodedPublicKey: String, deviceL
     }
     Promise.ofSuccess(Unit)
   } catch (e: Exception) {
-    Log.d("Loki", "Failed to send device link message to: $hexEncodedPublicKey.")
+    Log.d("Loki", "Failed to send device link message to $hexEncodedPublicKey: $e")
     Promise.ofFail(e)
   }
 }
diff --git a/src/org/thoughtcrime/securesms/loki/redesign/messaging/LokiPreKeyBundleDatabase.kt b/src/org/thoughtcrime/securesms/loki/redesign/messaging/LokiPreKeyBundleDatabase.kt
index 22a4ee1e7c..ed464c7a82 100644
--- a/src/org/thoughtcrime/securesms/loki/redesign/messaging/LokiPreKeyBundleDatabase.kt
+++ b/src/org/thoughtcrime/securesms/loki/redesign/messaging/LokiPreKeyBundleDatabase.kt
@@ -16,6 +16,7 @@ import org.thoughtcrime.securesms.loki.redesign.utilities.insertOrUpdate
 import org.thoughtcrime.securesms.util.Base64
 import org.thoughtcrime.securesms.util.TextSecurePreferences
 import org.whispersystems.libsignal.IdentityKey
+import org.whispersystems.libsignal.InvalidKeyException
 import org.whispersystems.libsignal.ecc.Curve
 import org.whispersystems.libsignal.state.PreKeyBundle
 import org.whispersystems.libsignal.util.KeyHelper
@@ -41,6 +42,25 @@ class LokiPreKeyBundleDatabase(context: Context, helper: SQLCipherOpenHelper) :
     }
 
     fun generatePreKeyBundle(hexEncodedPublicKey: String): PreKeyBundle? {
+        var failureCount = 0
+        while (failureCount < 3) {
+            try {
+                val preKey = generatePreKeyBundle(hexEncodedPublicKey, failureCount > 0) ?: return null
+                // Verify the bundle is correct
+                if (!Curve.verifySignature(preKey.identityKey.publicKey, preKey.signedPreKey.serialize(), preKey.signedPreKeySignature)) {
+                    throw InvalidKeyException()
+                }
+                return preKey;
+            } catch (e: InvalidKeyException) {
+                failureCount += 1
+            }
+        }
+        Log.w("Loki", "Failed to generate a valid PreKeyBundle for $hexEncodedPublicKey")
+        return null
+    }
+
+    private fun generatePreKeyBundle(hexEncodedPublicKey: String, forceClean: Boolean): PreKeyBundle? {
+        if (hexEncodedPublicKey.isEmpty()) return null
         var registrationID = TextSecurePreferences.getLocalRegistrationId(context)
         if (registrationID == 0) {
             registrationID = KeyHelper.generateRegistrationId(false)
@@ -49,7 +69,7 @@ class LokiPreKeyBundleDatabase(context: Context, helper: SQLCipherOpenHelper) :
         val deviceID = SignalServiceAddress.DEFAULT_DEVICE_ID
         val preKeyRecord = DatabaseFactory.getLokiPreKeyRecordDatabase(context).getOrCreatePreKeyRecord(hexEncodedPublicKey)
         val identityKeyPair = IdentityKeyUtil.getIdentityKeyPair(context)
-        if (TextSecurePreferences.isSignedPreKeyRegistered(context)) {
+        if (!forceClean && TextSecurePreferences.isSignedPreKeyRegistered(context)) {
             Log.d("Loki", "A signed pre key has already been registered.")
         } else {
             Log.d("Loki", "Registering a new signed pre key.")