Dima/matrix-js-sdk
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #3035 from clarkf/megolm-cancellation

Browse Source 
Reduce Megolm blocking and add cancellation
This commit is contained in:
Robin
2023-01-24 12:26:23 -05:00
committed by GitHub
parent 7ed787b86a 40a4c8d954
commit 70656e954f
4 changed files with 187 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
spec/unit/crypto/algorithms/megolm.spec.ts
+90 -5
View File
@@ -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
+34
View File
@@ -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
+52 -12
View File
@@ -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
+11 -3
View File
@@ -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;
}