Reduce Megolm blocking and add cancellation

spec/unit/crypto/algorithms/megolm.spec.ts

@@ -1,5 +1,5 @@
 /*
-Copyright 2022 The Matrix.org Foundation C.I.C.
+Copyright 2022 - 2023 The Matrix.org Foundation C.I.C.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -16,6 +16,7 @@ limitations under the License.
 
 import { mocked, MockedObject } from "jest-mock";
 
+import type { DeviceInfoMap } from "../../../../src/crypto/DeviceList";
 import "../../../olm-loader";
 import type { OutboundGroupSession } from "@matrix-org/olm";
 import * as algorithms from "../../../../src/crypto/algorithms";
@@ -33,6 +34,7 @@ import { ClientEvent, MatrixClient, RoomMember } from "../../../../src";
 import { DeviceInfo, IDevice } from "../../../../src/crypto/deviceinfo";
 import { DeviceTrustLevel } from "../../../../src/crypto/CrossSigning";
 import { MegolmEncryption as MegolmEncryptionClass } from "../../../../src/crypto/algorithms/megolm";
+import { sleep } from "../../../../src/utils";
 
 const MegolmDecryption = algorithms.DECRYPTION_CLASSES.get("m.megolm.v1.aes-sha2")!;
 const MegolmEncryption = algorithms.ENCRYPTION_CLASSES.get("m.megolm.v1.aes-sha2")!;
@@ -58,6 +60,12 @@ describe("MegolmDecryption", function () {
 
     beforeEach(async function () {
         mockCrypto = testUtils.mock(Crypto, "Crypto") as MockedObject<Crypto>;
+
+        // @ts-ignore assigning to readonly prop
+        mockCrypto.backupManager = {
+            backupGroupSession: () => {},
+        };
+
         mockBaseApis = {
             claimOneTimeKeys: jest.fn(),
             sendToDevice: jest.fn(),
@@ -314,10 +322,6 @@ describe("MegolmDecryption", function () {
             let olmDevice: OlmDevice;
 
             beforeEach(async () => {
-                // @ts-ignore assigning to readonly prop
-                mockCrypto.backupManager = {
-                    backupGroupSession: () => {},
-                };
                 const cryptoStore = new MemoryCryptoStore();
 
                 olmDevice = new OlmDevice(cryptoStore);
@@ -515,6 +519,87 @@ describe("MegolmDecryption", function () {
         });
     });
 
+    describe("prepareToEncrypt", () => {
+        let megolm: MegolmEncryptionClass;
+        let room: jest.Mocked<Room>;
+
+        const deviceMap: DeviceInfoMap = {
+            "user-a": {
+                "device-a": new DeviceInfo("device-a"),
+                "device-b": new DeviceInfo("device-b"),
+                "device-c": new DeviceInfo("device-c"),
+            },
+            "user-b": {
+                "device-d": new DeviceInfo("device-d"),
+                "device-e": new DeviceInfo("device-e"),
+                "device-f": new DeviceInfo("device-f"),
+            },
+            "user-c": {
+                "device-g": new DeviceInfo("device-g"),
+                "device-h": new DeviceInfo("device-h"),
+                "device-i": new DeviceInfo("device-i"),
+            },
+        };
+
+        beforeEach(() => {
+            room = testUtils.mock(Room, "Room") as jest.Mocked<Room>;
+            room.getEncryptionTargetMembers.mockImplementation(async () => [
+                new RoomMember(room.roomId, "@user:example.org"),
+            ]);
+            room.getBlacklistUnverifiedDevices.mockReturnValue(false);
+
+            mockCrypto.downloadKeys.mockImplementation(async () => deviceMap);
+
+            mockCrypto.checkDeviceTrust.mockImplementation(() => new DeviceTrustLevel(true, true, true, true));
+
+            const olmDevice = new OlmDevice(new MemoryCryptoStore());
+            megolm = new MegolmEncryptionClass({
+                userId: "@user:id",
+                deviceId: "12345",
+                crypto: mockCrypto,
+                olmDevice,
+                baseApis: mockBaseApis,
+                roomId: room.roomId,
+                config: {
+                    algorithm: "m.megolm.v1.aes-sha2",
+                    rotation_period_ms: 9_999_999,
+                },
+            });
+        });
+
+        it("checks each device", async () => {
+            megolm.prepareToEncrypt(room);
+            //@ts-ignore private member access, gross
+            await megolm.encryptionPreparation?.promise;
+
+            for (const userId in deviceMap) {
+                for (const deviceId in deviceMap[userId]) {
+                    expect(mockCrypto.checkDeviceTrust).toHaveBeenCalledWith(userId, deviceId);
+                }
+            }
+        });
+
+        it("defers before completing", async () => {
+            megolm.prepareToEncrypt(room);
+            // Ensure that `Crypto#checkDeviceTrust` has been called *fewer*
+            // than the full nine times, after yielding once.
+            await sleep(0);
+            const callCount = mockCrypto.checkDeviceTrust.mock.calls.length;
+            expect(callCount).toBeLessThan(9);
+        });
+
+        it("is cancellable", async () => {
+            const stop = megolm.prepareToEncrypt(room);
+
+            const before = mockCrypto.checkDeviceTrust.mock.calls.length;
+            stop();
+
+            // Ensure that no more devices were checked after cancellation.
+            await sleep(10);
+            expect(mockCrypto.checkDeviceTrust).toHaveBeenCalledTimes(before);
+        });
+    });
+
     it("notifies devices that have been blocked", async function () {
         const aliceClient = new TestClient("@alice:example.com", "alicedevice").client;
         const bobClient1 = new TestClient("@bob:example.com", "bobdevice1").client;

spec/unit/utils.spec.ts

@@ -1,3 +1,19 @@
+/*
+Copyright 2023 The Matrix.org Foundation C.I.C.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
 import * as utils from "../../src/utils";
 import {
     alphabetPad,
@@ -587,4 +603,22 @@ describe("utils", function () {
             expect(utils.isSupportedReceiptType("this is a receipt type")).toBeFalsy();
         });
     });
+
+    describe("sleep", () => {
+        it("resolves", async () => {
+            await utils.sleep(0);
+        });
+
+        it("resolves with the provided value", async () => {
+            const expected = Symbol("hi");
+            const result = await utils.sleep(0, expected);
+            expect(result).toBe(expected);
+        });
+    });
+
+    describe("immediate", () => {
+        it("resolves", async () => {
+            await utils.immediate();
+        });
+    });
 });

src/crypto/algorithms/megolm.ts

@@ -1,5 +1,5 @@
 /*
-Copyright 2015 - 2021 The Matrix.org Foundation C.I.C.
+Copyright 2015 - 2021, 2023 The Matrix.org Foundation C.I.C.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -43,6 +43,7 @@ import { IMegolmEncryptedContent, IncomingRoomKeyRequest, IEncryptedContent } fr
 import { RoomKeyRequestState } from "../OutgoingRoomKeyRequestManager";
 import { OlmGroupSessionExtraData } from "../../@types/crypto";
 import { MatrixError } from "../../http-api";
+import { immediate } from "../../utils";
 
 // determine whether the key can be shared with invitees
 export function isRoomSharedHistory(room: Room): boolean {
@@ -73,7 +74,6 @@ export interface IOlmDevice<T = DeviceInfo> {
     deviceInfo: T;
 }
 
-/* eslint-disable camelcase */
 export interface IOutboundGroupSessionKey {
     chain_index: number;
     key: string;
@@ -106,7 +106,6 @@ interface IPayload extends Partial<IMessage> {
     algorithm?: string;
     sender_key?: string;
 }
-/* eslint-enable camelcase */
 
 interface SharedWithData {
     // The identity key of the device we shared with
@@ -223,6 +222,7 @@ export class MegolmEncryption extends EncryptionAlgorithm {
     private encryptionPreparation?: {
         promise: Promise<void>;
         startTime: number;
+        cancel: () => void;
     };
 
     protected readonly roomId: string;
@@ -974,30 +974,36 @@ export class MegolmEncryption extends EncryptionAlgorithm {
      * send, in order to speed up sending of the message.
      *
      * @param room - the room the event is in
+     * @returns A function that, when called, will stop the preparation
      */
-    public prepareToEncrypt(room: Room): void {
+    public prepareToEncrypt(room: Room): () => void {
         if (room.roomId !== this.roomId) {
             throw new Error("MegolmEncryption.prepareToEncrypt called on unexpected room");
         }
 
         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.encryptionPreparation.startTime;
             this.prefixedLogger.debug(
                 `Already started preparing to encrypt for this room ${elapsedTime}ms ago, skipping`,
             );
-            return;
+            return this.encryptionPreparation.cancel;
         }
 
         this.prefixedLogger.debug("Preparing to encrypt events");
 
+        let cancelled = false;
+        const isCancelled = (): boolean => cancelled;
+
         this.encryptionPreparation = {
             startTime: Date.now(),
             promise: (async (): Promise<void> => {
                 try {
-                    const [devicesInRoom, blocked] = await this.getDevicesInRoom(room);
+                    // Attempt to enumerate the devices in room, and gracefully
+                    // handle cancellation if it occurs.
+                    const getDevicesResult = await this.getDevicesInRoom(room, false, isCancelled);
+                    if (getDevicesResult === null) return;
+                    const [devicesInRoom, blocked] = getDevicesResult;
 
                     if (this.crypto.globalErrorOnUnknownDevices) {
                         // Drop unknown devices for now.  When the message gets sent, we'll
@@ -1016,7 +1022,16 @@ export class MegolmEncryption extends EncryptionAlgorithm {
                     delete this.encryptionPreparation;
                 }
             })(),
+
+            cancel: (): void => {
+                // The caller has indicated that the process should be cancelled,
+                // so tell the promise that we'd like to halt, and reset the preparation state.
+                cancelled = true;
+                delete this.encryptionPreparation;
+            },
         };
+
+        return this.encryptionPreparation.cancel;
     }
 
     /**
@@ -1165,17 +1180,32 @@ export class MegolmEncryption extends EncryptionAlgorithm {
      *
      * @param forceDistributeToUnverified - if set to true will include the unverified devices
      * even if setting is set to block them (useful for verification)
+     * @param isCancelled - will cause the procedure to abort early if and when it starts
+     * returning `true`. If omitted, cancellation won't happen.
      *
-     * @returns Promise which resolves to an array whose
-     *     first element is a map from userId to deviceId to deviceInfo indicating
+     * @returns Promise which resolves to `null`, or an array whose
+     *     first element is a {@link DeviceInfoMap} indicating
      *     the devices that messages should be encrypted to, and whose second
      *     element is a map from userId to deviceId to data indicating the devices
-     *     that are in the room but that have been blocked
+     *     that are in the room but that have been blocked.
+     *     If `isCancelled` is provided and returns `true` while processing, `null`
+     *     will be returned.
+     *     If `isCancelled` is not provided, the Promise will never resolve to `null`.
      */
+    private async getDevicesInRoom(
+        room: Room,
+        forceDistributeToUnverified?: boolean,
+    ): Promise<[DeviceInfoMap, IBlockedMap]>;
+    private async getDevicesInRoom(
+        room: Room,
+        forceDistributeToUnverified?: boolean,
+        isCancelled?: () => boolean,
+    ): Promise<null | [DeviceInfoMap, IBlockedMap]>;
     private async getDevicesInRoom(
         room: Room,
         forceDistributeToUnverified = false,
-    ): Promise<[DeviceInfoMap, IBlockedMap]> {
+        isCancelled?: () => boolean,
+    ): Promise<null | [DeviceInfoMap, IBlockedMap]> {
         const members = await room.getEncryptionTargetMembers();
         this.prefixedLogger.debug(
             `Encrypting for users (shouldEncryptForInvitedMembers: ${room.shouldEncryptForInvitedMembers()}):`,
@@ -1201,6 +1231,11 @@ export class MegolmEncryption extends EncryptionAlgorithm {
         // See https://github.com/vector-im/element-web/issues/2305 for details.
         const devices = await this.crypto.downloadKeys(roomMembers, false);
         const blocked: IBlockedMap = {};
+
+        if (isCancelled?.() === true) {
+            return null;
+        }
+
         // remove any blocked devices
         for (const userId in devices) {
             if (!devices.hasOwnProperty(userId)) {
@@ -1213,6 +1248,11 @@ export class MegolmEncryption extends EncryptionAlgorithm {
                     continue;
                 }
 
+                // Yield prior to checking each device so that we don't block
+                // updating/rendering for too long.
+                // See https://github.com/vector-im/element-web/issues/21612
+                if (isCancelled !== undefined) await immediate();
+                if (isCancelled?.() === true) return null;
                 const deviceTrust = this.crypto.checkDeviceTrust(userId, deviceId);
 
                 if (

src/utils.ts

@@ -1,6 +1,5 @@
 /*
-Copyright 2015, 2016 OpenMarket Ltd
-Copyright 2019 The Matrix.org Foundation C.I.C.
+Copyright 2015, 2016, 2019, 2023 The Matrix.org Foundation C.I.C.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -392,13 +391,22 @@ export function ensureNoTrailingSlash(url?: string): string | undefined {
     }
 }
 
-// Returns a promise which resolves with a given value after the given number of ms
+/**
+ * Returns a promise which resolves with a given value after the given number of ms
+ */
 export function sleep<T>(ms: number, value?: T): Promise<T> {
     return new Promise((resolve) => {
         setTimeout(resolve, ms, value);
     });
 }
 
+/**
+ * Promise/async version of {@link setImmediate}.
+ */
+export function immediate(): Promise<void> {
+    return new Promise(setImmediate);
+}
+
 export function isNullOrUndefined(val: any): boolean {
     return val === null || val === undefined;
 }