Crypto: Rework device list tracking logic

Port matrix-org/matrix-js-sdk#425

element-hq/element-ios#955

MatrixSDK/Crypto/Data/MXDeviceList.h

@@ -27,6 +27,16 @@
 
 @class MXCrypto;
 
+// Constants for DeviceList.deviceTrackingStatus
+typedef enum : NSUInteger
+{
+    MXDeviceTrackingStatusNotTracked = 0,
+    MXDeviceTrackingStatusPendingDownload,
+    MXDeviceTrackingStatusDownloadInProgress,
+    MXDeviceTrackingStatusUpToDate
+
+} MXDeviceTrackingStatus;
+
 /**
  `MXDeviceList` manages the list of other users' devices.
  */
@@ -78,8 +88,22 @@
  */
 - (MXDeviceInfo*)deviceWithIdentityKey:(NSString*)senderKey forUser:(NSString*)userId andAlgorithm:(NSString*)algorithm;
 
+/**
+ Flag the given user for device-list tracking, if they are not already.
+ 
+ This will mean that a subsequent call to refreshOutdatedDeviceLists
+ will download the device list for the user, and that subsequent calls to
+ invalidateUserDeviceList will trigger more updates.
+ 
+ @param userId.
+ */
+- (void)startTrackingDeviceList:(NSString*)userId;
+
 /**
  Mark the cached device list for the given user outdated.
+ 
+ If we are not tracking this user's devices, we'll do nothing. Otherwise
+ we flag the user as needing an update.
 
  This doesn't set off an update, so that several users can be batched
  together. Call refreshOutdatedDeviceLists for that.
@@ -89,8 +113,15 @@
 - (void)invalidateUserDeviceList:(NSString*)userId;
 
 /**
- If there is not already a device list query in progress, and we have
- users who have outdated device lists, start a query now.
+ Mark all tracked device lists as outdated.
+ 
+ This will flag each user whose devices we are tracking as in need of an
+ update.
+ */
+- (void)invalidateAllDeviceLists;
+
+/**
+ If we have users who have outdated device lists, start key downloads for them.
  */
 - (void)refreshOutdatedDeviceLists;
 

MatrixSDK/Crypto/Data/MXDeviceList.m

@@ -22,12 +22,17 @@
 
 #import "MXDeviceListOperationsPool.h"
 
