handle partially-shared sessions better

- don't cancel key requests if we can't decrypt everything in the session
- overwrite the session key if we get a better version

spec/unit/crypto.spec.js

@@ -8,6 +8,10 @@ import expect from 'expect';
 import WebStorageSessionStore from '../../lib/store/session/webstorage';
 import MemoryCryptoStore from '../../lib/crypto/store/memory-crypto-store.js';
 import MockStorageApi from '../MockStorageApi';
+import TestClient from '../TestClient';
+import {MatrixEvent} from '../../lib/models/event';
+import Room from '../../lib/models/room';
+import olmlib from '../../lib/crypto/olmlib';
 
 const EventEmitter = require("events").EventEmitter;
 
@@ -119,4 +123,152 @@ describe("Crypto", function() {
             await prom;
         });
     });
+
+    describe('Key requests', function() {
+        let aliceClient;
+        let bobClient;
+
+        beforeEach(async function() {
+            aliceClient = (new TestClient(
+                "@alice:example.com", "alicedevice",
+            )).client;
+            bobClient = (new TestClient(
+                "@bob:example.com", "bobdevice",
+            )).client;
+            await aliceClient.initCrypto();
+            await bobClient.initCrypto();
+        });
+
+        it(
+            "does not cancel keyshare requests if some messages are not decrypted",
+            async function() {
+                function awaitEvent(emitter, event) {
+                    return new Promise((resolve, reject) => {
+                        emitter.once(event, (result) => {
+                            resolve(result);
+                        });
+                    });
+                }
+
+                async function keyshareEventForEvent(event, index) {
+                    const eventContent = event.getWireContent();
+                    const key = await aliceClient._crypto._olmDevice
+                        .getInboundGroupSessionKey(
+                            roomId, eventContent.sender_key, eventContent.session_id,
+                            index,
+                        );
+                    const ksEvent = new MatrixEvent({
+                        type: "m.forwarded_room_key",
+                        sender: "@alice:example.com",
+                        content: {
+                            algorithm: olmlib.MEGOLM_ALGORITHM,
+                            room_id: roomId,
+                            sender_key: eventContent.sender_key,
+                            sender_claimed_ed25519_key: key.sender_claimed_ed25519_key,
+                            session_id: eventContent.session_id,
+                            session_key: key.key,
+                            chain_index: key.chain_index,
+                            forwarding_curve25519_key_chain:
+                            key.forwarding_curve_key_chain,
+                        },
+                    });
+                    // make onRoomKeyEvent think this was an encrypted event
+                    ksEvent._senderCurve25519Key = "akey";
+                    return ksEvent;
+                }
+
+                const encryptionCfg = {
+                    "algorithm": "m.megolm.v1.aes-sha2",
+                };
+                const roomId = "!someroom";
+                const aliceRoom = new Room(roomId, aliceClient, "@alice:example.com", {});
+                const bobRoom = new Room(roomId, bobClient, "@bob:example.com", {});
+                aliceClient.store.storeRoom(aliceRoom);
+                bobClient.store.storeRoom(bobRoom);
+                await aliceClient.setRoomEncryption(roomId, encryptionCfg);
+                await bobClient.setRoomEncryption(roomId, encryptionCfg);
+                const events = [
+                    new MatrixEvent({
+                        type: "m.room.message",
+                        sender: "@alice:example.com",
+                        room_id: roomId,
+                        event_id: "$1",
+                        content: {
+                            msgtype: "m.text",
+                            body: "1",
+                        },
+                    }),
+                    new MatrixEvent({
+                        type: "m.room.message",
+                        sender: "@alice:example.com",
+                        room_id: roomId,
+                        event_id: "$2",
+                        content: {
+                            msgtype: "m.text",
+                            body: "2",
+                        },
+                    }),
+                ];
+                await Promise.all(events.map(async (event) => {
+                    // alice encrypts each event, and then bob tries to decrypt
+                    // them without any keys, so that they'll be in pending
+                    await aliceClient._crypto.encryptEvent(event, aliceRoom);
+                    event._clearEvent = {};
+                    event._senderCurve25519Key = null;
+                    event._claimedEd25519Key = null;
+                    try {
+                        await bobClient._crypto.decryptEvent(event);
+                    } catch (e) {
+                        // we expect this to fail because we don't have the
+                        // decryption keys yet
+                    }
+                }));
+
+                const bobDecryptor = bobClient._crypto._getRoomDecryptor(roomId, olmlib.MEGOLM_ALGORITHM);
+
+                let eventPromise = Promise.all(events.map((ev) => {
+                    return awaitEvent(ev, "Event.decrypted");
+                }));
+
+                // keyshare the session key starting at the second message, so
+                // the first message can't be decrypted yet, but the second one
+                // can
+                let ksEvent = await keyshareEventForEvent(events[1], 1);
+                await bobDecryptor.onRoomKeyEvent(ksEvent);
+                await eventPromise;
+                expect(events[0].getContent().msgtype).toBe("m.bad.encrypted");
+                expect(events[1].getContent().msgtype).toNotBe("m.bad.encrypted");
+
+                const cryptoStore = bobClient._cryptoStore;
+                const eventContent = events[0].getWireContent();
+                const senderKey = eventContent.sender_key;
+                const sessionId = eventContent.session_id;
+                const roomKeyRequestBody = {
+                    algorithm: olmlib.MEGOLM_ALGORITHM,
+                    room_id: roomId,
+                    sender_key: senderKey,
+                    session_id: sessionId,
+                };
+                // the room key request should still be there, since we haven't
+                // decrypted everything
+                expect(await cryptoStore.getOutgoingRoomKeyRequest(roomKeyRequestBody))
+                    .toExist();
+
+                // keyshare the session key starting at the first message, so
+                // that it can now be decrypted
+                eventPromise = awaitEvent(events[0], "Event.decrypted");
+                ksEvent = await keyshareEventForEvent(events[0], 0);
+                await bobDecryptor.onRoomKeyEvent(ksEvent);
+                await eventPromise;
+                expect(events[0].getContent().msgtype).toNotBe("m.bad.encrypted");
+                // the room key request should still be there, since we haven't
+                // decrypted everything
+                expect(await cryptoStore.getOutgoingRoomKeyRequest(roomKeyRequestBody))
+                    .toNotExist();
+
+                aliceClient.stopClient();
+                bobClient.stopClient();
+            },
+        );
+    });
 });

