From a90aa82710dbc73e4f40c38e076c68bb03b31d7e Mon Sep 17 00:00:00 2001
From: Mikunj <mikunj@live.com.au>
Date: Fri, 10 May 2019 15:22:03 +1000
Subject: [PATCH] Added PreKeyRecord <-> Contact pubkey mapping.

---
 Pods                                          |  2 +-
 .../Loki/Extensions/OWSPrimaryStorage+Loki.h  | 29 +++++++++++
 .../Loki/Extensions/OWSPrimaryStorage+Loki.m  | 48 +++++++++++++++++++
 3 files changed, 78 insertions(+), 1 deletion(-)
 create mode 100644 SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.h
 create mode 100644 SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.m

diff --git a/Pods b/Pods
index b3aff1eb5..21ee23a8e 160000
--- a/Pods
+++ b/Pods
@@ -1 +1 @@
-Subproject commit b3aff1eb5bab24ca0bfc545e9b8b1589fa79abe3
+Subproject commit 21ee23a8e22e91459835b682d821f862f75be233
diff --git a/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.h b/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.h
new file mode 100644
index 000000000..83986879d
--- /dev/null
+++ b/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.h
@@ -0,0 +1,29 @@
+#import "OWSPrimaryStorage.h"
+#import "PreKeyRecord.h"
+
+NS_ASSUME_NONNULL_BEGIN
+
+@interface OWSPrimaryStorage (Loki)
+
+# pragma mark - Prekey for contacts
+
+/**
+ Check if we have PreKeyRecord for the given contact.
+
+ @param pubKey The hex encoded public ket of the contact.
+ @return Whether we have a prekey or not.
+ */
+- (BOOL)hasPreKeyForContact:(NSString *)pubKey;
+
+/**
+ Get the PreKeyRecord associated with the given contact.
+ If the record doesn't exist then this will generate a new one.
+
+ @param pubKey The hex encoded public key of the contact.
+ @return The record associated with the contact.
+ */
+- (PreKeyRecord *)getPreKeyForContact:(NSString *)pubKey;
+
+@end
+
+NS_ASSUME_NONNULL_END
diff --git a/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.m b/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.m
new file mode 100644
index 000000000..a517d90e6
--- /dev/null
+++ b/SignalServiceKit/src/Loki/Extensions/OWSPrimaryStorage+Loki.m
@@ -0,0 +1,48 @@
+#import "OWSPrimaryStorage+Loki.h"
+#import "OWSPrimaryStorage+PreKeyStore.h"
+#import "YapDatabaseConnection+OWS.h"
+
+#define LokiPreKeyContactCollection @"LokiPreKeyContactCollection"
+
+@implementation OWSPrimaryStorage (Loki)
+
+# pragma mark - Prekey for contacts
+
+- (BOOL)hasPreKeyForContact:(NSString *)pubKey {
+    int preKeyId = [self.dbReadWriteConnection intForKey:pubKey inCollection:LokiPreKeyContactCollection];
+    return preKeyId > 0;
+}
+
+- (PreKeyRecord *)getPreKeyForContact:(NSString *)pubKey {
+    OWSAssertDebug(pubKey.length > 0);
+    int preKeyId = [self.dbReadWriteConnection intForKey:pubKey inCollection:LokiPreKeyContactCollection];
+    
+    // If we don't have an id then generate and store a new one
+    if (preKeyId < 1) {
+        return [self generateAndStorePreKeyForContact:pubKey];
+    }
+    
+    // Load the pre key otherwise just generate a new one
+    @try {
+        return [self throws_loadPreKey:preKeyId];
+    } @catch (NSException *exception) {
+        OWSLogWarn(@"[Loki] New prekey had to be generated for %@", pubKey);
+        return [self generateAndStorePreKeyForContact:pubKey];
+    }
+}
+
+/// Generate prekey for a contact and store it
+- (PreKeyRecord *)generateAndStorePreKeyForContact:(NSString *)pubKey {
+    OWSAssertDebug(pubKey.length > 0);
+    
+    NSArray<PreKeyRecord *> *records = [self generatePreKeyRecords:1];
+    [self storePreKeyRecords:records];
+    
+    OWSAssertDebug(records.count > 0);
+    PreKeyRecord *record = [records firstObject];
+    [self.dbReadWriteConnection setInt:record.Id forKey:pubKey inCollection:LokiPreKeyContactCollection];
+    
+    return record;
+}
+
+@end