Minor refactoring

SignalServiceKit/src/Loki/API/LokiAPI.swift

@@ -198,14 +198,14 @@ public final class LokiAPI : NSObject {
         // Uses a read/write connection because getting the last message hash value also removes expired messages as needed
         // TODO: This shouldn't be the case; a getter shouldn't have an unexpected side effect
         storage.dbReadWriteConnection.readWrite { transaction in
-            result = storage.getLastMessageHash(forServiceNode: target.address, transaction: transaction)
+            result = storage.getLastMessageHash(forSnode: target.address, transaction: transaction)
         }
         return result
     }
 
     private static func setLastMessageHashValue(for target: LokiAPITarget, hashValue: String, expirationDate: UInt64) {
         storage.dbReadWriteConnection.readWrite { transaction in
-            storage.setLastMessageHash(forServiceNode: target.address, hash: hashValue, expiresAt: expirationDate, transaction: transaction)
+            storage.setLastMessageHash(forSnode: target.address, hash: hashValue, expiresAt: expirationDate, transaction: transaction)
         }
     }
     

SignalServiceKit/src/Loki/Database/OWSPrimaryStorage+Loki.h

@@ -1,21 +1,22 @@
 #import "OWSPrimaryStorage.h"
-#import <AxolotlKit/PreKeyRecord.h>
+
+#import <AxolotlKit/AxolotlExceptions.h>
 #import <AxolotlKit/PreKeyBundle.h>
-#import <YapDatabase/YapDatabase.h>
+#import <AxolotlKit/PreKeyRecord.h>
 #import <Curve25519Kit/Ed25519.h>
-#import <AxolotlKit/AxolotlExceptions.h>
+#import <YapDatabase/YapDatabase.h>
 
 NS_ASSUME_NONNULL_BEGIN
 
 @interface OWSPrimaryStorage (Loki)
 
-# pragma mark - Pre Key for Contact
+# pragma mark - Pre Key Record Management
 
-- (BOOL)hasPreKeyForContact:(NSString *)pubKey;
-- (PreKeyRecord *_Nullable)getPreKeyForContact:(NSString *)pubKey transaction:(YapDatabaseReadTransaction *)transaction;
-- (PreKeyRecord *)getOrCreatePreKeyForContact:(NSString *)pubKey;
+- (BOOL)hasPreKeyRecordForContact:(NSString *)hexEncodedPublicKey;
+- (PreKeyRecord *_Nullable)getPreKeyRecordForContact:(NSString *)hexEncodedPublicKey transaction:(YapDatabaseReadTransaction *)transaction;
+- (PreKeyRecord *)getOrCreatePreKeyRecordForContact:(NSString *)hexEncodedPublicKey;
 
-# pragma mark - Pre Key Management
+# pragma mark - Pre Key Bundle Management
 
 /**
  * Generates a pre key bundle for the given contact. Doesn't store the pre key bundle (pre key bundles are supposed to be sent without ever being stored).
@@ -30,16 +31,16 @@ NS_ASSUME_NONNULL_BEGIN
 /**
  * Gets the last message hash and removes it if its `expiresAt` has already passed.
  */
-- (NSString *_Nullable)getLastMessageHashForServiceNode:(NSString *)serviceNode transaction:(YapDatabaseReadWriteTransaction *)transaction;
-- (void)setLastMessageHashForServiceNode:(NSString *)serviceNode hash:(NSString *)hash expiresAt:(u_int64_t)expiresAt transaction:(YapDatabaseReadWriteTransaction *)transaction NS_SWIFT_NAME(setLastMessageHash(forServiceNode:hash:expiresAt:transaction:));
+- (NSString *_Nullable)getLastMessageHashForSnode:(NSString *)snode transaction:(YapDatabaseReadWriteTransaction *)transaction;
+- (void)setLastMessageHashForSnode:(NSString *)snode hash:(NSString *)hash expiresAt:(u_int64_t)expiresAt transaction:(YapDatabaseReadWriteTransaction *)transaction NS_SWIFT_NAME(setLastMessageHash(forSnode:hash:expiresAt:transaction:));
 
-# pragma mark - Group Chat
+# pragma mark - Open Groups
 
 - (void)setIDForMessageWithServerID:(NSUInteger)serverID to:(NSString *)messageID in:(YapDatabaseReadWriteTransaction *)transaction;
 - (NSString *_Nullable)getIDForMessageWithServerID:(NSUInteger)serverID in:(YapDatabaseReadTransaction *)transaction;
 - (void)updateMessageIDCollectionByPruningMessagesWithIDs:(NSSet<NSString *> *)targetMessageIDs in:(YapDatabaseReadWriteTransaction *)transaction NS_SWIFT_NAME(updateMessageIDCollectionByPruningMessagesWithIDs(_:in:));
 
-# pragma mark - Restoration
+# pragma mark - Restoration from Seed
 
 - (void)setRestorationTime:(NSTimeInterval)time;
 - (NSTimeInterval)getRestorationTime;

SignalServiceKit/src/Loki/Database/OWSPrimaryStorage+Loki.m

@@ -17,9 +17,6 @@
 
 # pragma mark - Convenience
 