src/crypto/OlmDevice.js

@@ -915,14 +915,6 @@ OlmDevice.prototype.addInboundGroupSession = async function(
             this._getInboundGroupSession(
                 roomId, senderKey, sessionId, txn,
                 (existingSession, existingSessionData) => {
-                    if (existingSession) {
-                        logger.log(
-                            "Update for megolm session " + senderKey + "/" + sessionId,
-                        );
-                        // for now we just ignore updates. TODO: implement something here
-                        return;
-                    }
-
                     // new session.
                     const session = new global.Olm.InboundGroupSession();
                     try {
@@ -938,6 +930,19 @@ OlmDevice.prototype.addInboundGroupSession = async function(
                             );
                         }
 
+                        if (existingSession) {
+                            logger.log(
+                                "Update for megolm session " + senderKey + "/" + sessionId,
+                            );
+                            if (existingSession.first_known_index()
+                                <= session.first_known_index()) {
+                                // existing session has lower index (i.e. can
+                                // decrypt more), so keep it
+                                logger.log("Keeping existing session");
+                                return;
+                            }
+                        }
+
                         const sessionData = {
                             room_id: roomId,
                             session: session.pickle(this._pickleKey),
@@ -945,7 +950,7 @@ OlmDevice.prototype.addInboundGroupSession = async function(
                             forwardingCurve25519KeyChain: forwardingCurve25519KeyChain,
                         };
 
-                        this._cryptoStore.addEndToEndInboundGroupSession(
+                        this._cryptoStore.storeEndToEndInboundGroupSession(
                             senderKey, sessionId, sessionData, txn,
                         );
                     } finally {

src/crypto/algorithms/megolm.js

@@ -938,17 +938,6 @@ MegolmDecryption.prototype.onRoomKeyEvent = function(event) {
         content.session_key, keysClaimed,
         exportFormat,
     ).then(() => {
-        // cancel any outstanding room key requests for this session
-        this._crypto.cancelRoomKeyRequest({
-            algorithm: content.algorithm,
-            room_id: content.room_id,
-            session_id: content.session_id,
-            sender_key: senderKey,
-        });
-
-        // have another go at decrypting events sent with this session.
-        this._retryDecryption(senderKey, sessionId);
-    }).then(() => {
         if (this._crypto.backupInfo) {
             // don't wait for the keys to be backed up for the server
             this._crypto.backupGroupSession(
@@ -961,6 +950,25 @@ MegolmDecryption.prototype.onRoomKeyEvent = function(event) {
                 console.log("Failed to back up group session", e);
             });
         }
+    }).then(() => {
+        // have another go at decrypting events sent with this session.
+        this._retryDecryption(senderKey, sessionId)
+            .then((success) => {
+                // cancel any outstanding room key requests for this session.
+                // Only do this if we managed to decrypt every message in the
+                // session, because if we didn't, we leave the other key
+                // requests in the hopes that someone sends us a key that
+                // includes an earlier index.
+                if (success) {
+                    this._crypto.cancelRoomKeyRequest({
+                        algorithm: content.algorithm,
+                        room_id: content.room_id,
+                        session_id: content.session_id,
+                        sender_key: senderKey,
+                    });
+                }
+            });
+
     }).catch((e) => {
         logger.error(`Error handling m.room_key_event: ${e}`);
     });
@@ -1105,19 +1113,27 @@ MegolmDecryption.prototype.importRoomKey = function(session) {
  * @private
  * @param {String} senderKey
  * @param {String} sessionId
+ *
+ * @return {Boolean} whether all messages were successfully decrypted
  */
-MegolmDecryption.prototype._retryDecryption = function(senderKey, sessionId) {
+MegolmDecryption.prototype._retryDecryption = async function(senderKey, sessionId) {
     const k = senderKey + "|" + sessionId;
     const pending = this._pendingEvents[k];
     if (!pending) {
-        return;
+        return true;
     }
 
     delete this._pendingEvents[k];
 
-    for (const ev of pending) {
+    await Promise.all([...pending].map(async (ev) => {
-        ev.attemptDecryption(this._crypto);
+        try {
+            await ev.attemptDecryption(this._crypto);
+        } catch (e) {
+            // don't die if something goes wrong
         }
+    }));
+
+    return !this._pendingEvents[k];
 };
 
 base.registerAlgorithm(