Clean up prekey logic.

// FREEBIE

src/Account/TSAccountManager.m

@@ -1,9 +1,5 @@
 //
-//  TSAccountManagement.m
-//  TextSecureKit
-//
-//  Created by Frederic Jacobs on 27/10/14.
-//  Copyright (c) 2014 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
 //
 
 #import "TSAccountManager.h"
@@ -210,7 +206,9 @@ NS_ASSUME_NONNULL_BEGIN
                     [TSStorageManager storeServerToken:authToken signalingKey:signalingKey];
                     [self didRegister];
                     [TSSocketManager becomeActiveFromForeground];
-                    [TSPreKeyManager registerPreKeysWithSuccess:successBlock failure:failureBlock];
+                    [TSPreKeyManager registerPreKeysWithMode:RefreshPreKeysMode_SignedAndOneTime
+                                                     success:successBlock
+                                                     failure:failureBlock];
                     break;
                 }
                 default: {

src/Account/TSPreKeyManager.h

@@ -1,21 +1,31 @@
 //
-//  TSPrekeyManager.h
-//  TextSecureKit
-//
-//  Created by Frederic Jacobs on 07/11/14.
-//  Copyright (c) 2014 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
 //
 
 #import <Foundation/Foundation.h>
 #import "TSAccountManager.h"
 
-// Time before deletion of signed PreKeys (measured in seconds)
-#define SignedPreKeysDeletionTime 14 * 24 * 60 * 60
+typedef NS_ENUM(NSInteger, RefreshPreKeysMode) {
+    // Refresh the signed prekey AND the one-time prekeys.
+    RefreshPreKeysMode_SignedAndOneTime,
+    // Only refresh the signed prekey, which should happen around every 48 hours.
+    //
+    // Most users will refresh their signed prekeys much more often than their
+    // one-time prekeys, so we use a "signed only" mode to avoid updating the
+    // one-time keys in this case.
+    //
+    // We do not need a "one-time only" mode.
+    RefreshPreKeysMode_SignedOnly,
+};
 
 @interface TSPreKeyManager : NSObject
 
-+ (void)registerPreKeysWithSuccess:(void (^)())successHandler failure:(void (^)(NSError *error))failureHandler;
++ (void)registerPreKeysWithMode:(RefreshPreKeysMode)mode
+                        success:(void (^)())successHandler
+                        failure:(void (^)(NSError *error))failureHandler;
 
 + (void)refreshPreKeys;
 
++ (void)checkPreKeysIfNecessary;
+
 @end

src/Account/TSPreKeyManager.m

@@ -1,79 +1,185 @@
 //
-//  TSPrekeyManager.m
-//  TextSecureKit
-//
-//  Created by Frederic Jacobs on 07/11/14.
-//  Copyright (c) 2014 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
 //
 
 #import "TSPreKeyManager.h"
-
 #import "NSURLSessionDataTask+StatusCode.h"
 #import "TSNetworkManager.h"
+#import "TSRegisterSignedPrekeyRequest.h"
 #import "TSStorageHeaders.h"
 
-#define EPHEMERAL_PREKEYS_MINIMUM 15
+// Time before deletion of signed prekeys (measured in seconds)
+//
+// Currently we retain signed prekeys for at least 14 days.
+static const CGFloat kSignedPreKeysDeletionTime = 14 * 24 * 60 * 60;
+
+// Time before rotation of signed prekeys (measured in seconds)
+//
+// Currently we rotate signed prekeys every 2 days (48 hours).
+static const CGFloat kSignedPreKeysRotationTime = 2 * 24 * 60 * 60;
+
+// How often we check prekey state on app activation.
+//
+// Currently we check prekey state every 12 hours.
+static const CGFloat kPreKeyCheckFrequencySeconds = 12 * 60 * 60;
+
+// We generate 100 one-time prekeys at a time.  We should replenish
+// whenever ~2/3 of them have been consumed.
+static const NSUInteger kEphemeralPreKeysMinimumCount = 35;
+
+// This global should only be accessed on prekeyQueue.
+static NSDate *lastPreKeyCheckTimestamp = nil;
 
 @implementation TSPreKeyManager
 
-+ (void)registerPreKeysWithSuccess:(void (^)())success failure:(void (^)(NSError *))failureBlock
+// We should never dispatch sync to this queue.
++ (dispatch_queue_t)prekeyQueue
 {
-    TSStorageManager *storageManager = [TSStorageManager sharedManager];
-    ECKeyPair *identityKeyPair       = [storageManager identityKeyPair];
+    static dispatch_once_t onceToken;
+    static dispatch_queue_t queue;
+    dispatch_once(&onceToken, ^{
+        queue = dispatch_queue_create("org.whispersystems.signal.prekeyQueue", NULL);
+    });
+    return queue;
+}
 
-    if (!identityKeyPair) {
-        [storageManager generateNewIdentityKey];
-        identityKeyPair = [storageManager identityKeyPair];
-    }
++ (void)checkPreKeysIfNecessary
+{
+    OWSAssert([UIApplication sharedApplication].applicationState == UIApplicationStateActive);
 
-    PreKeyRecord *lastResortPreKey   = [storageManager getOrGenerateLastResortKey];
-    SignedPreKeyRecord *signedPreKey = [storageManager generateRandomSignedRecord];
+    // Update the prekey check timestamp.
+    dispatch_async(TSPreKeyManager.prekeyQueue, ^{
+        BOOL shouldCheck = (lastPreKeyCheckTimestamp == nil
+            || fabs([lastPreKeyCheckTimestamp timeIntervalSinceNow]) >= kPreKeyCheckFrequencySeconds);
+        if (shouldCheck) {
+            [[TSAccountManager sharedInstance] ifRegistered:YES
+                                                   runAsync:^{
+                                                       [TSPreKeyManager refreshPreKeys];
+                                                   }];
+        }
+    });
+}
 
-    NSArray *preKeys = [storageManager generatePreKeyRecords];
++ (void)registerPreKeysWithMode:(RefreshPreKeysMode)mode
+                        success:(void (^)())successHandler
+                        failure:(void (^)(NSError *error))failureHandler
+{
+    // We use prekeyQueue to serialize this logic and ensure that only
+    // one thread is "registering" or "clearing" prekeys at a time.
+    dispatch_async(TSPreKeyManager.prekeyQueue, ^{
+        RefreshPreKeysMode modeCopy = mode;
+        TSStorageManager *storageManager = [TSStorageManager sharedManager];
+        ECKeyPair *identityKeyPair = [storageManager identityKeyPair];
+
+        if (!identityKeyPair) {
+            [storageManager generateNewIdentityKey];
+            identityKeyPair = [storageManager identityKeyPair];
 
-    TSRegisterPrekeysRequest *request =
-        [[TSRegisterPrekeysRequest alloc] initWithPrekeyArray:preKeys
-                                                  identityKey:[storageManager identityKeyPair].publicKey
-                                           signedPreKeyRecord:signedPreKey
-                                             preKeyLastResort:lastResortPreKey];
+            // Switch modes if necessary.
+            modeCopy = RefreshPreKeysMode_SignedAndOneTime;
+        }
 
-    [[TSNetworkManager sharedManager] makeRequest:request
-        success:^(NSURLSessionDataTask *task, id responseObject) {
-            DDLogInfo(@"%@ Successfully registered pre keys.", self.tag);
-            [storageManager storePreKeyRecords:preKeys];
-            [storageManager storeSignedPreKey:signedPreKey.Id signedPreKeyRecord:signedPreKey];
+        SignedPreKeyRecord *signedPreKey = [storageManager generateRandomSignedRecord];
 
-            success();
+        NSArray *preKeys = nil;
+        TSRequest *request;
+        NSString *description;
+        if (modeCopy == RefreshPreKeysMode_SignedAndOneTime) {
+            description = @"signed and one-time prekeys";
+            PreKeyRecord *lastResortPreKey = [storageManager getOrGenerateLastResortKey];
+            preKeys = [storageManager generatePreKeyRecords];
+            request = [[TSRegisterPrekeysRequest alloc] initWithPrekeyArray:preKeys
+                                                                identityKey:[storageManager identityKeyPair].publicKey
+                                                         signedPreKeyRecord:signedPreKey
+                                                           preKeyLastResort:lastResortPreKey];
+        } else {
+            description = @"just signed prekey";
+            request = [[TSRegisterSignedPrekeyRequest alloc] initWithSignedPreKeyRecord:signedPreKey];
         }
-        failure:^(NSURLSessionDataTask *task, NSError *error) {
-            DDLogError(@"%@ Failed to register pre keys.", self.tag);
-            failureBlock(error);
-        }];
+
+        [[TSNetworkManager sharedManager] makeRequest:request
+            success:^(NSURLSessionDataTask *task, id responseObject) {
+                DDLogInfo(@"%@ Successfully registered %@.", self.tag, description);
+                if (modeCopy == RefreshPreKeysMode_SignedAndOneTime) {
+                    [storageManager storePreKeyRecords:preKeys];
+                }
+                [storageManager storeSignedPreKey:signedPreKey.Id signedPreKeyRecord:signedPreKey];
+
+                successHandler();
+            }
+            failure:^(NSURLSessionDataTask *task, NSError *error) {
+                DDLogError(@"%@ Failed to register %@.", self.tag, description);
+                failureHandler(error);
+            }];
+    });
 }
 
 + (void)refreshPreKeys {
+    // We want to update prekeys if either the one-time or signed prekeys need an update, so
+    // we check the status of both.
+    //
+    // Most users will refresh their signed prekeys much more often than their
+    // one-time PreKeys, so we use a "signed only" mode to avoid updating the
+    // one-time keys in this case.
+    //
+    // We do not need a "one-time only" mode.
     TSAvailablePreKeysCountRequest *preKeyCountRequest = [[TSAvailablePreKeysCountRequest alloc] init];
     [[TSNetworkManager sharedManager] makeRequest:preKeyCountRequest
         success:^(NSURLSessionDataTask *task, NSDictionary *responseObject) {
-          NSString *preKeyCountKey = @"count";
-          NSNumber *count          = [responseObject objectForKey:preKeyCountKey];
-
-          if (count.integerValue > EPHEMERAL_PREKEYS_MINIMUM) {
-              DDLogVerbose(@"Available prekeys sufficient: %@", count.stringValue);
-              return;
-          } else {
-              [self registerPreKeysWithSuccess:^{
-                DDLogInfo(@"New PreKeys registered with server.");
-
-                [self clearSignedPreKeyRecords];
-              }
-                  failure:^(NSError *error) {
-                    DDLogWarn(@"Failed to update prekeys with the server");
-                  }];
-          }
+            NSString *preKeyCountKey = @"count";
+            NSNumber *count = [responseObject objectForKey:preKeyCountKey];
+
+            void (^updatePreKeys)(RefreshPreKeysMode) = ^(RefreshPreKeysMode mode) {
+                [self registerPreKeysWithMode:mode
+                    success:^{
+                        DDLogInfo(@"%@ New prekeys registered with server.", self.tag);
+
+                        [self clearSignedPreKeyRecords];
+                    }
+                    failure:^(NSError *error) {
+                        DDLogWarn(@"%@ Failed to update prekeys with the server", self.tag);
+                    }];
+            };
+
+            BOOL shouldUpdateOneTimePreKeys = count.integerValue <= kEphemeralPreKeysMinimumCount;
+
+            if (shouldUpdateOneTimePreKeys) {
+                DDLogInfo(@"%@ Updating one-time and signed prekeys due to shortage of one-time prekeys.", self.tag);
+                updatePreKeys(RefreshPreKeysMode_SignedAndOneTime);
+            } else {
+                TSRequest *currentSignedPreKey = [[TSCurrentSignedPreKeyRequest alloc] init];
+                [[TSNetworkManager sharedManager] makeRequest:currentSignedPreKey
+                    success:^(NSURLSessionDataTask *task, NSDictionary *responseObject) {
+                        NSString *keyIdDictKey = @"keyId";
+                        NSNumber *keyId = [responseObject objectForKey:keyIdDictKey];
+                        OWSAssert(keyId);
+                        TSStorageManager *storageManager = [TSStorageManager sharedManager];
+                        SignedPreKeyRecord *currentRecord = [storageManager loadSignedPrekey:keyId.intValue];
+                        OWSAssert(currentRecord);
+
+                        BOOL shouldUpdateSignedPrekey
+                            = fabs([currentRecord.generatedAt timeIntervalSinceNow]) >= kSignedPreKeysRotationTime;
+                        if (shouldUpdateSignedPrekey) {
+                            DDLogInfo(@"%@ Updating signed prekey due to rotation period.", self.tag);
+                            updatePreKeys(RefreshPreKeysMode_SignedOnly);
+                        } else {
+                            DDLogDebug(@"%@ Not updating prekeys.", self.tag);
+                        }
+
+                        // Update the prekey check timestamp on success.
+                        dispatch_async(TSPreKeyManager.prekeyQueue, ^{
+                            lastPreKeyCheckTimestamp = [NSDate date];
+                        });
+                    }
+                    failure:^(NSURLSessionDataTask *task, NSError *error) {
+                        DDLogWarn(@"%@ Updating signed prekey because of failure to retrieve current signed prekey.",
+                            self.tag);
+                        updatePreKeys(RefreshPreKeysMode_SignedOnly);
+                    }];
+            }
         }
         failure:^(NSURLSessionDataTask *task, NSError *error) {
-          DDLogError(@"Failed to retrieve the number of available prekeys.");
+            DDLogError(@"%@ Failed to retrieve the number of available prekeys.", self.tag);
         }];
 }
 
@@ -81,37 +187,62 @@
     TSRequest *currentSignedPreKey = [[TSCurrentSignedPreKeyRequest alloc] init];
     [[TSNetworkManager sharedManager] makeRequest:currentSignedPreKey
         success:^(NSURLSessionDataTask *task, NSDictionary *responseObject) {
-          NSString *keyIdDictKey = @"keyId";
-          NSNumber *keyId        = [responseObject objectForKey:keyIdDictKey];
+            NSString *keyIdDictKey = @"keyId";
+            NSNumber *keyId = [responseObject objectForKey:keyIdDictKey];
 
-          [self clearSignedPreKeyRecordsWithKeyId:keyId];
+            [self clearSignedPreKeyRecordsWithKeyId:keyId];
 
         }
         failure:^(NSURLSessionDataTask *task, NSError *error) {
-          DDLogWarn(@"Failed to retrieve current prekey.");
+            DDLogWarn(@"%@ Failed to retrieve current prekey.", self.tag);
         }];
 }
 
 + (void)clearSignedPreKeyRecordsWithKeyId:(NSNumber *)keyId {
     if (!keyId) {
-        DDLogError(@"The server returned an incomplete ");
+        DDLogError(@"%@ The server returned an incomplete response", self.tag);
         return;
     }
 
-    TSStorageManager *storageManager  = [TSStorageManager sharedManager];
-    SignedPreKeyRecord *currentRecord = [storageManager loadSignedPrekey:keyId.intValue];
-    NSArray *allSignedPrekeys         = [storageManager loadSignedPreKeys];
-    NSArray *oldSignedPrekeys         = [self removeCurrentRecord:currentRecord fromRecords:allSignedPrekeys];
+    // We use prekeyQueue to serialize this logic and ensure that only
+    // one thread is "registering" or "clearing" prekeys at a time.
+    dispatch_async(TSPreKeyManager.prekeyQueue, ^{
+        TSStorageManager *storageManager = [TSStorageManager sharedManager];
+        SignedPreKeyRecord *currentRecord = [storageManager loadSignedPrekey:keyId.intValue];
+        NSArray *allSignedPrekeys = [storageManager loadSignedPreKeys];
+        NSArray *oldSignedPrekeys = [self removeCurrentRecord:currentRecord fromRecords:allSignedPrekeys];
+
+        NSDateFormatter *dateFormatter = [[NSDateFormatter alloc] init];
+        dateFormatter.dateStyle = NSDateFormatterMediumStyle;
+        dateFormatter.timeStyle = NSDateFormatterMediumStyle;
+
+        // Sort the signed prekeys in ascending order of generation time.
+        oldSignedPrekeys = [oldSignedPrekeys sortedArrayUsingComparator:^NSComparisonResult(
+            SignedPreKeyRecord *_Nonnull left, SignedPreKeyRecord *_Nonnull right) {
+            return [left.generatedAt compare:right.generatedAt];
+        }];
 
-    if ([oldSignedPrekeys count] > 3) {
+        NSUInteger deletedCount = 0;
+        // Iterate the signed prekeys in ascending order so that we try to delete older keys first.
         for (SignedPreKeyRecord *deletionCandidate in oldSignedPrekeys) {
-            DDLogInfo(@"Old signed prekey record: %@", deletionCandidate.generatedAt);
+            // Always keep the last three signed prekeys.
+            NSUInteger remainingCount = allSignedPrekeys.count - deletedCount;
+            if (remainingCount <= 3) {
+                break;
+            }
 
-            if ([deletionCandidate.generatedAt timeIntervalSinceNow] > SignedPreKeysDeletionTime) {
-                [storageManager removeSignedPreKey:deletionCandidate.Id];
+            // Never delete signed prekeys until they are N days old.
+            if (fabs([deletionCandidate.generatedAt timeIntervalSinceNow]) < kSignedPreKeysDeletionTime) {
+                break;
             }
+
+            DDLogInfo(@"%@ Deleting old signed prekey: %@",
+                self.tag,
+                [dateFormatter stringFromDate:deletionCandidate.generatedAt]);
+            [storageManager removeSignedPreKey:deletionCandidate.Id];
+            deletedCount++;
         }
-    }
+    });
 }
 
 + (NSArray *)removeCurrentRecord:(SignedPreKeyRecord *)currentRecord fromRecords:(NSArray *)allRecords {

src/Messages/PreKeyBundle+jsonDict.m

@@ -1,9 +1,5 @@
 //
-//  PreKeyBundle+jsonDict.m
-//  Signal
-//
-//  Created by Frederic Jacobs on 26/11/14.
-//  Copyright (c) 2014 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
 //
 
 #import "NSData+Base64.h"
@@ -40,15 +36,22 @@
         int registrationId = [registrationIdString intValue];
         int deviceId       = [deviceIdString intValue];
 
-        NSDictionary *preKey = [deviceDict objectForKey:@"preKey"];
+        NSDictionary *_Nullable preKeyDict;
+        id optionalPreKeyDict = [deviceDict objectForKey:@"preKey"];
+        // JSON parsing might give us NSNull, so we can't simply check for non-nil.
+        if ([optionalPreKeyDict isKindOfClass:[NSDictionary class]]) {
+            preKeyDict = (NSDictionary *)optionalPreKeyDict;
+        }
+
         int prekeyId;
-        NSData *preKeyPublic = nil;
+        NSData *_Nullable preKeyPublic;
 
-        if (!preKey) {
+        if (!preKeyDict) {
+            DDLogInfo(@"%@ No one-time prekey included in the bundle.", self.tag);
             prekeyId = -1;
         } else {
-            prekeyId                     = [[preKey objectForKey:@"keyId"] intValue];
-            NSString *preKeyPublicString = [preKey objectForKey:@"publicKey"];
+            prekeyId = [[preKeyDict objectForKey:@"keyId"] intValue];
+            NSString *preKeyPublicString = [preKeyDict objectForKey:@"publicKey"];
             preKeyPublic                 = [NSData dataFromBase64StringNoPadding:preKeyPublicString];
         }
 
@@ -92,4 +95,16 @@
     return bundle;
 }
 
+#pragma mark - Logging
+
++ (NSString *)tag
+{
+    return [NSString stringWithFormat:@"[%@]", self.class];
+}
+
+- (NSString *)tag
+{
+    return self.class.tag;
+}
+
 @end

src/Messages/TSMessagesManager.m

@@ -31,6 +31,7 @@
 #import "TSInfoMessage.h"
 #import "TSInvalidIdentityKeyReceivingErrorMessage.h"
 #import "TSNetworkManager.h"
+#import "TSPreKeyManager.h"
 #import "TSStorageHeaders.h"
 #import "TextSecureKitEnv.h"
 #import <AxolotlKit/AxolotlExceptions.h>
@@ -232,6 +233,7 @@ NS_ASSUME_NONNULL_BEGIN
                      completion:(void (^)(NSError *_Nullable error))completion
 {
     OWSAssert([NSThread isMainThread]);
+
     @synchronized(self) {
         TSStorageManager *storageManager = [TSStorageManager sharedManager];
         NSString *recipientId = preKeyEnvelope.source;
@@ -247,6 +249,9 @@ NS_ASSUME_NONNULL_BEGIN
         dispatch_async([OWSDispatch sessionCipher], ^{
             NSData *plaintextData;
             @try {
+                // Check whether we need to refresh our PreKeys every time we receive a PreKeyWhisperMessage.
+                [TSPreKeyManager refreshPreKeys];
+
                 PreKeyWhisperMessage *message = [[PreKeyWhisperMessage alloc] initWithData:encryptedData];
                 SessionCipher *cipher = [[SessionCipher alloc] initWithSessionStore:storageManager
                                                                         preKeyStore:storageManager

src/Network/API/Requests/TSRegisterSignedPrekeyRequest.h

@@ -0,0 +1,14 @@
+//
+//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
+//
+
+#import "TSRequest.h"
+
+@class SignedPreKeyRecord;
+@class PreKeyRecord;
+
+@interface TSRegisterSignedPrekeyRequest : TSRequest
+
+- (id)initWithSignedPreKeyRecord:(SignedPreKeyRecord *)signedRecord;
+
+@end

src/Network/API/Requests/TSRegisterSignedPrekeyRequest.m

@@ -0,0 +1,36 @@
+//
+//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
+//
+
+#import "TSRegisterSignedPrekeyRequest.h"
+#import "TSConstants.h"
+
+#import <25519/Curve25519.h>
+#import <AxolotlKit/NSData+keyVersionByte.h>
+#import <AxolotlKit/PreKeyRecord.h>
+#import <AxolotlKit/SignedPreKeyStore.h>
+
+@implementation TSRegisterSignedPrekeyRequest
+
+- (id)initWithSignedPreKeyRecord:(SignedPreKeyRecord *)signedRecord
+{
+    self = [super initWithURL:[NSURL URLWithString:textSecureSignedKeysAPI]];
+    self.HTTPMethod = @"PUT";
+
+    NSDictionary *serializedKeyRegistrationParameters = [self dictionaryFromSignedPreKey:signedRecord];
+
+    self.parameters = [serializedKeyRegistrationParameters mutableCopy];
+
+    return self;
+}
+
+- (NSDictionary *)dictionaryFromSignedPreKey:(SignedPreKeyRecord *)preKey
+{
+    return @{
+        @"keyId" : [NSNumber numberWithInt:preKey.Id],
+        @"publicKey" : [[preKey.keyPair.publicKey prependKeyType] base64EncodedStringWithOptions:0],
+        @"signature" : [preKey.signature base64EncodedStringWithOptions:0]
+    };
+}
+
+@end

src/Storage/AxolotlStore/TSStorageManager+PreKeyStore.h

@@ -1,9 +1,5 @@
 //
-//  TSStorageManager+PreKeyStore.h
-//  TextSecureKit
-//
-//  Created by Frederic Jacobs on 06/11/14.
-//  Copyright (c) 2014 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
 //
 
 #import <AxolotlKit/PreKeyStore.h>

src/Storage/AxolotlStore/TSStorageManager+PreKeyStore.m

@@ -1,29 +1,25 @@
 //
-//  TSStorageManager+PreKeyStore.m
-//  TextSecureKit
-//
-//  Created by Frederic Jacobs on 06/11/14.
-//  Copyright (c) 2014 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
 //
 
-#import <AxolotlKit/AxolotlExceptions.h>
 #import "TSStorageManager+PreKeyStore.h"
 #import "TSStorageManager+keyFromIntLong.h"
+#import <AxolotlKit/AxolotlExceptions.h>
+#import <AxolotlKit/SessionBuilder.h>
 
 #define TSStorageManagerPreKeyStoreCollection @"TSStorageManagerPreKeyStoreCollection"
 #define TSNextPrekeyIdKey @"TSStorageInternalSettingsNextPreKeyId"
 #define BATCH_SIZE 100
-#define MAX_VALUE_LASTRESORT 0xFFFFFF
 
 @implementation TSStorageManager (PreKeyStore)
 
 - (PreKeyRecord *)getOrGenerateLastResortKey {
-    if ([self containsPreKey:MAX_VALUE_LASTRESORT]) {
-        return [self loadPreKey:MAX_VALUE_LASTRESORT];
+    if ([self containsPreKey:kPreKeyOfLastResortId]) {
+        return [self loadPreKey:kPreKeyOfLastResortId];
     } else {
         PreKeyRecord *lastResort =
-            [[PreKeyRecord alloc] initWithId:MAX_VALUE_LASTRESORT keyPair:[Curve25519 generateKeyPair]];
-        [self storePreKey:MAX_VALUE_LASTRESORT preKeyRecord:lastResort];
+            [[PreKeyRecord alloc] initWithId:kPreKeyOfLastResortId keyPair:[Curve25519 generateKeyPair]];
+        [self storePreKey:kPreKeyOfLastResortId preKeyRecord:lastResort];
         return lastResort;
     }
 }
@@ -33,6 +29,8 @@
 
     @synchronized(self) {
         int preKeyId = [self nextPreKeyId];
+
+        DDLogInfo(@"%@ building %d new preKeys starting from preKeyId: %d", self.tag, BATCH_SIZE, preKeyId);
         for (int i = 0; i < BATCH_SIZE; i++) {
             ECKeyPair *keyPair   = [Curve25519 generateKeyPair];
             PreKeyRecord *record = [[PreKeyRecord alloc] initWithId:preKeyId keyPair:keyPair];
@@ -81,11 +79,29 @@
 - (int)nextPreKeyId {
     int lastPreKeyId = [self intForKey:TSNextPrekeyIdKey inCollection:TSStorageInternalSettingsCollection];
 
-    while (lastPreKeyId < 1 || (lastPreKeyId > (MAX_VALUE_LASTRESORT - BATCH_SIZE))) {
-        lastPreKeyId = rand();
+    if (lastPreKeyId < 1) {
+        // One-time prekey ids must be > 0 and < kPreKeyOfLastResortId.
+        lastPreKeyId = 1 + arc4random_uniform(kPreKeyOfLastResortId - (BATCH_SIZE + 1));
+    } else if (lastPreKeyId > kPreKeyOfLastResortId - BATCH_SIZE) {
+        // We want to "overflow" to 1 when we reach the "prekey of last resort" id
+        // to avoid biasing towards higher values.
+        lastPreKeyId = 1;
     }
+    OWSCAssert(lastPreKeyId > 0 && lastPreKeyId < kPreKeyOfLastResortId);
 
     return lastPreKeyId;
 }
 
+#pragma mark - Logging
+
++ (NSString *)tag
+{
+    return [NSString stringWithFormat:@"[%@+PreKeyStore]", self.class];
+}
+
+- (NSString *)tag
+{
+    return self.class.tag;
+}
+
 @end

src/Storage/AxolotlStore/TSStorageManager+SignedPreKeyStore.m

@@ -1,13 +1,9 @@
 //
-//  TSStorageManager+SignedPreKeyStore.m
-//  TextSecureKit
+//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
 //
-//  Created by Frederic Jacobs on 06/11/14.
-//  Copyright (c) 2014 Open Whisper Systems. All rights reserved.
-//
-
 
 #import "TSStorageManager+IdentityKeyStore.h"
+#import "TSStorageManager+PreKeyStore.h"
 #import "TSStorageManager+SignedPreKeyStore.h"
 #import "TSStorageManager+keyFromIntLong.h"
 
@@ -19,8 +15,11 @@
 
 - (SignedPreKeyRecord *)generateRandomSignedRecord {
     ECKeyPair *keyPair = [Curve25519 generateKeyPair];
+
+    // Signed prekey ids must be > 0.
+    int preKeyId = 1 + arc4random_uniform(INT32_MAX - 1);
     return [[SignedPreKeyRecord alloc]
-         initWithId:rand()
+         initWithId:preKeyId
             keyPair:keyPair
           signature:[Ed25519 sign:keyPair.publicKey.prependKeyType withKeyPair:[self identityKeyPair]]
         generatedAt:[NSDate date]];