-#define OWSPrimaryStoragePreKeyStoreCollection @"TSStorageManagerPreKeyStoreCollection"
-#define LKPreKeyContactCollection @"LKPreKeyContactCollection"
-
 - (OWSIdentityManager *)identityManager {
     return OWSIdentityManager.sharedManager;
 }
@@ -28,57 +25,59 @@
     return TSAccountManager.sharedInstance;
 }
 
-# pragma mark - Pre Key for Contact
+# pragma mark - Pre Key Record Management
+
+#define LKPreKeyContactCollection @"LKPreKeyContactCollection"
+#define OWSPrimaryStoragePreKeyStoreCollection @"TSStorageManagerPreKeyStoreCollection"
 
-- (BOOL)hasPreKeyForContact:(NSString *)pubKey {
+- (BOOL)hasPreKeyRecordForContact:(NSString *)pubKey {
     int preKeyId = [self.dbReadWriteConnection intForKey:pubKey inCollection:LKPreKeyContactCollection];
     return preKeyId > 0;
 }
 
-- (PreKeyRecord *_Nullable)getPreKeyForContact:(NSString *)pubKey transaction:(YapDatabaseReadTransaction *)transaction {
-    OWSAssertDebug(pubKey.length > 0);
-    int preKeyId = [transaction intForKey:pubKey inCollection:LKPreKeyContactCollection];
-    
-    // If we don't have an id then return nil
-    if (preKeyId <= 0) { return nil; }
+- (PreKeyRecord *_Nullable)getPreKeyRecordForContact:(NSString *)hexEncodedPublicKey transaction:(YapDatabaseReadTransaction *)transaction {
+    OWSAssertDebug(hexEncodedPublicKey.length > 0);
+    int preKeyID = [transaction intForKey:hexEncodedPublicKey inCollection:LKPreKeyContactCollection];
+
+    if (preKeyID <= 0) { return nil; }
     
-    /// throws_loadPreKey doesn't allow us to pass transaction ;(
-    return [transaction preKeyRecordForKey:[self keyFromInt:preKeyId] inCollection:OWSPrimaryStoragePreKeyStoreCollection];
+    // throws_loadPreKey doesn't allow us to pass transaction
+    // FIXME: This seems like it could be a pretty big issue?
+    return [transaction preKeyRecordForKey:[self keyFromInt:preKeyID] inCollection:OWSPrimaryStoragePreKeyStoreCollection];
 }
 
-- (PreKeyRecord *)getOrCreatePreKeyForContact:(NSString *)pubKey {
-    OWSAssertDebug(pubKey.length > 0);
-    int preKeyId = [self.dbReadWriteConnection intForKey:pubKey inCollection:LKPreKeyContactCollection];
+- (PreKeyRecord *)getOrCreatePreKeyRecordForContact:(NSString *)hexEncodedPublicKey {
+    OWSAssertDebug(hexEncodedPublicKey.length > 0);
+    int preKeyID = [self.dbReadWriteConnection intForKey:hexEncodedPublicKey inCollection:LKPreKeyContactCollection];
     
-    // If we don't have an id then generate and store a new one
-    if (preKeyId <= 0) {
-        return [self generateAndStorePreKeyForContact:pubKey];
+    // If we don't have an ID then generate and store a new one
+    if (preKeyID <= 0) {
+        return [self generateAndStorePreKeyRecordForContact:hexEncodedPublicKey];
     }
     
-    // Load the prekey otherwise just generate a new one
+    // Load existing pre key record if possible; generate a new one otherwise
     @try {
-        return [self throws_loadPreKey:preKeyId];
+        return [self throws_loadPreKey:preKeyID];
     } @catch (NSException *exception) {
-        return [self generateAndStorePreKeyForContact:pubKey];
+        return [self generateAndStorePreKeyRecordForContact:hexEncodedPublicKey];
     }
 }
 
-/// Generate prekey for a contact and store it
-- (PreKeyRecord *)generateAndStorePreKeyForContact:(NSString *)pubKey {
-    [LKLogger print:[NSString stringWithFormat:@"[Loki] Generating new pre key for: %@.", pubKey]];
+- (PreKeyRecord *)generateAndStorePreKeyRecordForContact:(NSString *)pubKey {
+    [LKLogger print:[NSString stringWithFormat:@"[Loki] Generating new pre key record for: %@.", pubKey]];
     OWSAssertDebug(pubKey.length > 0);
     
     NSArray<PreKeyRecord *> *records = [self generatePreKeyRecords:1];
-    [self storePreKeyRecords:records];
-    
     OWSAssertDebug(records.count > 0);
+    [self storePreKeyRecords:records];
+
     PreKeyRecord *record = records.firstObject;
     [self.dbReadWriteConnection setInt:record.Id forKey:pubKey inCollection:LKPreKeyContactCollection];
     
     return record;
 }
 
-# pragma mark - Pre Key Management
+# pragma mark - Pre Key Bundle Management
 
 #define LKPreKeyBundleCollection @"LKPreKeyBundleCollection"
 
@@ -103,7 +102,7 @@
         OWSFailDebug(@"Signed pre key is nil.");
     }
     
-    PreKeyRecord *preKey = [self getOrCreatePreKeyForContact:hexEncodedPublicKey];
+    PreKeyRecord *preKey = [self getOrCreatePreKeyRecordForContact:hexEncodedPublicKey];
     uint32_t registrationID = [self.accountManager getOrGenerateRegistrationId];
     
     PreKeyBundle *bundle = [[PreKeyBundle alloc] initWithRegistrationId:registrationID
@@ -159,35 +158,34 @@
 
 #define LKLastMessageHashCollection @"LKLastMessageHashCollection"
 
-- (NSString *_Nullable)getLastMessageHashForServiceNode:(NSString *)serviceNode transaction:(YapDatabaseReadWriteTransaction *)transaction {
-    NSDictionary *_Nullable dict = [transaction objectForKey:serviceNode inCollection:LKLastMessageHashCollection];
-    if (!dict) { return nil; }
+- (NSString *_Nullable)getLastMessageHashForSnode:(NSString *)snode transaction:(YapDatabaseReadWriteTransaction *)transaction {
+    NSDictionary *_Nullable dict = [transaction objectForKey:snode inCollection:LKLastMessageHashCollection];
+    if (dict == nil) { return nil; }
     
     NSString *_Nullable hash = dict[@"hash"];
-    if (!hash) { return nil; }
+    if (hash == nil) { return nil; }
     
-    // Check if the hash isn't expired
+    // Check if the hash has expired
     uint64_t now = NSDate.ows_millisecondTimeStamp;
     NSNumber *_Nullable expiresAt = dict[@"expiresAt"];
     if (expiresAt && expiresAt.unsignedLongLongValue <= now) {
-        // The last message has expired from the storage server
-        [self removeLastMessageHashForServiceNode:serviceNode transaction:transaction];
+        [self removeLastMessageHashForSnode:snode transaction:transaction];
         return nil;
     }
     
     return hash;
 }
 
-- (void)setLastMessageHashForServiceNode:(NSString *)serviceNode hash:(NSString *)hash expiresAt:(u_int64_t)expiresAt transaction:(YapDatabaseReadWriteTransaction *)transaction {
+- (void)setLastMessageHashForSnode:(NSString *)snode hash:(NSString *)hash expiresAt:(u_int64_t)expiresAt transaction:(YapDatabaseReadWriteTransaction *)transaction {
     NSDictionary *dict = @{ @"hash" : hash, @"expiresAt": @(expiresAt) };
-    [transaction setObject:dict forKey:serviceNode inCollection:LKLastMessageHashCollection];
+    [transaction setObject:dict forKey:snode inCollection:LKLastMessageHashCollection];
 }
 
-- (void)removeLastMessageHashForServiceNode:(NSString *)serviceNode transaction:(YapDatabaseReadWriteTransaction *)transaction {
-    [transaction removeObjectForKey:serviceNode inCollection:LKLastMessageHashCollection];
+- (void)removeLastMessageHashForSnode:(NSString *)snode transaction:(YapDatabaseReadWriteTransaction *)transaction {
+    [transaction removeObjectForKey:snode inCollection:LKLastMessageHashCollection];
 }
 
-# pragma mark - Group Chat
+# pragma mark - Open Groups
 
 #define LKMessageIDCollection @"LKMessageIDCollection"
 
@@ -212,7 +210,7 @@
     [transaction removeObjectsForKeys:serverIDs inCollection:LKMessageIDCollection];
 }
 
-# pragma mark - Restoration
+# pragma mark - Restoration from Seed
 
 #define LKGeneralCollection @"Loki"
 

SignalServiceKit/src/Loki/Protocol/Session Management/LokiSessionResetImplementation.swift

@@ -20,7 +20,7 @@ public class LokiSessionResetImplementation : NSObject, SessionResetProtocol {
             return
         }
         guard let preKeyMessage = whisperMessage as? PreKeyWhisperMessage else { return }
-        guard let storedPreKey = storage.getPreKey(forContact: recipientID, transaction: transaction) else {
+        guard let storedPreKey = storage.getPreKeyRecord(forContact: recipientID, transaction: transaction) else {
             print("[Loki] Received a friend request from a public key for which no pre key bundle was created.")
             throw Errors.invalidPreKey
         }

SignalServiceKit/src/Messages/OWSMessageManager.m

@@ -457,9 +457,9 @@ NS_ASSUME_NONNULL_BEGIN
         } else if (contentProto.dataMessage) {
 
             // Loki: Don't process session request messages any further
-            if ([LKSessionManagementProtocol isSessionRequestMessage:dataMessage]) { return; }
+            if ([LKSessionManagementProtocol isSessionRequestMessage:contentProto.dataMessage]) { return; }
             // Loki: Don't process session restore messages any further
-            if ([LKSessionManagementProtocol isSessionRestoreMessage:dataMessage]) { return; }
+            if ([LKSessionManagementProtocol isSessionRestoreMessage:contentProto.dataMessage]) { return; }
 
             [self handleIncomingEnvelope:envelope
                          withDataMessage:contentProto.dataMessage