Minor refactoring

This commit is contained in:
nielsandriesse 2020-04-29 09:02:45 +10:00
parent d6747ebb14
commit b0233e4cc0
5 changed files with 59 additions and 60 deletions

View file

@ -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)
}
}

View file

@ -1,21 +1,22 @@
#import "OWSPrimaryStorage.h"
#import <AxolotlKit/PreKeyRecord.h>
#import <AxolotlKit/PreKeyBundle.h>
#import <YapDatabase/YapDatabase.h>
#import <Curve25519Kit/Ed25519.h>
#import <AxolotlKit/AxolotlExceptions.h>
#import <AxolotlKit/PreKeyBundle.h>
#import <AxolotlKit/PreKeyRecord.h>
#import <Curve25519Kit/Ed25519.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;

View file

@ -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<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"

View file

@ -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
}

View file

@ -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