Various changes to src/crypto files for correctness (#2137)

* make various changes for correctness

* apply some review feedback

* Address some review feedback

* add some more correctness

* refactor ensureOutboundSession to fit types better

* change variable naming slightly to prevent confusion

* some wording around exception-catching

* Tidy test

* Simplify

* Add tests

* Add more test coverage

* Apply suggestions from code review

Co-authored-by: Travis Ralston <travpc@gmail.com>

* Update crypto.spec.js

* Update spec/unit/crypto.spec.js

Co-authored-by: Faye Duxovni <duxovni@duxovni.org>

Co-authored-by: Michael Telatynski <7t3chguy@gmail.com>
Co-authored-by: Travis Ralston <travpc@gmail.com>
Co-authored-by: Faye Duxovni <duxovni@duxovni.org>

spec/unit/crypto.spec.js

@@ -17,6 +17,43 @@ import { logger } from '../../src/logger';
 
 const Olm = global.Olm;
 
+function awaitEvent(emitter, event) {
+    return new Promise((resolve, reject) => {
+        emitter.once(event, (result) => {
+            resolve(result);
+        });
+    });
+}
+
+async function keyshareEventForEvent(client, event, index) {
+    const roomId = event.getRoomId();
+    const eventContent = event.getWireContent();
+    const key = await client.crypto.olmDevice.getInboundGroupSessionKey(
+        roomId,
+        eventContent.sender_key,
+        eventContent.session_id,
+        index,
+    );
+    const ksEvent = new MatrixEvent({
+        type: "m.forwarded_room_key",
+        sender: client.getUserId(),
+        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;
+}
+
 describe("Crypto", function() {
     if (!CRYPTO_ENABLED) {
         return;
@@ -203,136 +240,141 @@ describe("Crypto", function() {
             bobClient.stopClient();
         });
 
-        it(
+        it("does not cancel keyshare requests if some messages are not decrypted", async function() {
-            "does not cancel keyshare requests if some messages are not decrypted",
+            const encryptionCfg = {
-            async function() {
+                "algorithm": "m.megolm.v1.aes-sha2",
-                function awaitEvent(emitter, event) {
+            };
-                    return new Promise((resolve, reject) => {
+            const roomId = "!someroom";
-                        emitter.once(event, (result) => {
+            const aliceRoom = new Room(roomId, aliceClient, "@alice:example.com", {});
-                            resolve(result);
+            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);
-                async function keyshareEventForEvent(event, index) {
+            const events = [
-                    const eventContent = event.getWireContent();
+                new MatrixEvent({
-                    const key = await aliceClient.crypto.olmDevice
+                    type: "m.room.message",
-                        .getInboundGroupSessionKey(
+                    sender: "@alice:example.com",
-                            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 = undefined;
-                    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).not.toBe("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,
+                    event_id: "$1",
-                    session_id: sessionId,
+                    content: {
-                };
+                        msgtype: "m.text",
-                // the room key request should still be there, since we haven't
+                        body: "1",
-                // decrypted everything
+                    },
-                expect(await cryptoStore.getOutgoingRoomKeyRequest(roomKeyRequestBody))
+                }),
-                    .toBeDefined();
+                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 = undefined;
+                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
+                }
+            }));
 
-                // keyshare the session key starting at the first message, so
+            const bobDecryptor = bobClient.crypto.getRoomDecryptor(
-                // that it can now be decrypted
+                roomId, olmlib.MEGOLM_ALGORITHM,
-                eventPromise = awaitEvent(events[0], "Event.decrypted");
+            );
-                ksEvent = await keyshareEventForEvent(events[0], 0);
+
-                await bobDecryptor.onRoomKeyEvent(ksEvent);
+            let eventPromise = Promise.all(events.map((ev) => {
-                await eventPromise;
+                return awaitEvent(ev, "Event.decrypted");
-                expect(events[0].getContent().msgtype).not.toBe("m.bad.encrypted");
+            }));
-                await sleep(1);
+
-                // the room key request should be gone since we've now decrypted everything
+            // keyshare the session key starting at the second message, so
-                expect(await cryptoStore.getOutgoingRoomKeyRequest(roomKeyRequestBody))
+            // the first message can't be decrypted yet, but the second one
-                    .toBeFalsy();
+            // can
-            },
+            let ksEvent = await keyshareEventForEvent(aliceClient, events[1], 1);
-        );
+            await bobDecryptor.onRoomKeyEvent(ksEvent);
+            await eventPromise;
+            expect(events[0].getContent().msgtype).toBe("m.bad.encrypted");
+            expect(events[1].getContent().msgtype).not.toBe("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)).toBeDefined();
+
+            // 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(aliceClient, events[0], 0);
+            await bobDecryptor.onRoomKeyEvent(ksEvent);
+            await eventPromise;
+            expect(events[0].getContent().msgtype).not.toBe("m.bad.encrypted");
+            await sleep(1);
+            // the room key request should be gone since we've now decrypted everything
+            expect(await cryptoStore.getOutgoingRoomKeyRequest(roomKeyRequestBody)).toBeFalsy();
+        });
+
+        it("should error if a forwarded room key lacks a content.sender_key", async function() {
+            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 event = new MatrixEvent({
+                type: "m.room.message",
+                sender: "@alice:example.com",
+                room_id: roomId,
+                event_id: "$1",
+                content: {
+                    msgtype: "m.text",
+                    body: "1",
+                },
+            });
+            // 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 = undefined;
+            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,
+            );
+
+            const ksEvent = await keyshareEventForEvent(aliceClient, event, 1);
+            ksEvent.getContent().sender_key = undefined; // test
+            bobClient.crypto.addInboundGroupSession = jest.fn();
+            await bobDecryptor.onRoomKeyEvent(ksEvent);
+            expect(bobClient.crypto.addInboundGroupSession).not.toHaveBeenCalled();
+        });
 
         it("creates a new keyshare request if we request a keyshare", async function() {
             // make sure that cancelAndResend... creates a new keyshare request

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

@@ -257,6 +257,8 @@ describe("MegolmDecryption", function() {
         });
 
         describe("session reuse and key reshares", () => {
+            const rotationPeriodMs = 999 * 24 * 60 * 60 * 1000; // 999 days, so we don't have to deal with it
+
             let megolmEncryption;
             let aliceDeviceInfo;
             let mockRoom;
@@ -318,7 +320,7 @@ describe("MegolmDecryption", function() {
                     baseApis: mockBaseApis,
                     roomId: ROOM_ID,
                     config: {
-                        rotation_period_ms: 9999999999999,
+                        rotation_period_ms: rotationPeriodMs,
                     },
                 });
                 mockRoom = {
@@ -329,6 +331,31 @@ describe("MegolmDecryption", function() {
                 };
             });
 
+            it("should use larger otkTimeout when preparing to encrypt room", async () => {
+                megolmEncryption.prepareToEncrypt(mockRoom);
+                await megolmEncryption.encryptMessage(mockRoom, "a.fake.type", {
+                    body: "Some text",
+                });
+                expect(mockRoom.getEncryptionTargetMembers).toHaveBeenCalled();
+
+                expect(mockBaseApis.claimOneTimeKeys).toHaveBeenCalledWith(
+                    [['@alice:home.server', 'aliceDevice']], 'signed_curve25519', 10000,
+                );
+            });
+
+            it("should generate a new session if this one needs rotation", async () => {
+                const session = await megolmEncryption.prepareNewSession(false);
+                session.creationTime -= rotationPeriodMs + 10000; // a smidge over the rotation time
+                // Inject expired session which needs rotation
+                megolmEncryption.setupPromise = Promise.resolve(session);
+
+                const prepareNewSessionSpy = jest.spyOn(megolmEncryption, "prepareNewSession");
+                await megolmEncryption.encryptMessage(mockRoom, "a.fake.type", {
+                    body: "Some text",
+                });
+                expect(prepareNewSessionSpy).toHaveBeenCalledTimes(1);
+            });
+
             it("re-uses sessions for sequential messages", async function() {
                 const ct1 = await megolmEncryption.encryptMessage(mockRoom, "a.fake.type", {
                     body: "Some text",

src/crypto/algorithms/base.ts

@@ -179,7 +179,7 @@ export abstract class DecryptionAlgorithm {
      *
      * @param {module:models/event.MatrixEvent} params event key event
      */
-    public onRoomKeyEvent(params: MatrixEvent): void {
+    public async onRoomKeyEvent(params: MatrixEvent): Promise<void> {
         // ignore by default
     }
 

src/crypto/algorithms/megolm.ts

@@ -213,6 +213,8 @@ class OutboundSessionInfo {
                 }
             }
         }
+
+        return false;
     }
 }
 
@@ -231,7 +233,7 @@ class MegolmEncryption extends EncryptionAlgorithm {
     // are using, and which devices we have shared the keys with. It resolves
     // with an OutboundSessionInfo (or undefined, for the first message in the
     // room).
-    private setupPromise = Promise.resolve<OutboundSessionInfo>(undefined);
+    private setupPromise = Promise.resolve<OutboundSessionInfo | null>(null);
 
     // Map of outbound sessions by sessions ID. Used if we need a particular
     // session (the session we're currently using to send is always obtained
@@ -240,8 +242,8 @@ class MegolmEncryption extends EncryptionAlgorithm {
 
     private readonly sessionRotationPeriodMsgs: number;
     private readonly sessionRotationPeriodMs: number;
-    private encryptionPreparation: Promise<void>;
+    private encryptionPreparation?: {
-    private encryptionPreparationMetadata: {
+        promise: Promise<void>;
         startTime: number;
     };
 
@@ -270,193 +272,209 @@ class MegolmEncryption extends EncryptionAlgorithm {
         blocked: IBlockedMap,
         singleOlmCreationPhase = false,
     ): Promise<OutboundSessionInfo> {
-        let session: OutboundSessionInfo;
-
         // takes the previous OutboundSessionInfo, and considers whether to create
         // a new one. Also shares the key with any (new) devices in the room.
-        // Updates `session` to hold the final OutboundSessionInfo.
+        //
+        // Returns the successful session whether keyshare succeeds or not.
         //
         // returns a promise which resolves once the keyshare is successful.
-        const prepareSession = async (oldSession: OutboundSessionInfo) => {
+        const setup = async (oldSession: OutboundSessionInfo | null): Promise<OutboundSessionInfo> => {
-            session = oldSession;
-
             const sharedHistory = isRoomSharedHistory(room);
 
-            // history visibility changed
+            const session = await this.prepareSession(devicesInRoom, sharedHistory, oldSession);
-            if (session && sharedHistory !== session.sharedHistory) {
+
-                session = null;
+            try {
+                await this.shareSession(devicesInRoom, sharedHistory, singleOlmCreationPhase, blocked, session);
+            } catch (e) {
+                logger.error(`Failed to ensure outbound session in ${this.roomId}`, e);
             }
 
-            // need to make a brand new session?
+            return session;
-            if (session && session.needsRotation(this.sessionRotationPeriodMsgs,
-                this.sessionRotationPeriodMs)
-            ) {
-                logger.log("Starting new megolm session because we need to rotate.");
-                session = null;
-            }
-
-            // determine if we have shared with anyone we shouldn't have
-            if (session && session.sharedWithTooManyDevices(devicesInRoom)) {
-                session = null;
-            }
-
-            if (!session) {
-                logger.log(`Starting new megolm session for room ${this.roomId}`);
-                session = await this.prepareNewSession(sharedHistory);
-                logger.log(`Started new megolm session ${session.sessionId} ` +
-                    `for room ${this.roomId}`);
-                this.outboundSessions[session.sessionId] = session;
-            }
-
-            // now check if we need to share with any devices
-            const shareMap: Record<string, DeviceInfo[]> = {};
-
-            for (const [userId, userDevices] of Object.entries(devicesInRoom)) {
-                for (const [deviceId, deviceInfo] of Object.entries(userDevices)) {
-                    const key = deviceInfo.getIdentityKey();
-                    if (key == this.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);
-                    }
-                }
-            }
-
-            const key = this.olmDevice.getOutboundGroupSessionKey(session.sessionId);
-            const payload: IPayload = {
-                type: "m.room_key",
-                content: {
-                    "algorithm": olmlib.MEGOLM_ALGORITHM,
-                    "room_id": this.roomId,
-                    "session_id": session.sessionId,
-                    "session_key": key.key,
-                    "chain_index": key.chain_index,
-                    "org.matrix.msc3061.shared_history": sharedHistory,
-                },
-            };
-            const [devicesWithoutSession, olmSessions] = await olmlib.getExistingOlmSessions(
-                this.olmDevice, this.baseApis, shareMap,
-            );
-
-            await Promise.all([
-                (async () => {
-                    // share keys with devices that we already have a session for
-                    logger.debug(`Sharing keys with existing Olm sessions in ${this.roomId}`, olmSessions);
-                    await this.shareKeyWithOlmSessions(session, key, payload, olmSessions);
-                    logger.debug(`Shared keys with existing Olm sessions in ${this.roomId}`);
-                })(),
-                (async () => {
-                    logger.debug(
-                        `Sharing keys (start phase 1) with new Olm sessions in ${this.roomId}`,
-                        devicesWithoutSession,
-                    );
-                    const errorDevices: IOlmDevice[] = [];
-
-                    // meanwhile, establish olm sessions for devices that we don't
-                    // already have a session for, and share keys with them.  If
-                    // we're doing two phases of olm session creation, use a
-                    // shorter timeout when fetching one-time keys for the first
-                    // phase.
-                    const start = Date.now();
-                    const failedServers: string[] = [];
-                    await this.shareKeyWithDevices(
-                        session, key, payload, devicesWithoutSession, errorDevices,
-                        singleOlmCreationPhase ? 10000 : 2000, failedServers,
-                    );
-                    logger.debug(`Shared keys (end phase 1) with new Olm sessions in ${this.roomId}`);
-
-                    if (!singleOlmCreationPhase && (Date.now() - start < 10000)) {
-                        // perform the second phase of olm session creation if requested,
-                        // and if the first phase didn't take too long
-                        (async () => {
-                            // Retry sending keys to devices that we were unable to establish
-                            // an olm session for.  This time, we use a longer timeout, but we
-                            // do this in the background and don't block anything else while we
-                            // do this.  We only need to retry users from servers that didn't
-                            // respond the first time.
-                            const retryDevices: Record<string, DeviceInfo[]> = {};
-                            const failedServerMap = new Set;
-                            for (const server of failedServers) {
-                                failedServerMap.add(server);
-                            }
-                            const failedDevices = [];
-                            for (const { userId, deviceInfo } of errorDevices) {
-                                const userHS = userId.slice(userId.indexOf(":") + 1);
-                                if (failedServerMap.has(userHS)) {
-                                    retryDevices[userId] = retryDevices[userId] || [];
-                                    retryDevices[userId].push(deviceInfo);
-                                } else {
-                                    // if we aren't going to retry, then handle it
-                                    // as a failed device
-                                    failedDevices.push({ userId, deviceInfo });
-                                }
-                            }
-
-                            logger.debug(`Sharing keys (start phase 2) with new Olm sessions in ${this.roomId}`);
-                            await this.shareKeyWithDevices(
-                                session, key, payload, retryDevices, failedDevices, 30000,
-                            );
-                            logger.debug(`Shared keys (end phase 2) with new Olm sessions in ${this.roomId}`);
-
-                            await this.notifyFailedOlmDevices(session, key, failedDevices);
-                        })();
-                    } else {
-                        await this.notifyFailedOlmDevices(session, key, errorDevices);
-                    }
-                    logger.debug(`Shared keys (all phases done) with new Olm sessions in ${this.roomId}`);
-                })(),
-                (async () => {
-                    logger.debug(`There are ${Object.entries(blocked).length} blocked devices in ${this.roomId}`,
-                        Object.entries(blocked));
-
-                    // also, notify newly blocked devices that they're blocked
-                    logger.debug(`Notifying newly blocked devices in ${this.roomId}`);
-                    const blockedMap: Record<string, Record<string, { device: IBlockedDevice }>> = {};
-                    let blockedCount = 0;
-                    for (const [userId, userBlockedDevices] of Object.entries(blocked)) {
-                        for (const [deviceId, device] of Object.entries(userBlockedDevices)) {
-                            if (
-                                !session.blockedDevicesNotified[userId] ||
-                                session.blockedDevicesNotified[userId][deviceId] === undefined
-                            ) {
-                                blockedMap[userId] = blockedMap[userId] || {};
-                                blockedMap[userId][deviceId] = { device };
-                                blockedCount++;
-                            }
-                        }
-                    }
-
-                    await this.notifyBlockedDevices(session, blockedMap);
-                    logger.debug(`Notified ${blockedCount} newly blocked devices in ${this.roomId}`, blockedMap);
-                })(),
-            ]);
         };
 
-        // helper which returns the session prepared by prepareSession
-        function returnSession() {
-            return session;
-        }
-
         // first wait for the previous share to complete
-        const prom = this.setupPromise.then(prepareSession);
+        const prom = this.setupPromise.then(setup);
 
         // Ensure any failures are logged for debugging
         prom.catch(e => {
-            logger.error(`Failed to ensure outbound session in ${this.roomId}`, e);
+            logger.error(`Failed to setup outbound session in ${this.roomId}`, e);
         });
 
         // setupPromise resolves to `session` whether or not the share succeeds
-        this.setupPromise = prom.then(returnSession, returnSession);
+        this.setupPromise = prom;
 
         // but we return a promise which only resolves if the share was successful.
-        return prom.then(returnSession);
+        return prom;
+    }
+
+    private async prepareSession(
+        devicesInRoom: DeviceInfoMap,
+        sharedHistory: boolean,
+        session: OutboundSessionInfo | null,
+    ): Promise<OutboundSessionInfo> {
+        // history visibility changed
+        if (session && sharedHistory !== session.sharedHistory) {
+            session = null;
+        }
+
+        // need to make a brand new session?
+        if (session?.needsRotation(this.sessionRotationPeriodMsgs, this.sessionRotationPeriodMs)) {
+            logger.log("Starting new megolm session because we need to rotate.");
+            session = null;
+        }
+
+        // determine if we have shared with anyone we shouldn't have
+        if (session?.sharedWithTooManyDevices(devicesInRoom)) {
+            session = null;
+        }
+
+        if (!session) {
+            logger.log(`Starting new megolm session for room ${this.roomId}`);
+            session = await this.prepareNewSession(sharedHistory);
+            logger.log(`Started new megolm session ${session.sessionId} ` +
+                `for room ${this.roomId}`);
+            this.outboundSessions[session.sessionId] = session;
+        }
+
+        return session;
+    }
+
+    private async shareSession(
+        devicesInRoom: DeviceInfoMap,
+        sharedHistory: boolean,
+        singleOlmCreationPhase: boolean,
+        blocked: IBlockedMap,
+        session: OutboundSessionInfo,
+    ) {
+        // now check if we need to share with any devices
+        const shareMap: Record<string, DeviceInfo[]> = {};
+
+        for (const [userId, userDevices] of Object.entries(devicesInRoom)) {
+            for (const [deviceId, deviceInfo] of Object.entries(userDevices)) {
+                const key = deviceInfo.getIdentityKey();
+                if (key == this.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);
+                }
+            }
+        }
+
+        const key = this.olmDevice.getOutboundGroupSessionKey(session.sessionId);
+        const payload: IPayload = {
+            type: "m.room_key",
+            content: {
+                "algorithm": olmlib.MEGOLM_ALGORITHM,
+                "room_id": this.roomId,
+                "session_id": session.sessionId,
+                "session_key": key.key,
+                "chain_index": key.chain_index,
+                "org.matrix.msc3061.shared_history": sharedHistory,
+            },
+        };
+        const [devicesWithoutSession, olmSessions] = await olmlib.getExistingOlmSessions(
+            this.olmDevice, this.baseApis, shareMap,
+        );
+
+        await Promise.all([
+            (async () => {
+                // share keys with devices that we already have a session for
+                logger.debug(`Sharing keys with existing Olm sessions in ${this.roomId}`, olmSessions);
+                await this.shareKeyWithOlmSessions(session, key, payload, olmSessions);
+                logger.debug(`Shared keys with existing Olm sessions in ${this.roomId}`);
+            })(),
+            (async () => {
+                logger.debug(
+                    `Sharing keys (start phase 1) with new Olm sessions in ${this.roomId}`,
+                    devicesWithoutSession,
+                );
+                const errorDevices: IOlmDevice[] = [];
+
+                // meanwhile, establish olm sessions for devices that we don't
+                // already have a session for, and share keys with them.  If
+                // we're doing two phases of olm session creation, use a
+                // shorter timeout when fetching one-time keys for the first
+                // phase.
+                const start = Date.now();
+                const failedServers: string[] = [];
+                await this.shareKeyWithDevices(
+                    session, key, payload, devicesWithoutSession, errorDevices,
+                    singleOlmCreationPhase ? 10000 : 2000, failedServers,
+                );
+                logger.debug(`Shared keys (end phase 1) with new Olm sessions in ${this.roomId}`);
+
+                if (!singleOlmCreationPhase && (Date.now() - start < 10000)) {
+                    // perform the second phase of olm session creation if requested,
+                    // and if the first phase didn't take too long
+                    (async () => {
+                        // Retry sending keys to devices that we were unable to establish
+                        // an olm session for.  This time, we use a longer timeout, but we
+                        // do this in the background and don't block anything else while we
+                        // do this.  We only need to retry users from servers that didn't
+                        // respond the first time.
+                        const retryDevices: Record<string, DeviceInfo[]> = {};
+                        const failedServerMap = new Set;
+                        for (const server of failedServers) {
+                            failedServerMap.add(server);
+                        }
+                        const failedDevices: IOlmDevice[] = [];
+                        for (const { userId, deviceInfo } of errorDevices) {
+                            const userHS = userId.slice(userId.indexOf(":") + 1);
+                            if (failedServerMap.has(userHS)) {
+                                retryDevices[userId] = retryDevices[userId] || [];
+                                retryDevices[userId].push(deviceInfo);
+                            } else {
+                                // if we aren't going to retry, then handle it
+                                // as a failed device
+                                failedDevices.push({ userId, deviceInfo });
+                            }
+                        }
+
+                        logger.debug(`Sharing keys (start phase 2) with new Olm sessions in ${this.roomId}`);
+                        await this.shareKeyWithDevices(
+                            session, key, payload, retryDevices, failedDevices, 30000,
+                        );
+                        logger.debug(`Shared keys (end phase 2) with new Olm sessions in ${this.roomId}`);
+
+                        await this.notifyFailedOlmDevices(session, key, failedDevices);
+                    })();
+                } else {
+                    await this.notifyFailedOlmDevices(session, key, errorDevices);
+                }
+                logger.debug(`Shared keys (all phases done) with new Olm sessions in ${this.roomId}`);
+            })(),
+            (async () => {
+                logger.debug(`There are ${Object.entries(blocked).length} blocked devices in ${this.roomId}`,
+                    Object.entries(blocked));
+
+                // also, notify newly blocked devices that they're blocked
+                logger.debug(`Notifying newly blocked devices in ${this.roomId}`);
+                const blockedMap: Record<string, Record<string, { device: IBlockedDevice }>> = {};
+                let blockedCount = 0;
+                for (const [userId, userBlockedDevices] of Object.entries(blocked)) {
+                    for (const [deviceId, device] of Object.entries(userBlockedDevices)) {
+                        if (
+                            !session.blockedDevicesNotified[userId] ||
+                            session.blockedDevicesNotified[userId][deviceId] === undefined
+                        ) {
+                            blockedMap[userId] = blockedMap[userId] || {};
+                            blockedMap[userId][deviceId] = { device };
+                            blockedCount++;
+                        }
+                    }
+                }
+
+                await this.notifyBlockedDevices(session, blockedMap);
+                logger.debug(`Notified ${blockedCount} newly blocked devices in ${this.roomId}`, blockedMap);
+            })(),
+        ]);
     }
 
     /**
@@ -866,7 +884,7 @@ class MegolmEncryption extends EncryptionAlgorithm {
         logger.debug(`Ensuring Olm sessions for devices in ${this.roomId}`);
         const devicemap = await olmlib.ensureOlmSessionsForDevices(
             this.olmDevice, this.baseApis, devicesByUser, false, otkTimeout, failedServers,
-            logger.withPrefix(`[${this.roomId}]`),
+            logger.withPrefix?.(`[${this.roomId}]`),
         );
         logger.debug(`Ensured Olm sessions for devices in ${this.roomId}`);
 
@@ -1006,11 +1024,11 @@ class MegolmEncryption extends EncryptionAlgorithm {
      * @param {module:models/room} room the room the event is in
      */
     public prepareToEncrypt(room: Room): void {
-        if (this.encryptionPreparation) {
+        if (this.encryptionPreparation != null) {
             // We're already preparing something, so don't do anything else.
             // FIXME: check if we need to restart
             // (https://github.com/matrix-org/matrix-js-sdk/issues/1255)
-            const elapsedTime = Date.now() - this.encryptionPreparationMetadata.startTime;
+            const elapsedTime = Date.now() - this.encryptionPreparation.startTime;
             logger.debug(
                 `Already started preparing to encrypt for ${this.roomId} ` +
                 `${elapsedTime} ms ago, skipping`,
@@ -1020,32 +1038,31 @@ class MegolmEncryption extends EncryptionAlgorithm {
 
         logger.debug(`Preparing to encrypt events for ${this.roomId}`);
 
-        this.encryptionPreparationMetadata = {
+        this.encryptionPreparation = {
             startTime: Date.now(),
-        };
+            promise: (async () => {
-        this.encryptionPreparation = (async () => {
+                try {
-            try {
+                    logger.debug(`Getting devices in ${this.roomId}`);
-                logger.debug(`Getting devices in ${this.roomId}`);
+                    const [devicesInRoom, blocked] = await this.getDevicesInRoom(room);
-                const [devicesInRoom, blocked] = await this.getDevicesInRoom(room);
 
-                if (this.crypto.getGlobalErrorOnUnknownDevices()) {
+                    if (this.crypto.getGlobalErrorOnUnknownDevices()) {
-                    // Drop unknown devices for now.  When the message gets sent, we'll
+                        // Drop unknown devices for now.  When the message gets sent, we'll
-                    // throw an error, but we'll still be prepared to send to the known
+                        // throw an error, but we'll still be prepared to send to the known
-                    // devices.
+                        // devices.
-                    this.removeUnknownDevices(devicesInRoom);
+                        this.removeUnknownDevices(devicesInRoom);
+                    }
+
+                    logger.debug(`Ensuring outbound session in ${this.roomId}`);
+                    await this.ensureOutboundSession(room, devicesInRoom, blocked, true);
+
+                    logger.debug(`Ready to encrypt events for ${this.roomId}`);
+                } catch (e) {
+                    logger.error(`Failed to prepare to encrypt events for ${this.roomId}`, e);
+                } finally {
+                    delete this.encryptionPreparation;
                 }
-
+            })(),
-                logger.debug(`Ensuring outbound session in ${this.roomId}`);
+        };
-                await this.ensureOutboundSession(room, devicesInRoom, blocked, true);
-
-                logger.debug(`Ready to encrypt events for ${this.roomId}`);
-            } catch (e) {
-                logger.error(`Failed to prepare to encrypt events for ${this.roomId}`, e);
-            } finally {
-                delete this.encryptionPreparationMetadata;
-                delete this.encryptionPreparation;
-            }
-        })();
     }
 
     /**
@@ -1060,12 +1077,12 @@ class MegolmEncryption extends EncryptionAlgorithm {
     public async encryptMessage(room: Room, eventType: string, content: object): Promise<object> {
         logger.log(`Starting to encrypt event for ${this.roomId}`);
 
-        if (this.encryptionPreparation) {
+        if (this.encryptionPreparation != null) {
             // If we started sending keys, wait for it to be done.
             // FIXME: check if we need to cancel
             // (https://github.com/matrix-org/matrix-js-sdk/issues/1255)
             try {
-                await this.encryptionPreparation;
+                await this.encryptionPreparation.promise;
             } catch (e) {
                 // ignore any errors -- if the preparation failed, we'll just
                 // restart everything here
@@ -1405,7 +1422,7 @@ class MegolmDecryption extends DecryptionAlgorithm {
         if (!senderPendingEvents.has(sessionId)) {
             senderPendingEvents.set(sessionId, new Set());
         }
-        senderPendingEvents.get(sessionId).add(event);
+        senderPendingEvents.get(sessionId)?.add(event);
     }
 
     /**
@@ -1439,17 +1456,17 @@ class MegolmDecryption extends DecryptionAlgorithm {
      *
      * @param {module:models/event.MatrixEvent} event key event
      */
-    public onRoomKeyEvent(event: MatrixEvent): Promise<void> {
+    public async onRoomKeyEvent(event: MatrixEvent): Promise<void> {
-        const content = event.getContent();
+        const content = event.getContent<Partial<IMessage["content"]>>();
-        const sessionId = content.session_id;
         let senderKey = event.getSenderKey();
-        let forwardingKeyChain = [];
+        let forwardingKeyChain: string[] = [];
         let exportFormat = false;
-        let keysClaimed;
+        let keysClaimed: ReturnType<MatrixEvent["getKeysClaimed"]>;
 
         if (!content.room_id ||
-            !sessionId ||
+            !content.session_key ||
-            !content.session_key
+            !content.session_id ||
+            !content.algorithm
         ) {
             logger.error("key event is missing fields");
             return;
@@ -1462,20 +1479,18 @@ class MegolmDecryption extends DecryptionAlgorithm {
 
         if (event.getType() == "m.forwarded_room_key") {
             exportFormat = true;
-            forwardingKeyChain = content.forwarding_curve25519_key_chain;
+            forwardingKeyChain = Array.isArray(content.forwarding_curve25519_key_chain) ?
-            if (!Array.isArray(forwardingKeyChain)) {
+                content.forwarding_curve25519_key_chain : [];
-                forwardingKeyChain = [];
-            }
 
             // copy content before we modify it
             forwardingKeyChain = forwardingKeyChain.slice();
             forwardingKeyChain.push(senderKey);
 
-            senderKey = content.sender_key;
+            if (!content.sender_key) {
-            if (!senderKey) {
                 logger.error("forwarded_room_key event is missing sender_key field");
                 return;
             }
+            senderKey = content.sender_key;
 
             const ed25519Key = content.sender_claimed_ed25519_key;
             if (!ed25519Key) {
@@ -1496,34 +1511,39 @@ class MegolmDecryption extends DecryptionAlgorithm {
         if (content["org.matrix.msc3061.shared_history"]) {
             extraSessionData.sharedHistory = true;
         }
-        return this.olmDevice.addInboundGroupSession(
+
-            content.room_id, senderKey, forwardingKeyChain, sessionId,
+        try {
-            content.session_key, keysClaimed,
+            await this.olmDevice.addInboundGroupSession(
-            exportFormat, extraSessionData,
+                content.room_id,
-        ).then(() => {
+                senderKey,
+                forwardingKeyChain,
+                content.session_id,
+                content.session_key,
+                keysClaimed,
+                exportFormat,
+                extraSessionData,
+            );
+
             // have another go at decrypting events sent with this session.
-            this.retryDecryption(senderKey, sessionId)
+            if (await this.retryDecryption(senderKey, content.session_id)) {
-                .then((success) => {
+                // cancel any outstanding room key requests for this session.
-                    // cancel any outstanding room key requests for this session.
+                // Only do this if we managed to decrypt every message in the
-                    // Only do this if we managed to decrypt every message in the
+                // session, because if we didn't, we leave the other key
-                    // session, because if we didn't, we leave the other key
+                // requests in the hopes that someone sends us a key that
-                    // requests in the hopes that someone sends us a key that
+                // includes an earlier index.
-                    // includes an earlier index.
+                this.crypto.cancelRoomKeyRequest({
-                    if (success) {
+                    algorithm: content.algorithm,
-                        this.crypto.cancelRoomKeyRequest({
+                    room_id: content.room_id,
-                            algorithm: content.algorithm,
+                    session_id: content.session_id,
-                            room_id: content.room_id,
+                    sender_key: senderKey,
-                            session_id: content.session_id,
-                            sender_key: senderKey,
-                        });
-                    }
                 });
-        }).then(() => {
+            }
+
             // don't wait for the keys to be backed up for the server
-            this.crypto.backupManager.backupGroupSession(senderKey, content.session_id);
+            await this.crypto.backupManager.backupGroupSession(senderKey, content.session_id);
-        }).catch((e) => {
+        } catch (e) {
             logger.error(`Error handling m.room_key_event: ${e}`);
-        });
+        }
     }
 
     /**
@@ -1716,7 +1736,10 @@ class MegolmDecryption extends DecryptionAlgorithm {
      * @param {boolean} [opts.untrusted] whether the key should be considered as untrusted
      * @param {string} [opts.source] where the key came from
      */
-    public importRoomKey(session: IMegolmSessionData, opts: any = {}): Promise<void> {
+    public importRoomKey(
+        session: IMegolmSessionData,
+        opts: { untrusted?: boolean, source?: string } = {},
+    ): Promise<void> {
         const extraSessionData: any = {};
         if (opts.untrusted || session.untrusted) {
             extraSessionData.untrusted = true;

src/crypto/algorithms/olm.ts

@@ -52,7 +52,7 @@ interface IMessage {
  */
 class OlmEncryption extends EncryptionAlgorithm {
     private sessionPrepared = false;
-    private prepPromise: Promise<void> = null;
+    private prepPromise: Promise<void> | null = null;
 
     /**
      * @private
@@ -117,11 +117,11 @@ class OlmEncryption extends EncryptionAlgorithm {
             ciphertext: {},
         };
 
-        const promises = [];
+        const promises: Promise<void>[] = [];
 
         for (let i = 0; i < users.length; ++i) {
             const userId = users[i];
-            const devices = this.crypto.getStoredDevicesForUser(userId);
+            const devices = this.crypto.getStoredDevicesForUser(userId) || [];
 
             for (let j = 0; j < devices.length; ++j) {
                 const deviceInfo = devices[j];
@@ -240,7 +240,7 @@ class OlmDecryption extends DecryptionAlgorithm {
             throw new DecryptionError(
                 "OLM_BAD_ROOM",
                 "Message intended for room " + payload.room_id, {
-                    reported_room: event.getRoomId(),
+                    reported_room: event.getRoomId() || "ROOM_ID_UNDEFINED",
                 },
             );
         }

src/crypto/index.ts

@@ -2590,7 +2590,7 @@ export class Crypto extends TypedEventEmitter<CryptoEvent, CryptoEventHandlerMap
         // because it first stores in memory. We should await the promise only
         // after all the in-memory state (roomEncryptors and _roomList) has been updated
         // to avoid races when calling this method multiple times. Hence keep a hold of the promise.
-        let storeConfigPromise = null;
+        let storeConfigPromise: Promise<void> = null;
         if (!existingConfig) {
             storeConfigPromise = this.roomList.setRoomEncryption(roomId, config);
         }

src/crypto/olmlib.ts

@@ -76,7 +76,7 @@ export interface IOlmSessionResult {
 export async function encryptMessageForDevice(
     resultsObject: Record<string, string>,
     ourUserId: string,
-    ourDeviceId: string,
+    ourDeviceId: string | undefined,
     olmDevice: OlmDevice,
     recipientUserId: string,
     recipientDevice: DeviceInfo,