Check devices to share keys with on each send

Instead of trying to maintain a list of devices we need to share with, just
check all the devices for all the users on each send.

This should fix https://github.com/vector-im/vector-web/issues/2568, and
generally mean we're less likely to get out of sync.

lib/crypto/algorithms/megolm.js

@@ -38,12 +38,17 @@ var base = require("./base");
  * @property {Number} creationTime when the session was created (ms since the epoch)
  * @property {module:client.Promise?} sharePromise  If a share operation is in progress,
  *    a promise which resolves when it is complete.
+ *
+ * @property {object} sharedWithDevices
+ *    devices with which we have shared the session key
+ *        userId -> {deviceId -> msgindex}
  */
 function OutboundSessionInfo(sessionId) {
     this.sessionId = sessionId;
     this.useCount = 0;
     this.creationTime = new Date().getTime();
     this.sharePromise = null;
+    this.sharedWithDevices = {};
 }
 
 
@@ -90,11 +95,6 @@ function MegolmEncryption(params) {
     // case _outboundSession.sharePromise will be non-null.)
     this._outboundSession = null;
 
-    // devices which have joined since we last sent a message.
-    // userId -> {deviceId -> true}, or
-    // userId -> true
-    this._devicesPendingKeyShare = {};
-
     // default rotation periods
     this._sessionRotationPeriodMsgs = 100;
     this._sessionRotationPeriodMs = 7 * 24 * 3600 * 1000;
@@ -134,32 +134,55 @@ MegolmEncryption.prototype._ensureOutboundSession = function(room) {
         return session.sharePromise;
     }
 
-    // no share in progress: check for new devices
+    // no share in progress: check if we need to share with any devices
-    var shareMap = this._devicesPendingKeyShare;
+    var prom = this._getDevicesInRoom(room).then(function(devicesInRoom) {
-    this._devicesPendingKeyShare = {};
+        var shareMap = {};
 
-    // check each user is (still) a member of the room
+        for (var userId in devicesInRoom) {
-    for (var userId in shareMap) {
+            if (!devicesInRoom.hasOwnProperty(userId)) {
-        if (!shareMap.hasOwnProperty(userId)) {
+                continue;
-            continue;
+            }
+
+            var userDevices = devicesInRoom[userId];
+
+            for (var deviceId in userDevices) {
+                if (!userDevices.hasOwnProperty(deviceId)) {
+                    continue;
+                }
+
+                var deviceInfo = userDevices[deviceId];
+
+                if (deviceInfo.isBlocked()) {
+                    continue;
+                }
+
+                var key = deviceInfo.getIdentityKey();
+                if (key == self._olmDevice.deviceCurve25519Key) {
+                    // don't bother sending to ourself
+                    continue;
+                }
+
+                if (
+                    !session.sharedWithDevices[userId] ||
+                        session.sharedWithDevices[userId][deviceId] === undefined
+                ) {
+                    shareMap[userId] = shareMap[userId] || [];
+                    shareMap[userId].push(deviceInfo);
+                }
+            }
         }
 
-        // XXX what about rooms where invitees can see the content?
+        return self._shareKeyWithDevices(
-        var member = room.getMember(userId);
+            session, shareMap
-        if (member.membership !== "join") {
+        );
-            delete shareMap[userId];
+    }).finally(function() {
-        }
-    }
-
-    session.sharePromise = this._shareKeyWithDevices(
-        session.sessionId, shareMap
-    ).finally(function() {
         session.sharePromise = null;
     }).then(function() {
         return session;
     });
 
-    return session.sharePromise;
+    session.sharePromise = prom;
+    return prom;
 };
 
 /**
@@ -178,95 +201,53 @@ MegolmEncryption.prototype._prepareNewSession = function(room) {
         key.key, {ed25519: this._olmDevice.deviceEd25519Key}
     );
 
-    // we're going to share the key with all current members of the room,
+    return new OutboundSessionInfo(session_id);
-    // so we can reset this.
-    this._devicesPendingKeyShare = {};
-
-    var session = new OutboundSessionInfo(session_id);
-
-    var roomMembers = utils.map(room.getJoinedMembers(), function(u) {
-        return u.userId;
-    });
-
-    var shareMap = {};
-    for (var i = 0; i < roomMembers.length; i++) {
-        var userId = roomMembers[i];
-        shareMap[userId] = true;
-    }
-
-    var self = this;
-
-    // TODO: we need to give the user a chance to block any devices or users
-    // before we send them the keys; it's too late to download them here.
-    session.sharePromise = this._crypto.downloadKeys(
-        roomMembers, false
-    ).then(function(res) {
-        return self._shareKeyWithDevices(session_id, shareMap);
-    }).then(function() {
-        return session;
-    }).finally(function() {
-        session.sharePromise = null;
-    });
-
-    return session;
 };
 
 /**
  * @private
  *
- * @param {string} session_id
+ * @param {module:crypto/algorithms/megolm.OutboundSessionInfo} session
  *
- * @param {Object<string, Object<string, boolean>|boolean>} shareMap
+ * @param {object<string, module:crypto/deviceinfo[]>} devicesByUser
- *    Map from userid to either: true (meaning this is a new user in the room,
+ *    map from userid to list of devices
- *    so all of his devices need the keys); or a map from deviceid to true
- *    (meaning this user has one or more new devices, which need the keys).
  *
  * @return {module:client.Promise} Promise which resolves once the key sharing
  *     message has been sent.
  */
-MegolmEncryption.prototype._shareKeyWithDevices = function(session_id, shareMap) {
+MegolmEncryption.prototype._shareKeyWithDevices = function(session, devicesByUser) {
     var self = this;
 
-    var key = this._olmDevice.getOutboundGroupSessionKey(session_id);
+    var key = this._olmDevice.getOutboundGroupSessionKey(session.sessionId);
     var payload = {
         type: "m.room_key",
         content: {
             algorithm: olmlib.MEGOLM_ALGORITHM,
             room_id: this._roomId,
-            session_id: session_id,
+            session_id: session.sessionId,
             session_key: key.key,
             chain_index: key.chain_index,
         }
     };
 
-    // we downloaded the user's device list when they joined the room, or when
+    var contentMap = {};
-    // the new device announced itself, so there is no need to do so now.
 
-    return self._crypto.ensureOlmSessionsForUsers(
+    return olmlib.ensureOlmSessionsForDevices(
-        utils.keys(shareMap)
+        this._olmDevice, this._baseApis, devicesByUser
     ).then(function(devicemap) {
-        var contentMap = {};
         var haveTargets = false;
 
-        for (var userId in devicemap) {
+        for (var userId in devicesByUser) {
-            if (!devicemap.hasOwnProperty(userId)) {
+            if (!devicesByUser.hasOwnProperty(userId)) {
                 continue;
             }
 
-            var devicesToShareWith = shareMap[userId];
+            var devicesToShareWith = devicesByUser[userId];
             var sessionResults = devicemap[userId];
 
-            for (var deviceId in sessionResults) {
+            for (var i = 0; i < devicesToShareWith.length; i++) {
-                if (!sessionResults.hasOwnProperty(deviceId)) {
+                var deviceInfo = devicesToShareWith[i];
-                    continue;
+                var deviceId = deviceInfo.deviceId;
-                }
-
-                if (devicesToShareWith === true) {
-                    // all devices
-                } else if (!devicesToShareWith[deviceId]) {
-                    // not a new device
-                    continue;
-                }
 
                 var sessionResult = sessionResults[deviceId];
                 if (!sessionResult.sessionId) {
@@ -288,8 +269,6 @@ MegolmEncryption.prototype._shareKeyWithDevices = function(session_id, shareMap)
                     "sharing keys with device " + userId + ":" + deviceId
                 );
 
-                var deviceInfo = sessionResult.device;
-
                 var encryptedContent = {
                     algorithm: olmlib.OLM_ALGORITHM,
                     sender_key: self._olmDevice.deviceCurve25519Key,
@@ -321,6 +300,27 @@ MegolmEncryption.prototype._shareKeyWithDevices = function(session_id, shareMap)
 
         // TODO: retries
         return self._baseApis.sendToDevice("m.room.encrypted", contentMap);
+    }).then(function() {
+        // Add the devices we have shared with to session.sharedWithDevices.
+        //
+        // we deliberately iterate over devicesByUser (ie, the devices we
+        // attempted to share with) rather than the contentMap (those we did
+        // share with), because we don't want to try to claim a one-time-key
+        // for dead devices on every message.
+        for (var userId in devicesByUser) {
+            if (!devicesByUser.hasOwnProperty(userId)) {
+                continue;
+            }
+            if (!session.sharedWithDevices[userId]) {
+                session.sharedWithDevices[userId] = {};
+            }
+            var devicesToShareWith = devicesByUser[userId];
+            for (var i = 0; i < devicesToShareWith.length; i++) {
+                var deviceInfo = devicesToShareWith[i];
+                session.sharedWithDevices[userId][deviceInfo.deviceId] =
+                    key.chain_index;
+            }
+        }
     });
 };
 
@@ -369,20 +369,9 @@ MegolmEncryption.prototype.encryptMessage = function(room, eventType, content) {
  * @param {string=} oldMembership  previous membership
  */
 MegolmEncryption.prototype.onRoomMembership = function(event, member, oldMembership) {
-    // if we haven't yet made a session, there's nothing to do here.
-    if (!this._outboundSession) {
-        return;
-    }
-
     var newMembership = member.membership;
 
-    if (newMembership === 'join') {
+    if (newMembership === 'join' || newMembership === 'invite') {
-        this._onNewRoomMember(member.userId);
-        return;
-    }
-
-    if (newMembership === 'invite' && oldMembership !== 'join') {
-        // we don't (yet) share keys with invited members, so nothing to do yet
         return;
     }
 
@@ -396,44 +385,26 @@ MegolmEncryption.prototype.onRoomMembership = function(event, member, oldMembers
 };
 
 /**
- * handle a new user joining a room
+ * Get the list of devices for all users in the room
  *
- * @param {string} userId   new member
+ * @param {module:models/room} room
- */
-MegolmEncryption.prototype._onNewRoomMember = function(userId) {
-    // make sure we have a list of this user's devices. We are happy to use a
-    // cached version here: we assume that if we already have a list of the
-    // user's devices, then we already share an e2e room with them, which means
-    // that they will have announced any new devices via an m.new_device.
-    this._crypto.downloadKeys([userId], false).done();
-
-    // also flag this user up for needing a keyshare.
-    this._devicesPendingKeyShare[userId] = true;
-};
-
-
-/**
- * @inheritdoc
  *
- * @param {string} userId    owner of the device
+ * @return {module:client.Promise} Promise which resolves to a map
- * @param {string} deviceId  deviceId of the device
+ *     from userId to deviceId to deviceInfo
  */
-MegolmEncryption.prototype.onNewDevice = function(userId, deviceId) {
+MegolmEncryption.prototype._getDevicesInRoom = function(room) {
-    var d = this._devicesPendingKeyShare[userId];
+    // XXX what about rooms where invitees can see the content?
+    var roomMembers = utils.map(room.getJoinedMembers(), function(u) {
+        return u.userId;
+    });
 
-    if (d === true) {
+    // We are happy to use a cached version here: we assume that if we already
-        // we already want to share keys with all devices for this user
+    // have a list of the user's devices, then we already share an e2e room
-        return;
+    // with them, which means that they will have announced any new devices via
-    }
+    // an m.new_device.
-
+    return this._crypto.downloadKeys(roomMembers, false);
-    if (!d) {
-        this._devicesPendingKeyShare[userId] = d = {};
-    }
-
-    d[deviceId] = true;
 };
 
-
 /**
  * Megolm decryption implementation
  *

spec/integ/megolm.spec.js

@@ -541,4 +541,77 @@ describe("megolm", function() {
         }).nodeify(done);
     });
 
+
+    it("We shouldn't attempt to send to blocked devices", function(done) {
+        // establish an olm session with alice
+        var p2pSession = createOlmSession(testOlmAccount, aliceTestClient);
+
+        var olmEvent = encryptOlmEvent({
+            senderKey: testSenderKey,
+            recipient: aliceTestClient,
+            p2pSession: p2pSession,
+        });
+
+        var syncResponse = {
+            next_batch: 1,
+            to_device: {
+                events: [olmEvent],
+            },
+            rooms: {
+                join: {},
+            },
+        };
+
+        syncResponse.rooms.join[ROOM_ID] = {
+            state: {
+                events: [
+                    test_utils.mkEvent({
+                        type: 'm.room.encryption',
+                        skey: '',
+                        content: {
+                            algorithm: 'm.megolm.v1.aes-sha2',
+                        },
+                    }),
+                    test_utils.mkMembership({
+                        mship: 'join',
+                        sender: '@bob:xyz',
+                    }),
+                ],
+            },
+        };
+        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');
+
+            aliceTestClient.httpBackend.when('POST', '/keys/query').respond(200, {
+                device_keys: {
+                    '@bob:xyz': {
+                        'DEVICE_ID': testDeviceKeys,
+                    },
+                }
+            });
+
+            return q.all([
+                aliceTestClient.client.downloadKeys(['@bob:xyz']),
+                aliceTestClient.httpBackend.flush('/keys/query', 1),
+            ]);
+        }).then(function() {
+            console.log('Telling alice to block our device');
+            aliceTestClient.client.setDeviceBlocked('@bob:xyz', 'DEVICE_ID');
+
+            console.log('Telling alice to send a megolm message');
+            aliceTestClient.httpBackend.when(
+                'PUT', '/send/'
+            ).respond(200, {
+                event_id: '$event_id',
+            });
+
+            return q.all([
+                aliceTestClient.client.sendTextMessage(ROOM_ID, 'test'),
+                aliceTestClient.httpBackend.flush(),
+            ]);
+        }).nodeify(done);
+    });
+
 });