From b76e7ca7826ba3073a42e539204a88f6104d372e Mon Sep 17 00:00:00 2001 From: Clark Fischer Date: Sun, 8 Jan 2023 10:43:49 -0800 Subject: [PATCH] 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 (