Update recipient devices on successful decrypt to avoid wasting a valid session

created by sender.

Make device set immutable.

SignalServiceKit/src/Contacts/SignalRecipient.h

@@ -1,5 +1,5 @@
 //
-//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
 //
 
 #import "TSYapDatabaseObject.h"
@@ -12,16 +12,21 @@ NS_ASSUME_NONNULL_BEGIN
                                        relay:(nullable NSString *)relay;
 
 + (instancetype)selfRecipient;
+
++ (void)ensureRecipientExistsWithRecipientId:(NSString *)recipientId
+                                    deviceId:(UInt32)deviceId
+                                       relay:(NSString *)relay
+                                 transaction:(YapDatabaseReadWriteTransaction *)transaction;
+
 + (nullable instancetype)recipientWithTextSecureIdentifier:(NSString *)textSecureIdentifier;
 + (nullable instancetype)recipientWithTextSecureIdentifier:(NSString *)textSecureIdentifier
                                            withTransaction:(YapDatabaseReadTransaction *)transaction;
 
+@property (readonly) NSOrderedSet *devices;
 - (void)addDevices:(NSSet *)set;
-
 - (void)removeDevices:(NSSet *)set;
 
 @property (nonatomic, nullable) NSString *relay;
-@property (nonatomic) NSMutableOrderedSet *devices;
 
 - (BOOL)supportsVoice;
 // This property indicates support for both WebRTC audio and video calls.

SignalServiceKit/src/Contacts/SignalRecipient.m

@@ -9,12 +9,47 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+@interface SignalRecipient ()
+
+@property NSOrderedSet *devices;
+
+@end
+
 @implementation SignalRecipient
 
 + (NSString *)collection {
     return @"SignalRecipient";
 }
 
++ (void)ensureRecipientExistsWithRecipientId:(NSString *)recipientId
+                                    deviceId:(UInt32)deviceId
+                                       relay:(NSString *)relay
+                                 transaction:(YapDatabaseReadWriteTransaction *)transaction
+{
+    SignalRecipient *_Nullable existingRecipient =
+        [self recipientWithTextSecureIdentifier:recipientId withTransaction:transaction];
+    if (!existingRecipient) {
+        DDLogDebug(
+            @"%@ in %s creating recipient with deviceId: %u", self.logTag, __PRETTY_FUNCTION__, (unsigned int)deviceId);
+
+        SignalRecipient *newRecipient = [[self alloc] initWithTextSecureIdentifier:recipientId relay:relay];
+        [newRecipient addDevices:[NSSet setWithObject:@(deviceId)]];
+        [newRecipient saveWithTransaction:transaction];
+
+        return;
+    }
+
+    if (![existingRecipient.devices containsObject:@(deviceId)]) {
+        DDLogDebug(@"%@ in %s adding device %u to existing recipient.",
+            self.logTag,
+            __PRETTY_FUNCTION__,
+            (unsigned int)deviceId);
+
+        [existingRecipient addDevices:[NSSet setWithObject:@(deviceId)]];
+        [existingRecipient saveWithTransaction:transaction];
+    }
+}
+
 - (instancetype)initWithTextSecureIdentifier:(NSString *)textSecureIdentifier
                                        relay:(nullable NSString *)relay
 {
@@ -31,13 +66,13 @@ NS_ASSUME_NONNULL_BEGIN
         // sync message to linked devices.  We shouldn't have any linked devices
         // yet when we create the "self" SignalRecipient, and we don't need to
         // send sync messages to the primary - we ARE the primary.
-        _devices = [NSMutableOrderedSet new];
+        _devices = [NSOrderedSet new];
     } else {
         // Default to sending to just primary device.
         //
         // OWSMessageSender will correct this if it is wrong the next time
         // we send a message to this recipient.
-        _devices = [NSMutableOrderedSet orderedSetWithObject:@(1)];
+        _devices = [NSOrderedSet orderedSetWithObject:@(1)];
     }
 
     _relay = [relay isEqualToString:@""] ? nil : relay;
@@ -45,6 +80,24 @@ NS_ASSUME_NONNULL_BEGIN
     return self;
 }
 
