Restart broken Olm sessions

 * Start a new Olm sessions with a device when we get an undecryptable
   message on it.
 * Send a dummy message on that sessions such that the other end knows
   about it.
 * Re-send any outstanding keyshare requests for that device.

Also includes a unit test for megolm that isn't very related but came
out as a result anyway.

Includes #776
Fixes element-hq/element-web#3822

spec/unit/crypto.spec.js

@@ -1,9 +1,13 @@
-
-"use strict";
 import 'source-map-support/register';
 import Crypto from '../../lib/crypto';
 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';
+
+const EventEmitter = require("events").EventEmitter;
+
 const sdk = require("../..");
 
 const Olm = global.Olm;
@@ -20,4 +24,96 @@ describe("Crypto", function() {
     it("Crypto exposes the correct olm library version", function() {
         expect(Crypto.getOlmVersion()[0]).toEqual(3);
     });
+
+
+    describe('Session management', function() {
+        const otkResponse = {
+            one_time_keys: {
+                '@alice:home.server': {
+                    aliceDevice: {
+                        'signed_curve25519:FLIBBLE': {
+                            key: 'YmJiYmJiYmJiYmJiYmJiYmJiYmJiYmJiYmJiYmJiYmI',
+                            signatures: {
+                                '@alice:home.server': {
+                                    'ed25519:aliceDevice': 'totally a valid signature',
+                                },
+                            },
+                        },
+                    },
+                },
+            },
+        };
+        let crypto;
+        let mockBaseApis;
+        let mockRoomList;
+
+        let fakeEmitter;
+
+        beforeEach(async function() {
+            const mockStorage = new MockStorageApi();
+            const sessionStore = new WebStorageSessionStore(mockStorage);
+            const cryptoStore = new MemoryCryptoStore(mockStorage);
+
+            cryptoStore.storeEndToEndDeviceData({
+                devices: {
+                    '@bob:home.server': {
+                        'BOBDEVICE': {
+                            keys: {
+                                'curve25519:BOBDEVICE': 'this is a key',
+                            },
+                        },
+                    },
+                },
+                trackingStatus: {},
+            });
+
+            mockBaseApis = {
+                sendToDevice: expect.createSpy(),
+            };
+            mockRoomList = {};
+
+            fakeEmitter = new EventEmitter();
+
+            crypto = new Crypto(
+                mockBaseApis,
+                sessionStore,
+                "@alice:home.server",
+                "FLIBBLE",
+                sessionStore,
+                cryptoStore,
+                mockRoomList,
+            );
+            crypto.registerEventHandlers(fakeEmitter);
+            await crypto.init();
+        });
+
+        afterEach(async function() {
+            await crypto.stop();
+        });
+
+        it("restarts wedged Olm sessions", async function() {
+            const prom = new Promise((resolve) => {
+                mockBaseApis.claimOneTimeKeys = function() {
+                    resolve();
+                    return otkResponse;
+                };
+            });
+
+            fakeEmitter.emit('toDeviceEvent', {
+                getType: expect.createSpy().andReturn('m.room.message'),
+                getContent: expect.createSpy().andReturn({
+                    msgtype: 'm.bad.encrypted',
+                }),
+                getWireContent: expect.createSpy().andReturn({
+                    algorithm: 'm.olm.v1.curve25519-aes-sha2',
+                    sender_key: 'this is a key',
+                }),
+                getSender: expect.createSpy().andReturn('@bob:home.server'),
+            });
+
+            console.log("waiting");
+            await prom;
+            console.log("done");
+        });
+    });
 });

spec/unit/crypto/algorithms/megolm.spec.js

@@ -18,6 +18,7 @@ import Crypto from '../../../../lib/crypto';
 
 const MatrixEvent = sdk.MatrixEvent;
 const MegolmDecryption = algorithms.DECRYPTION_CLASSES['m.megolm.v1.aes-sha2'];
+const MegolmEncryption = algorithms.ENCRYPTION_CLASSES['m.megolm.v1.aes-sha2'];
 
 const ROOM_ID = '!ROOM:ID';
 
