e2e: Wait for pending device lists

When we send a megolm message, wait for any existing key download to complete.

lib/crypto/index.js

@@ -57,6 +57,8 @@ function Crypto(baseApis, eventEmitter, sessionStore, userId, deviceId) {
     this._initialSyncCompleted = false;
     // userId -> true
     this._pendingUsersWithNewDevices = {};
+    // userId -> [promise, ...]
+    this._keyDownloadsInProgressByUser = {};
 
     this._olmDevice = new OlmDevice(sessionStore);
 
@@ -285,52 +287,66 @@ function _uploadOneTimeKeys(crypto) {
 Crypto.prototype.downloadKeys = function(userIds, forceDownload) {
     var self = this;
 
-    // map from userid -> deviceid -> DeviceInfo
-    var stored = {};
-    function storeDev(userId, dev) {
-        stored[userId][dev.deviceId] = dev;
-    }
+    // promises we need to wait for while the download happens
+    var promises = [];
 
     // list of userids we need to download keys for
     var downloadUsers = [];
 
+    function perUserCatch(u) {
+        return function(e) {
+            console.warn('Error downloading keys for user ' + u + ':', e);
+        };
+    }
+
     if (forceDownload) {
         downloadUsers = userIds;
     } else {
         for (var i = 0; i < userIds.length; ++i) {
-            var userId = userIds[i];
-            var devices = this.getStoredDevicesForUser(userId);
+            var u = userIds[i];
 
-            if (!devices) {
-                downloadUsers.push(userId);
-            } else {
-                stored[userId] = {};
-                devices.map(storeDev.bind(null, userId));
+            var inprogress = this._keyDownloadsInProgressByUser[u];
+            if (inprogress) {
+                // wait for the download to complete
+                promises.push(q.any(inprogress).catch(perUserCatch(u)));
+            } else if (!this.getStoredDevicesForUser(u)) {
+                downloadUsers.push(u);
             }
         }
     }
 
-    if (downloadUsers.length === 0) {
-        return q(stored);
+    if (downloadUsers.length > 0) {
+        var r = this._doKeyDownloadForUsers(downloadUsers);
+        downloadUsers.map(function(u) {
+            promises.push(r[u].catch(perUserCatch(u)));
+        });
     }
 
-    var r = this._doKeyDownloadForUsers(downloadUsers);
-    var promises = [];
-    downloadUsers.map(function(u) {
-        promises.push(r[u].catch(function(e) {
-            console.warn('Error downloading keys for user ' + u + ':', e);
-        }).then(function() {
-            stored[u] = {};
-            var devices = self.getStoredDevicesForUser(u) || [];
-            devices.map(storeDev.bind(null, u));
-        }));
-    });
-
     return q.all(promises).then(function() {
-        return stored;
+        return self._getDevicesFromStore(userIds);
     });
 };
 
+
+/**
+ * Get the stored device keys for a list of user ids
+ *
+ * @param {string[]} userIds the list of users to list keys for.
+ *
+ * @return {Object} userId->deviceId->{@link module:crypto/deviceinfo|DeviceInfo}.
+ */
+Crypto.prototype._getDevicesFromStore = function(userIds) {
+    var stored = {};
+    userIds.map(function(u) {
+        stored[u] = {};
+        var devices = this.getStoredDevicesForUser(u) || [];
+        devices.map(function(dev) {
+            stored[u][dev.deviceId] = dev;
+        });
+    });
+    return stored;
+};
+
 /**
  * @param {string[]} downloadUsers list of userIds
  *
@@ -345,8 +361,23 @@ Crypto.prototype._doKeyDownloadForUsers = function(downloadUsers) {
     var promiseMap = {};
 
     downloadUsers.map(function(u) {
-        deferMap[u] = q.defer();
-        promiseMap[u] = deferMap[u].promise;
+        var deferred = q.defer();
+        var promise = deferred.promise.finally(function() {
+            var inProgress = self._keyDownloadsInProgressByUser[u];
+            utils.removeElement(inProgress, function(e) { return e === promise; });
+            if (inProgress.length === 0) {
+                // no more downloads for this user; remove the element
+                delete self._keyDownloadsInProgressByUser[u];
+            }
+        });
+
+        if (!self._keyDownloadsInProgressByUser[u]) {
+            self._keyDownloadsInProgressByUser[u] = [];
+        }
+        self._keyDownloadsInProgressByUser[u].push(promise);
+
+        deferMap[u] = deferred;
+        promiseMap[u] = promise;
     });
 
     this._baseApis.downloadKeysForUsers(

spec/integ/megolm.spec.js

@@ -811,4 +811,77 @@ describe("megolm", function() {
             expect(decrypted.content.body).toEqual('test');
         }).nodeify(done);
     });
+
+
+    it('Alice should wait for device list to complete when sending a megolm message',
+    function(done) {
+        var p2pSession;
+        var inboundGroupSession;
+
+        var downloadPromise;
+        var sendPromise;
+
+        aliceTestClient.httpBackend.when(
+            'PUT', '/sendToDevice/m.room.encrypted/'
+        ).respond(200, function(path, content) {
+            var m = content.messages['@bob:xyz'].DEVICE_ID;
+            var ct = m.ciphertext[testSenderKey];
+            var decrypted = JSON.parse(p2pSession.decrypt(ct.type, ct.body));
+
+            expect(decrypted.type).toEqual('m.room_key');
+            inboundGroupSession = new Olm.InboundGroupSession();
+            inboundGroupSession.create(decrypted.content.session_key);
+            return {};
+        });
+
+        aliceTestClient.httpBackend.when(
+            'PUT', '/send/'
+        ).respond(200, function(path, content) {
+            var ct = content.ciphertext;
+            var decrypted = JSON.parse(inboundGroupSession.decrypt(ct));
+
+            console.log('Decrypted received megolm message', decrypted);
+            expect(decrypted.type).toEqual('m.room.message');
+            expect(decrypted.content.body).toEqual('test');
+
+            return {
+                event_id: '$event_id',
+            };
+        });
+
+        return aliceTestClient.start().then(function() {
+            var syncResponse = getSyncResponse(['@bob:xyz']);
+
+            // establish an olm session with alice
+            p2pSession = createOlmSession(testOlmAccount, aliceTestClient);
+
+            var olmEvent = encryptOlmEvent({
+                senderKey: testSenderKey,
+                recipient: aliceTestClient,
+                p2pSession: p2pSession,
+            });
+
+            syncResponse.to_device = { events: [olmEvent] };
+
+            aliceTestClient.httpBackend.when('GET', '/sync').respond(200, syncResponse);
+            return aliceTestClient.httpBackend.flush('/sync', 1);
+        }).then(function() {
+            console.log('Forcing alice to download our device keys');
+
+            // this will block
+            downloadPromise = aliceTestClient.client.downloadKeys(['@bob:xyz']);
+        }).then(function() {
+
+            // so will this.
+            sendPromise = aliceTestClient.client.sendTextMessage(ROOM_ID, 'test');
+        }).then(function() {
+            aliceTestClient.httpBackend.when('POST', '/keys/query').respond(
+                200, getTestKeysQueryResponse('@bob:xyz')
+            );
+
+            return aliceTestClient.httpBackend.flush();
+        }).then(function() {
+            return q.all([downloadPromise, sendPromise]);
+        }).nodeify(done);
+    });
 });