+- (nullable instancetype)initWithCoder:(NSCoder *)coder
+{
+    self = [super initWithCoder:coder];
+    if (!self) {
+        return self;
+    }
+
+    if (_devices == nil) {
+        _devices = [NSOrderedSet new];
+    }
+
+    if ([self.uniqueId isEqual:[TSAccountManager localNumber]] && [self.devices containsObject:@(1)]) {
+        OWSFail(@"%@ in %s self as recipient device", self.logTag, __PRETTY_FUNCTION__);
+    }
+
+    return self;
+}
+
 + (nullable instancetype)recipientWithTextSecureIdentifier:(NSString *)textSecureIdentifier
                                            withTransaction:(YapDatabaseReadTransaction *)transaction
 {
@@ -70,24 +123,25 @@ NS_ASSUME_NONNULL_BEGIN
     return myself;
 }
 
-- (NSMutableOrderedSet *)devices {
-    return [_devices copy];
-}
+- (void)addDevices:(NSSet *)set
+{
+    if ([self.uniqueId isEqual:[TSAccountManager localNumber]] && [set containsObject:@(1)]) {
+        OWSFail(@"%@ in %s adding self as recipient device", self.logTag, __PRETTY_FUNCTION__);
+        return;
+    }
 
-- (void)addDevices:(NSSet *)set {
-    [self checkDevices];
-    [_devices unionSet:set];
-}
+    NSMutableOrderedSet *updatedDevices = [self.devices mutableCopy];
+    [updatedDevices unionSet:set];
 
-- (void)removeDevices:(NSSet *)set {
-    [self checkDevices];
-    [_devices minusSet:set];
+    self.devices = [updatedDevices copy];
 }
 
-- (void)checkDevices {
-    if (_devices == nil || ![_devices isKindOfClass:[NSMutableOrderedSet class]]) {
-        _devices = [NSMutableOrderedSet orderedSetWithObject:[NSNumber numberWithInt:1]];
-    }
+- (void)removeDevices:(NSSet *)set
+{
+    NSMutableOrderedSet *updatedDevices = [self.devices mutableCopy];
+    [updatedDevices minusSet:set];
+
+    self.devices = [updatedDevices copy];
 }
 
 - (BOOL)supportsVoice

SignalServiceKit/src/Messages/OWSMessageDecrypter.m

@@ -14,6 +14,7 @@
 #import "OWSPrimaryStorage+SignedPreKeyStore.h"
 #import "OWSPrimaryStorage.h"
 #import "OWSSignalServiceProtos.pb.h"
+#import "SignalRecipient.h"
 #import "TSAccountManager.h"
 #import "TSContactThread.h"
 #import "TSErrorMessage.h"
@@ -89,11 +90,11 @@ NS_ASSUME_NONNULL_BEGIN
 #pragma mark - Decryption
 
 - (void)decryptEnvelope:(OWSSignalServiceProtosEnvelope *)envelope
-           successBlock:(DecryptSuccessBlock)successBlock
+           successBlock:(DecryptSuccessBlock)successBlockParameter
            failureBlock:(DecryptFailureBlock)failureBlockParameter
 {
     OWSAssert(envelope);
-    OWSAssert(successBlock);
+    OWSAssert(successBlockParameter);
     OWSAssert(failureBlockParameter);
     OWSAssert([TSAccountManager isRegistered]);
 
@@ -107,6 +108,16 @@ NS_ASSUME_NONNULL_BEGIN
         });
     };
 
+    DecryptSuccessBlock successBlock
+        = ^(NSData *_Nullable plaintextData, YapDatabaseReadWriteTransaction *transaction) {
+              [SignalRecipient ensureRecipientExistsWithRecipientId:envelope.source
+                                                           deviceId:envelope.sourceDevice
+                                                              relay:envelope.relay
+                                                        transaction:transaction];
+
+              successBlockParameter(plaintextData, transaction);
+          };
+
     @try {
         DDLogInfo(@"%@ decrypting envelope: %@", self.logTag, [self descriptionForEnvelope:envelope]);