diff --git a/SignalServiceKit/src/Loki/API/LokiAPI.swift b/SignalServiceKit/src/Loki/API/LokiAPI.swift index f53fa2564..ccd3f277a 100644 --- a/SignalServiceKit/src/Loki/API/LokiAPI.swift +++ b/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) } } diff --git a/SignalServiceKit/src/Loki/Database/OWSPrimaryStorage+Loki.h b/SignalServiceKit/src/Loki/Database/OWSPrimaryStorage+Loki.h index 44dc2cf43..fb8d51ac1 100644 --- a/SignalServiceKit/src/Loki/Database/OWSPrimaryStorage+Loki.h +++ b/SignalServiceKit/src/Loki/Database/OWSPrimaryStorage+Loki.h @@ -1,21 +1,22 @@ #import "OWSPrimaryStorage.h" -#import -#import -#import -#import + #import +#import +#import +#import +#import 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 *)targetMessageIDs in:(YapDatabaseReadWriteTransaction *)transaction NS_SWIFT_NAME(updateMessageIDCollectionByPruningMessagesWithIDs(_:in:)); -# pragma mark - Restoration +# pragma mark - Restoration from Seed - (void)setRestorationTime:(NSTimeInterval)time; - (NSTimeInterval)getRestorationTime; diff --git a/SignalServiceKit/src/Loki/Database/OWSPrimaryStorage+Loki.m b/SignalServiceKit/src/Loki/Database/OWSPrimaryStorage+Loki.m index b539fc705..fa41abbfd 100644 --- a/SignalServiceKit/src/Loki/Database/OWSPrimaryStorage+Loki.m +++ b/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 -- (BOOL)hasPreKeyForContact:(NSString *)pubKey { +#define LKPreKeyContactCollection @"LKPreKeyContactCollection" +#define OWSPrimaryStoragePreKeyStoreCollection @"TSStorageManagerPreKeyStoreCollection" + +- (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]; +- (PreKeyRecord *_Nullable)getPreKeyRecordForContact:(NSString *)hexEncodedPublicKey transaction:(YapDatabaseReadTransaction *)transaction { + OWSAssertDebug(hexEncodedPublicKey.length > 0); + int preKeyID = [transaction intForKey:hexEncodedPublicKey inCollection:LKPreKeyContactCollection]; + + if (preKeyID <= 0) { return nil; } - // If we don't have an id then return nil - 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 *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" diff --git a/SignalServiceKit/src/Loki/Protocol/Session Management/LokiSessionResetImplementation.swift b/SignalServiceKit/src/Loki/Protocol/Session Management/LokiSessionResetImplementation.swift index e5bfd99b7..7c08d0d1e 100644 --- a/SignalServiceKit/src/Loki/Protocol/Session Management/LokiSessionResetImplementation.swift +++ b/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 } diff --git a/SignalServiceKit/src/Messages/OWSMessageManager.m b/SignalServiceKit/src/Messages/OWSMessageManager.m index 9c9aa40fd..0911fec3c 100644 --- a/SignalServiceKit/src/Messages/OWSMessageManager.m +++ b/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