@@ -34,9 +35,11 @@ describe("MegolmDecryption", function() {
     let mockCrypto;
     let mockBaseApis;
 
-    beforeEach(function() {
+    beforeEach(async function() {
         testUtils.beforeEach(this); // eslint-disable-line no-invalid-this
 
+        await Olm.init();
+
         mockCrypto = testUtils.mock(Crypto, 'Crypto');
         mockBaseApis = {};
 
@@ -66,7 +69,6 @@ describe("MegolmDecryption", function() {
     describe('receives some keys:', function() {
         let groupSession;
         beforeEach(async function() {
-            await Olm.init();
             groupSession = new global.Olm.OutboundGroupSession();
             groupSession.create();
 
@@ -263,5 +265,92 @@ describe("MegolmDecryption", function() {
                 // test is successful if no exception is thrown
             });
         });
+
+        it("re-uses sessions for sequential messages", async function() {
+            const mockStorage = new MockStorageApi();
+            const sessionStore = new WebStorageSessionStore(mockStorage);
+            const cryptoStore = new MemoryCryptoStore(mockStorage);
+
+            const olmDevice = new OlmDevice(sessionStore, cryptoStore);
+            olmDevice.verifySignature = expect.createSpy();
+            await olmDevice.init();
+
+            mockBaseApis.claimOneTimeKeys = expect.createSpy().andReturn(Promise.resolve({
+                one_time_keys: {
+                    '@alice:home.server': {
+                        aliceDevice: {
+                            'signed_curve25519:flooble': {
+                                key: 'YmJiYmJiYmJiYmJiYmJiYmJiYmJiYmJiYmJiYmJiYmI',
+                                signatures: {
+                                    '@alice:home.server': {
+                                        'ed25519:aliceDevice': 'totally valid',
+                                    },
+                                },
+                            },
+                        },
+                    },
+                },
+            }));
+            mockBaseApis.sendToDevice = expect.createSpy().andReturn(Promise.resolve());
+
+            mockCrypto.downloadKeys.andReturn(Promise.resolve({
+                '@alice:home.server': {
+                    aliceDevice: {
+                        deviceId: 'aliceDevice',
+                        isBlocked: expect.createSpy().andReturn(false),
+                        isUnverified: expect.createSpy().andReturn(false),
+                        getIdentityKey: expect.createSpy().andReturn(
+                            'YWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWE',
+                        ),
+                        getFingerprint: expect.createSpy().andReturn(''),
+                    },
+                },
+            }));
+
+            const megolmEncryption = new MegolmEncryption({
+                userId: '@user:id',
+                crypto: mockCrypto,
+                olmDevice: olmDevice,
+                baseApis: mockBaseApis,
+                roomId: ROOM_ID,
+                config: {
+                    rotation_period_ms: 9999999999999,
+                },
+            });
+            const mockRoom = {
+                getEncryptionTargetMembers: expect.createSpy().andReturn(
+                    [{userId: "@alice:home.server"}],
+                ),
+                getBlacklistUnverifiedDevices: expect.createSpy().andReturn(false),
+            };
+            const ct1 = await megolmEncryption.encryptMessage(mockRoom, "a.fake.type", {
+                body: "Some text",
+            });
+            expect(mockRoom.getEncryptionTargetMembers).toHaveBeenCalled();
+
+            // this should have claimed a key for alice as it's starting a new session
+            expect(mockBaseApis.claimOneTimeKeys).toHaveBeenCalled(
+                [['@alice:home.server', 'aliceDevice']], 'signed_curve25519',
+            );
+            expect(mockCrypto.downloadKeys).toHaveBeenCalledWith(
+                ['@alice:home.server'], false,
+            );
+            expect(mockBaseApis.sendToDevice).toHaveBeenCalled();
+            expect(mockBaseApis.claimOneTimeKeys).toHaveBeenCalled(
+                [['@alice:home.server', 'aliceDevice']], 'signed_curve25519',
+            );
+
+            mockBaseApis.claimOneTimeKeys.reset();
+
+            const ct2 = await megolmEncryption.encryptMessage(mockRoom, "a.fake.type", {
+                body: "Some more text",
+            });
+
+            // this should *not* have claimed a key as it should be using the same session
+            expect(mockBaseApis.claimOneTimeKeys).toNotHaveBeenCalled();
+
+            // likewise they should show the same session ID
+            expect(ct2.session_id).toEqual(ct1.session_id);
+        });
     });
 });

src/crypto/OlmDevice.js

@@ -461,6 +461,12 @@ OlmDevice.prototype.createOutboundSession = async function(
                     session.create_outbound(account, theirIdentityKey, theirOneTimeKey);
                     newSessionId = session.session_id();
                     this._storeAccount(txn, account);
+                    // Pretend we've received a message at this point, otherwise
+                    // if we try to send a message to the device, it won't use
+                    // this session (storing the creation time separately would
+                    // make the pickle longer and would not be useful otherwise).
+                    session.set_last_received_message_ts(Date.now());
+
                     this._saveSession(theirIdentityKey, session, txn);
                 } finally {
                     session.free();
@@ -725,7 +731,7 @@ OlmDevice.prototype._saveOutboundGroupSession = function(session) {
  */
 OlmDevice.prototype._getOutboundGroupSession = function(sessionId, func) {
     const pickled = this._outboundGroupSessionStore[sessionId];
-    if (pickled === null) {
+    if (pickled === undefined) {
         throw new Error("Unknown outbound group session " + sessionId);
     }
 
@@ -1059,16 +1065,21 @@ OlmDevice.prototype.hasInboundSessionKeys = async function(roomId, senderKey, se
  * @param {string} roomId    room in which the message was received
  * @param {string} senderKey base64-encoded curve25519 key of the sender
  * @param {string} sessionId session identifier
+ * @param {integer} chainIndex The chain index at which to export the session.
+ *     If omitted, export at the first index we know about.
  *
  * @returns {Promise<{chain_index: number, key: string,
  *        forwarding_curve25519_key_chain: Array<string>,
  *        sender_claimed_ed25519_key: string
  *    }>}
  *    details of the session key. The key is a base64-encoded megolm key in
  *    export format.
+ *
+ * @throws Error If the given chain index could not be obtained from the known
+ *     index (ie. the given chain index is before the first we have).
  */
 OlmDevice.prototype.getInboundGroupSessionKey = async function(
-    roomId, senderKey, sessionId,
+    roomId, senderKey, sessionId, chainIndex,
 ) {
     let result;
     await this._cryptoStore.doTxn(
@@ -1079,14 +1090,19 @@ OlmDevice.prototype.getInboundGroupSessionKey = async function(
                         result = null;
                         return;
                     }
-                    const messageIndex = session.first_known_index();
+
+                    if (chainIndex === undefined) {
+                        chainIndex = session.first_known_index();
+                    }
+
+                    const exportedSession = session.export_session(chainIndex);
 
                     const claimedKeys = sessionData.keysClaimed || {};
                     const senderEd25519Key = claimedKeys.ed25519 || null;
 
                     result = {
-                        "chain_index": messageIndex,
-                        "key": session.export_session(messageIndex),
+                        "chain_index": chainIndex,
+                        "key": exportedSession,
                         "forwarding_curve25519_key_chain":
                             sessionData.forwardingCurve25519KeyChain || [],
                         "sender_claimed_ed25519_key": senderEd25519Key,

src/crypto/OutgoingRoomKeyRequestManager.js

@@ -244,6 +244,21 @@ export default class OutgoingRoomKeyRequestManager {
         });
     }
 
+    /**
+     * Look for room key requests by target device and state
+     *
+     * @param {string} userId Target user ID
+     * @param {string} deviceId Target device ID
+     *
+     * @return {Promise} resolves to a list of all the
+     *    {@link module:crypto/store/base~OutgoingRoomKeyRequest}
+     */
+    getOutgoingSentRoomKeyRequest(userId, deviceId) {
+        return this._cryptoStore.getOutgoingRoomKeyRequestsByTarget(
+            userId, deviceId, [ROOM_KEY_REQUEST_STATES.SENT],
+        );
+    }
+
     // start the background timer to send queued requests, if the timer isn't
     // already running
     _startTimer() {