+// Helper to transform a NSNumber stored in a NSDictionary to MXDeviceTrackingStatus
+#define MXDeviceTrackingStatusFromNSNumber(aNSNumberObject) ((MXDeviceTrackingStatus)[aNSNumberObject integerValue])
+
+
 @interface MXDeviceList ()
 {
     MXCrypto *crypto;
 
-    // Users with new devices
-    NSMutableSet<NSString*> *pendingUsersWithNewDevices;
+    // Users we are tracking device status for.
+    // userId -> MXDeviceTrackingStatus*
+    NSMutableDictionary<NSString*, NSNumber*> *deviceTrackingStatus;
 
     /**
      The pool which the http request is currenlty being processed.
@@ -57,7 +62,17 @@ - (id)initWithCrypto:(MXCrypto *)theCrypto
     {
         crypto = theCrypto;
 
-        pendingUsersWithNewDevices = [NSMutableSet set];
+        // Retrieve tracking status from the store
+        deviceTrackingStatus = [NSMutableDictionary dictionaryWithDictionary:[crypto.store deviceTrackingStatus]];
+
+        for (NSString *userId in deviceTrackingStatus)
+        {
+            // if a download was in progress when we got shut down, it isn't any more.
+            if (MXDeviceTrackingStatusFromNSNumber(deviceTrackingStatus[userId]) == MXDeviceTrackingStatusDownloadInProgress)
+            {
+                deviceTrackingStatus[userId] = @(MXDeviceTrackingStatusPendingDownload);
+            }
+        }
     }
     return self;
 }
@@ -68,64 +83,52 @@ - (MXHTTPOperation*)downloadKeys:(NSArray<NSString*>*)userIds forceDownload:(BOO
 {
     NSLog(@"[MXDeviceList] downloadKeys(forceDownload: %tu) : %@", forceDownload, userIds);
 
-    BOOL needsRefresh = NO;
-    BOOL waitForCurrentQuery = NO;
+    NSMutableArray *usersToDownload = [NSMutableArray array];
+    BOOL doANewQuery = NO;
 
     for (NSString *userId in userIds)
     {
-        if ([pendingUsersWithNewDevices containsObject:userId])
-        {
-            // we already know this user's devices are outdated
-            needsRefresh = YES;
-        }
-        else if ([currentQueryPool.userIds containsObject:userId])
-        {
-            // already a download in progress - just wait for it.
-            // (even if forceDownload is true)
-            waitForCurrentQuery = true;
-        }
-        else if (forceDownload)
+        MXDeviceTrackingStatus trackingStatus = MXDeviceTrackingStatusFromNSNumber(deviceTrackingStatus[userId]);
+
+        if ([currentQueryPool.userIds containsObject:userId]) // equivalent to (trackingStatus == MXDeviceTrackingStatusDownloadInProgress)
         {
-            NSLog(@"[MXDeviceList] downloadKeys: Invalidating device list for %@ for forceDownload", userId);
-            [self invalidateUserDeviceList:userId];
-            needsRefresh = true;
+            // already a key download in progress/queued for this user; its results
+            // will be good enough for us.
+            [usersToDownload addObject:userId];
         }
-        else if (![self storedDevicesForUser:userId])
+        else if (forceDownload || trackingStatus != MXDeviceTrackingStatusUpToDate)
         {
-            NSLog(@"[MXDeviceList] downloadKeys: Invalidating device list for %@ due to empty cache", userId);
-            [self invalidateUserDeviceList:userId];
-            needsRefresh = true;
+            [usersToDownload addObject:userId];
+            doANewQuery = YES;
         }
     }
 
     MXDeviceListOperation *operation;
 
-    if (needsRefresh)
+    if (usersToDownload.count)
     {
-        NSLog(@"[MXDeviceList] downloadKeys: waiting for next key query");
+        for (NSString *userId in usersToDownload)
+        {
+            deviceTrackingStatus[userId] = @(MXDeviceTrackingStatusDownloadInProgress);
+        }
 
-        operation = [[MXDeviceListOperation alloc] initWithUserIds:userIds success:^(NSArray<NSString *> *succeededUserIds, NSArray<NSString *> *failedUserIds) {
+        operation = [[MXDeviceListOperation alloc] initWithUserIds:usersToDownload success:^(NSArray<NSString *> *succeededUserIds, NSArray<NSString *> *failedUserIds) {
 
-            NSLog(@"[MXDeviceList] downloadKeys: waiting for next key query -> DONE");
-            if (success)
+            NSLog(@"[MXDeviceList] downloadKeys -> DONE");
+
+            for (NSString *userId in succeededUserIds)
             {
-                MXUsersDevicesMap<MXDeviceInfo *> *usersDevicesInfoMap = [self devicesForUsers:userIds];
-                success(usersDevicesInfoMap);
+                MXDeviceTrackingStatus trackingStatus = MXDeviceTrackingStatusFromNSNumber(deviceTrackingStatus[userId]);
+                if (trackingStatus == MXDeviceTrackingStatusDownloadInProgress)
+                {
+                    // we didn't get any new invalidations since this download started:
+                    // this user's device list is now up to date.
+                    deviceTrackingStatus[userId] = @(MXDeviceTrackingStatusUpToDate);
+                }
             }
 
-        } failure:failure];
-
-        [self startOrQueueDeviceQuery:operation];
-
-        return operation;
-    }
-    else if (waitForCurrentQuery)
-    {
-        NSLog(@"[MXDeviceList] downloadKeys: waiting for in-flight query to complete");
+            [self persistDeviceTrackingStatus];
 
-        operation = [[MXDeviceListOperation alloc] initWithUserIds:userIds success:^(NSArray<NSString *> *succeededUserIds, NSArray<NSString *> *failedUserIds) {
-
-            NSLog(@"[MXDeviceList] downloadKeys: waiting for in-flight query to complete -> DONE");
             if (success)
             {
                 MXUsersDevicesMap<MXDeviceInfo *> *usersDevicesInfoMap = [self devicesForUsers:userIds];
@@ -134,9 +137,19 @@ - (MXHTTPOperation*)downloadKeys:(NSArray<NSString*>*)userIds forceDownload:(BOO
 
         } failure:failure];
 
-        [operation addToPool:currentQueryPool];
+        if (doANewQuery)
+        {
+            NSLog(@"[MXDeviceList] downloadKeys: waiting for next key query");
 
-        return operation;
+            [self startOrQueueDeviceQuery:operation];
+        }
+        else
+        {
+
+            NSLog(@"[MXDeviceList] downloadKeys: waiting for in-flight query to complete");
+
+            [operation addToPool:currentQueryPool];
+        }
     }
     else
     {
@@ -182,24 +195,60 @@ - (MXDeviceInfo *)deviceWithIdentityKey:(NSString *)senderKey forUser:(NSString
     return nil;
 }
 
+- (void)startTrackingDeviceList:(NSString*)userId
+{
+    MXDeviceTrackingStatus trackingStatus = MXDeviceTrackingStatusFromNSNumber(deviceTrackingStatus[userId]);
+
+    if (!trackingStatus)
+    {
+        NSLog(@"[MXDeviceList] Now tracking device list for %@", userId);
+        deviceTrackingStatus[userId] = @(MXDeviceTrackingStatusPendingDownload);
+    }
+    // we don't yet persist the tracking status, since there may be a lot
+    // of calls; instead we wait for the forthcoming
+    // refreshOutdatedDeviceLists.
+}
+
 - (void)invalidateUserDeviceList:(NSString *)userId
 {
-    [pendingUsersWithNewDevices addObject:userId];
+    MXDeviceTrackingStatus trackingStatus = MXDeviceTrackingStatusFromNSNumber(deviceTrackingStatus[userId]);
+
+    if (trackingStatus)
+    {
+        NSLog(@"[MXDeviceList] Marking device list outdated for %@", userId);
+        deviceTrackingStatus[userId] = @(MXDeviceTrackingStatusPendingDownload);
+    }
+    // we don't yet persist the tracking status, since there may be a lot
+    // of calls; instead we wait for the forthcoming
+    // refreshOutdatedDeviceLists.
+}
+
+- (void)invalidateAllDeviceLists;
+{
+    for (NSString *userId in deviceTrackingStatus.allKeys)
+    {
+        [self invalidateUserDeviceList:userId];
+    }
 }
 
 - (void)refreshOutdatedDeviceLists
 {
-    NSArray *users = pendingUsersWithNewDevices.allObjects;
-    if (users.count == 0)
+    NSMutableArray *users = [NSMutableArray array];
+    for (NSString *userId in deviceTrackingStatus)
     {
-        // That means we're up-to-date with the lastKnownSyncToken
-        if (_lastKnownSyncToken)
+        MXDeviceTrackingStatus trackingStatus = MXDeviceTrackingStatusFromNSNumber(deviceTrackingStatus[userId]);
+        if (trackingStatus == MXDeviceTrackingStatusPendingDownload)
         {
-            [crypto.store storeDeviceSyncToken:_lastKnownSyncToken];
+            [users addObject:userId];
         }
     }
-    else
+
+    if (users)
     {
+        // we didn't persist the tracking status during
+        // invalidateUserDeviceList, so do it now.
+        [self persistDeviceTrackingStatus];
+
         MXDeviceListOperation *operation = [[MXDeviceListOperation alloc] initWithUserIds:users success:^(NSArray<NSString *> *succeededUserIds, NSArray<NSString *> *failedUserIds) {
 
             NSLog(@"[MXDeviceList] refreshOutdatedDeviceLists: %@", succeededUserIds);
@@ -214,15 +263,24 @@ - (void)refreshOutdatedDeviceLists
 
         } failure:^(NSError *error) {
 
-            NSLog(@"[MXDeviceList] refreshOutdatedDeviceLists: ERROR updating device keys for users %@", pendingUsersWithNewDevices);
-            [pendingUsersWithNewDevices addObjectsFromArray:users];
+            NSLog(@"[MXDeviceList] refreshOutdatedDeviceLists: ERROR updating device keys for users %@", users);
+            for (NSString *userId in users)
+            {
+                deviceTrackingStatus[userId] = @(MXDeviceTrackingStatusPendingDownload);
+            }
 
+            [self persistDeviceTrackingStatus];
         } ];
 
         [self startOrQueueDeviceQuery:operation];
     }
 }
 
+- (void)persistDeviceTrackingStatus
+{
+    [crypto.store storeDeviceTrackingStatus:deviceTrackingStatus];
+}
+
 /**
  Get the stored device keys for a list of user ids.
 
@@ -285,10 +343,6 @@ - (void)startCurrentPoolQuery
     {
         NSString *token = _lastKnownSyncToken;
 
-        // We've kicked off requests to these users: remove their
-        // pending flag for now.
-        [pendingUsersWithNewDevices minusSet:currentQueryPool.userIds];
-
         // Add token
         [currentQueryPool downloadKeys:token complete:^(NSDictionary<NSString *,NSDictionary *> *failedUserIds) {
 

MatrixSDK/Crypto/Data/Store/MXCryptoStore.h

@@ -154,6 +154,20 @@
  */
 - (NSDictionary<NSString*, MXDeviceInfo*>*)devicesForUser:(NSString*)userId;
 
+/**
+ The device tracking status.
+
+ @return A map from user id to MXDeviceTrackingStatus.
+ */
+- (NSDictionary<NSString*, NSNumber*>*)deviceTrackingStatus;
+
+/**
+ Store the device tracking status.
+
+ @param statusMap A map from user id to MXDeviceTrackingStatus.
+ */
+- (void)storeDeviceTrackingStatus:(NSDictionary<NSString*, NSNumber*>*)statusMap;
+
 /**
  Store the crypto algorithm for a room.
 

MatrixSDK/Crypto/Data/Store/MXFileCryptoStore/MXFileCryptoStore.m

@@ -315,6 +315,16 @@ - (void)storeDevicesForUser:(NSString *)userId devices:(NSDictionary<NSString *,
     return usersDevicesInfoMap.map[userId];
 }
 
+- (NSDictionary<NSString*, NSNumber*>*)deviceTrackingStatus
+{
+    return nil;
+}
+
+- (void)storeDeviceTrackingStatus:(NSDictionary<NSString*, NSNumber*>*)statusMap
+{
+
+}
+
 - (void)storeAlgorithmForRoom:(NSString *)roomId algorithm:(NSString *)algorithm
 {
     roomsAlgorithms[roomId] = algorithm;

MatrixSDK/Crypto/Data/Store/MXRealmCryptoStore/MXRealmCryptoStore.m

@@ -126,6 +126,12 @@ @interface MXRealmOlmAccount : RLMObject
  */
 @property (nonatomic) NSString *deviceSyncToken;
 
+/**
+ NSData serialisation of users we are tracking device status for.
+ userId -> MXDeviceTrackingStatus*
+ */
+@property (nonatomic)  NSData *deviceTrackingStatusData;
+
 /**
  Settings for blacklisting unverified devices.
  */
@@ -429,6 +435,21 @@ - (void)storeDevicesForUser:(NSString*)userID devices:(NSDictionary<NSString*, M
     return devicesForUser;
 }
 
+- (NSDictionary<NSString*, NSNumber*>*)deviceTrackingStatus
+{
+    MXRealmOlmAccount *account = self.accountInCurrentThread;
+    return [NSKeyedUnarchiver unarchiveObjectWithData:account.deviceTrackingStatusData];
+}
+
+- (void)storeDeviceTrackingStatus:(NSDictionary<NSString*, NSNumber*>*)statusMap
+{
+    MXRealmOlmAccount *account = self.accountInCurrentThread;
+    [account.realm transactionWithBlock:^{
+
+        account.deviceTrackingStatusData = [NSKeyedArchiver archivedDataWithRootObject:statusMap];
+    }];
+}
+
 - (void)storeAlgorithmForRoom:(NSString*)roomId algorithm:(NSString*)algorithm
 {
     BOOL isNew = NO;