Make prepareToEncrypt cancellable.

NOTE: This commit introduces a backwards-compatible API change.

Adds the ability to cancel `MegolmEncryption#prepareToEncrypt` by returning
a cancellation function. The bulk of the processing happens in
`getDevicesInRoom`, which now accepts a 'getter' that allows the caller to
indicate cancellation.

See matrix-org#1255
Closes matrix-org#1255

Signed-off-by: Clark Fischer <[email protected]>

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

@@ -585,6 +585,17 @@ describe("MegolmDecryption", function () {
             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 () {

src/crypto/algorithms/megolm.ts

@@ -57,6 +57,8 @@ export function isRoomSharedHistory(room: Room): boolean {
     return ["world_readable", "shared"].includes(visibility);
 }
 
+type Stop = () => void;
+
 interface IBlockedDevice {
     code: string;
     reason: string;
@@ -224,6 +226,7 @@ export class MegolmEncryption extends EncryptionAlgorithm {
     private encryptionPreparation?: {
         promise: Promise<void>;
         startTime: number;
+        cancel: Stop;
     };
 
     protected readonly roomId: string;
@@ -947,27 +950,32 @@ 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): Stop {
         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;
-            logger.debug(
-                `Already started preparing to encrypt for ${this.roomId} ` + `${elapsedTime} ms ago, skipping`,
-            );
-            return;
+            logger.debug(`Already started preparing to encrypt for ${this.roomId} ${elapsedTime} ms ago, skipping`);
+            return this.encryptionPreparation.cancel;
         }
 
         logger.debug(`Preparing to encrypt events for ${this.roomId}`);
 
+        let cancelled = false;
+        const isCancelled = (): boolean => cancelled;
+
         this.encryptionPreparation = {
             startTime: Date.now(),
             promise: (async (): Promise<void> => {
                 try {
                     logger.debug(`Getting devices in ${this.roomId}`);
-                    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
@@ -986,7 +994,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;
     }
 
     /**
@@ -1014,7 +1031,7 @@ export class MegolmEncryption extends EncryptionAlgorithm {
          * clients should ensure that encryption does not hinder the verification.
          */
         const forceDistributeToUnverified = this.isVerificationEvent(eventType, content);
-        const [devicesInRoom, blocked] = await this.getDevicesInRoom(room, forceDistributeToUnverified);
+        const [devicesInRoom, blocked] = (await this.getDevicesInRoom(room, forceDistributeToUnverified))!;
 
         // check if any of these devices are not yet known to the user.
         // if so, warn the user so they can verify or ignore.
@@ -1135,17 +1152,23 @@ 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 = false,
-    ): Promise<[DeviceInfoMap, IBlockedMap]> {
+        isCancelled?: () => boolean,
+    ): Promise<null | [DeviceInfoMap, IBlockedMap]> {
         const members = await room.getEncryptionTargetMembers();
         const roomMembers = members.map(function (u) {
             return u.userId;
@@ -1166,6 +1189,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)) {
@@ -1182,6 +1210,7 @@ export class MegolmEncryption extends EncryptionAlgorithm {
                 // updating/rendering for too long.
                 // See https://github.com/vector-im/element-web/issues/21612
                 await sleep(0);
+                if (isCancelled?.() === true) return null;
                 const deviceTrust = this.crypto.checkDeviceTrust(userId, deviceId);
 
                 if (