From ddce1bcd28c4fd0d756694fc2f4e6e77af2e4600 Mon Sep 17 00:00:00 2001 From: Clark Fischer Date: Mon, 16 Jan 2023 07:35:23 -0800 Subject: [PATCH 1/3] Add async `setImmediate` util Adds an async/promise-based version of `setImmediate`. Note that, despite being poorly adopted, `setImmediate` is polyfilled, and should be more performant than `sleep(0)`. Signed-off-by: Clark Fischer --- spec/unit/utils.spec.ts | 34 ++++++++++++++++++++++++++++++++++ src/utils.ts | 14 +++++++++++--- 2 files changed, 45 insertions(+), 3 deletions(-) diff --git a/spec/unit/utils.spec.ts b/spec/unit/utils.spec.ts index 8104cba08..caf15a43e 100644 --- a/spec/unit/utils.spec.ts +++ b/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(); + }); + }); }); diff --git a/src/utils.ts b/src/utils.ts index 5134c8a4d..6a15e9744 100644 --- a/src/utils.ts +++ b/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(ms: number, value?: T): Promise { return new Promise((resolve) => { setTimeout(resolve, ms, value); }); } +/** + * Promise/async version of {@link setImmediate}. + */ +export function immediate(): Promise { + return new Promise(setImmediate); +} + export function isNullOrUndefined(val: any): boolean { return val === null || val === undefined; } From b76e7ca7826ba3073a42e539204a88f6104d372e Mon Sep 17 00:00:00 2001 From: Clark Fischer Date: Sun, 8 Jan 2023 10:43:49 -0800 Subject: [PATCH 2/3] Reduce blocking while pre-fetching Megolm keys Currently, calling `Client#prepareToEncrypt` in a megolm room has the potential to block for multiple seconds while it crunches numbers. Sleeping for 0 seconds (approximating `setImmediate`) allows the engine to process other events, updates, or re-renders in between checks. See - https://github.com/vector-im/element-web/issues/21612 - https://github.com/vector-im/element-web/issues/11836 Signed-off-by: Clark Fischer --- spec/unit/crypto/algorithms/megolm.spec.ts | 84 ++++++++++++++++++++-- src/crypto/algorithms/megolm.ts | 9 ++- 2 files changed, 85 insertions(+), 8 deletions(-) diff --git a/spec/unit/crypto/algorithms/megolm.spec.ts b/spec/unit/crypto/algorithms/megolm.spec.ts index 973ec0bd2..0008501af 100644 --- a/spec/unit/crypto/algorithms/megolm.spec.ts +++ b/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; + + // @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,76 @@ describe("MegolmDecryption", function () { }); }); + describe("prepareToEncrypt", () => { + let megolm: MegolmEncryptionClass; + let room: jest.Mocked; + + 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.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("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; diff --git a/src/crypto/algorithms/megolm.ts b/src/crypto/algorithms/megolm.ts index 163d3953d..e7daf4b7c 100644 --- a/src/crypto/algorithms/megolm.ts +++ b/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 { deviceInfo: T; } -/* eslint-disable camelcase */ export interface IOutboundGroupSessionKey { chain_index: number; key: string; @@ -106,7 +106,6 @@ interface IPayload extends Partial { algorithm?: string; sender_key?: string; } -/* eslint-enable camelcase */ interface SharedWithData { // The identity key of the device we shared with @@ -1213,6 +1212,10 @@ 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 + await immediate(); const deviceTrust = this.crypto.checkDeviceTrust(userId, deviceId); if ( From 1ee487a2ff48b33d6250fc1828f4d2e02f0fa534 Mon Sep 17 00:00:00 2001 From: Clark Fischer Date: Sun, 8 Jan 2023 13:55:03 -0800 Subject: [PATCH 3/3] 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 https://github.com/matrix-org/matrix-js-sdk/issues/1255 Closes #1255 Signed-off-by: Clark Fischer --- spec/unit/crypto/algorithms/megolm.spec.ts | 11 +++++ src/crypto/algorithms/megolm.ts | 57 ++++++++++++++++++---- 2 files changed, 58 insertions(+), 10 deletions(-) diff --git a/spec/unit/crypto/algorithms/megolm.spec.ts b/spec/unit/crypto/algorithms/megolm.spec.ts index 0008501af..9f0ffae99 100644 --- a/spec/unit/crypto/algorithms/megolm.spec.ts +++ b/spec/unit/crypto/algorithms/megolm.spec.ts @@ -587,6 +587,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 () { diff --git a/src/crypto/algorithms/megolm.ts b/src/crypto/algorithms/megolm.ts index e7daf4b7c..ff7e29264 100644 --- a/src/crypto/algorithms/megolm.ts +++ b/src/crypto/algorithms/megolm.ts @@ -222,6 +222,7 @@ export class MegolmEncryption extends EncryptionAlgorithm { private encryptionPreparation?: { promise: Promise; startTime: number; + cancel: () => void; }; protected readonly roomId: string; @@ -973,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 => { 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 @@ -1015,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; } /** @@ -1164,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; private async getDevicesInRoom( room: Room, forceDistributeToUnverified = false, - ): Promise<[DeviceInfoMap, IBlockedMap]> { + isCancelled?: () => boolean, + ): Promise { const members = await room.getEncryptionTargetMembers(); this.prefixedLogger.debug( `Encrypting for users (shouldEncryptForInvitedMembers: ${room.shouldEncryptForInvitedMembers()}):`, @@ -1200,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)) { @@ -1215,7 +1251,8 @@ export class MegolmEncryption extends EncryptionAlgorithm { // 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 - await immediate(); + if (isCancelled !== undefined) await immediate(); + if (isCancelled?.() === true) return null; const deviceTrust = this.crypto.checkDeviceTrust(userId, deviceId); if (