From 912b617a117a9fbc7ffeebb0b1aa068dd036ecf8 Mon Sep 17 00:00:00 2001
From: Frederic Jacobs <github@fredericjacobs.com>
Date: Sun, 7 Jun 2015 21:49:35 +0200
Subject: [PATCH] Support for Mismatched Devices.

---
 Signal/src/textsecure/Contacts/TSRecipient.h  |  3 +-
 Signal/src/textsecure/Contacts/TSRecipient.m  | 18 ++++----
 .../Messages/TSMessagesManager+sendMessages.m | 43 ++++++++++++++++++-
 3 files changed, 52 insertions(+), 12 deletions(-)

diff --git a/Signal/src/textsecure/Contacts/TSRecipient.h b/Signal/src/textsecure/Contacts/TSRecipient.h
index 046315712..0f1dfad38 100644
--- a/Signal/src/textsecure/Contacts/TSRecipient.h
+++ b/Signal/src/textsecure/Contacts/TSRecipient.h
@@ -14,12 +14,11 @@
 
 + (instancetype)recipientWithTextSecureIdentifier:(NSString*)textSecureIdentifier withTransaction:(YapDatabaseReadTransaction*)transaction;
 
-- (NSSet*)devices; //NSNumbers
-
 - (void)addDevices:(NSSet *)set;
 
 - (void)removeDevices:(NSSet *)set;
 
 @property (nonatomic, readonly) NSString *relay;
+@property (nonatomic, retain)   NSMutableOrderedSet *devices;
 
 @end
diff --git a/Signal/src/textsecure/Contacts/TSRecipient.m b/Signal/src/textsecure/Contacts/TSRecipient.m
index 60ccb4388..caf3aa30e 100644
--- a/Signal/src/textsecure/Contacts/TSRecipient.m
+++ b/Signal/src/textsecure/Contacts/TSRecipient.m
@@ -9,12 +9,6 @@
 #import "TSStorageManager+IdentityKeyStore.h"
 #import "TSRecipient.h"
 
-@interface TSRecipient ()
-
-@property (nonatomic, retain) NSMutableSet *devices;
-
-@end
-
 @implementation TSRecipient
 
 + (NSString*)collection{
@@ -25,7 +19,7 @@
     self = [super initWithUniqueId:textSecureIdentifier];
     
     if (self) {
-        _devices     = [NSMutableSet setWithObject:[NSNumber numberWithInt:1]];
+        _devices     = [NSMutableOrderedSet orderedSetWithObject:[NSNumber numberWithInt:1]];
         _relay       = relay;
     }
     
@@ -36,16 +30,24 @@
     return [self fetchObjectWithUniqueID:textSecureIdentifier transaction:transaction];
 }
 
-- (NSSet*)devices{
+- (NSMutableOrderedSet*)devices{
     return [_devices copy];
 }
 
 - (void)addDevices:(NSSet *)set{
+    [self checkDevices];
     [_devices unionSet:set];
 }
 
 - (void)removeDevices:(NSSet *)set{
+    [self checkDevices];
     [_devices minusSet:set];
 }
 
+- (void)checkDevices {
+    if (_devices == nil || ![_devices isKindOfClass:[NSMutableOrderedSet class]]) {
+        _devices = [NSMutableOrderedSet orderedSetWithObject:[NSNumber numberWithInt:1]];
+    }
+}
+
 @end
diff --git a/Signal/src/textsecure/Messages/TSMessagesManager+sendMessages.m b/Signal/src/textsecure/Messages/TSMessagesManager+sendMessages.m
index 54acdbd04..4ca9d3400 100644
--- a/Signal/src/textsecure/Messages/TSMessagesManager+sendMessages.m
+++ b/Signal/src/textsecure/Messages/TSMessagesManager+sendMessages.m
@@ -136,10 +136,28 @@ dispatch_queue_t sendingQueue() {
                         }];
                         break;
                     }
-                    case 409:
+                    case 409:{
                         // Mismatched devices
-                        DDLogError(@"Missing some devices");
+                        DDLogWarn(@"Mismatch Devices.");
+                        
+                        NSError *e;
+                        NSDictionary *serializedResponse = [NSJSONSerialization JSONObjectWithData:responseData
+                                                                                           options:0
+                                                                                             error:&e];
+                        
+                        if (e) {
+                            DDLogError(@"Failed to serialize response of mismatched devices: %@", e.description);
+                        } else {
+                            [self handleMismatchedDevices:serializedResponse
+                                                recipient:recipient];
+                        }
+                        
+                        dispatch_async(sendingQueue(), ^{
+                            [self sendMessage:message toRecipient:recipient inThread:thread withAttemps:remainingAttempts];
+                        });
+                        
                         break;
+                    }
                     case 410:{
                         // staledevices
                         DDLogWarn(@"Stale devices");
@@ -168,6 +186,27 @@ dispatch_queue_t sendingQueue() {
     }
 }
 
+- (void)handleMismatchedDevices:(NSDictionary*)dictionary recipient:(TSRecipient*)recipient {
+    NSArray *extraDevices   = [dictionary objectForKey:@"extraDevices"];
+    NSArray *missingDevices = [dictionary objectForKey:@"missingDevices"];
+    
+    if (extraDevices && [extraDevices count] > 0) {
+        for (NSNumber *extraDeviceId in extraDevices) {
+            [[TSStorageManager sharedManager] deleteSessionForContact:recipient.uniqueId deviceId:[extraDeviceId intValue]];
+        }
+        
+        [recipient removeDevices:[NSSet setWithArray:extraDevices]];
+    }
+    
+    if (missingDevices && [missingDevices count] > 0) {
+        [recipient addDevices:[NSSet setWithArray:missingDevices]];
+    }
+    
+    [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
+        [recipient saveWithTransaction:transaction];
+    }];
+}
+
 - (void)handleMessageSent:(TSOutgoingMessage*)message {
     [self saveMessage:message withState:TSOutgoingMessageStateSent];
 }