113 lines
3.9 KiB
Objective-C
113 lines
3.9 KiB
Objective-C
//
|
|
// Copyright (c) 2019 Open Whisper Systems. All rights reserved.
|
|
//
|
|
|
|
#import "OWSPrimaryStorage+PreKeyStore.h"
|
|
#import "OWSPrimaryStorage+keyFromIntLong.h"
|
|
#import "TSStorageKeys.h"
|
|
#import "YapDatabaseConnection+OWS.h"
|
|
#import <SessionProtocolKit/SessionProtocolKit.h>
|
|
|
|
#define OWSPrimaryStoragePreKeyStoreCollection @"TSStorageManagerPreKeyStoreCollection"
|
|
#define TSNextPrekeyIdKey @"TSStorageInternalSettingsNextPreKeyId"
|
|
#define BATCH_SIZE 100
|
|
|
|
NS_ASSUME_NONNULL_BEGIN
|
|
|
|
@implementation OWSPrimaryStorage (PreKeyStore)
|
|
|
|
- (NSArray<PreKeyRecord *> *)generatePreKeyRecords
|
|
{
|
|
return [self generatePreKeyRecords:BATCH_SIZE];
|
|
}
|
|
|
|
- (NSArray<PreKeyRecord *> *)generatePreKeyRecords:(int)batchSize
|
|
{
|
|
NSMutableArray *preKeyRecords = [NSMutableArray array];
|
|
|
|
@synchronized(self)
|
|
{
|
|
int preKeyId = [self nextPreKeyId:batchSize];
|
|
|
|
OWSLogInfo(@"building %d new preKeys starting from preKeyId: %d", batchSize, preKeyId);
|
|
for (int i = 0; i < batchSize; i++) {
|
|
ECKeyPair *keyPair = [Curve25519 generateKeyPair];
|
|
PreKeyRecord *record = [[PreKeyRecord alloc] initWithId:preKeyId keyPair:keyPair];
|
|
|
|
[preKeyRecords addObject:record];
|
|
preKeyId++;
|
|
}
|
|
|
|
[self.dbReadWriteConnection setInt:preKeyId
|
|
forKey:TSNextPrekeyIdKey
|
|
inCollection:TSStorageInternalSettingsCollection];
|
|
}
|
|
return preKeyRecords;
|
|
}
|
|
|
|
- (void)storePreKeyRecords:(NSArray<PreKeyRecord *> *)preKeyRecords
|
|
{
|
|
for (PreKeyRecord *record in preKeyRecords) {
|
|
[self.dbReadWriteConnection setObject:record
|
|
forKey:[self keyFromInt:record.Id]
|
|
inCollection:OWSPrimaryStoragePreKeyStoreCollection];
|
|
}
|
|
}
|
|
|
|
- (PreKeyRecord *)throws_loadPreKey:(int)preKeyId
|
|
{
|
|
PreKeyRecord *preKeyRecord = [self.dbReadConnection preKeyRecordForKey:[self keyFromInt:preKeyId]
|
|
inCollection:OWSPrimaryStoragePreKeyStoreCollection];
|
|
|
|
if (!preKeyRecord) {
|
|
OWSRaiseException(InvalidKeyIdException, @"No pre key found matching key id");
|
|
} else {
|
|
return preKeyRecord;
|
|
}
|
|
}
|
|
|
|
- (void)storePreKey:(int)preKeyId preKeyRecord:(PreKeyRecord *)record
|
|
{
|
|
[self.dbReadWriteConnection setObject:record
|
|
forKey:[self keyFromInt:preKeyId]
|
|
inCollection:OWSPrimaryStoragePreKeyStoreCollection];
|
|
}
|
|
|
|
- (BOOL)containsPreKey:(int)preKeyId
|
|
{
|
|
PreKeyRecord *preKeyRecord = [self.dbReadConnection preKeyRecordForKey:[self keyFromInt:preKeyId]
|
|
inCollection:OWSPrimaryStoragePreKeyStoreCollection];
|
|
return (preKeyRecord != nil);
|
|
}
|
|
|
|
- (void)removePreKey:(int)preKeyId protocolContext:(nullable id)protocolContext
|
|
{
|
|
if ([protocolContext isKindOfClass:YapDatabaseReadWriteTransaction.class]) {
|
|
[(YapDatabaseReadWriteTransaction *)protocolContext removeObjectForKey:[self keyFromInt:preKeyId] inCollection:OWSPrimaryStoragePreKeyStoreCollection];
|
|
} else {
|
|
[self.dbReadWriteConnection removeObjectForKey:[self keyFromInt:preKeyId] inCollection:OWSPrimaryStoragePreKeyStoreCollection];
|
|
}
|
|
}
|
|
|
|
- (int)nextPreKeyId:(int)batchSize
|
|
{
|
|
int lastPreKeyId =
|
|
[self.dbReadConnection intForKey:TSNextPrekeyIdKey inCollection:TSStorageInternalSettingsCollection];
|
|
|
|
if (lastPreKeyId < 1) {
|
|
// One-time prekey ids must be > 0 and < kPreKeyOfLastResortId.
|
|
lastPreKeyId = 1 + arc4random_uniform(kPreKeyOfLastResortId - (batchSize + 1));
|
|
} else if (lastPreKeyId > kPreKeyOfLastResortId - batchSize) {
|
|
// We want to "overflow" to 1 when we reach the "prekey of last resort" id
|
|
// to avoid biasing towards higher values.
|
|
lastPreKeyId = 1;
|
|
}
|
|
OWSCAssertDebug(lastPreKeyId > 0 && lastPreKeyId < kPreKeyOfLastResortId);
|
|
|
|
return lastPreKeyId;
|
|
}
|
|
|
|
@end
|
|
|
|
NS_ASSUME_NONNULL_END
|