From 06e8cea63df852f70793f3f9e83fa78d2cdce1a2 Mon Sep 17 00:00:00 2001 From: Richard van der Hoff <1389908+richvdh@users.noreply.github.com> Date: Tue, 16 Jan 2024 13:31:21 +0000 Subject: [PATCH] Emit events during migration from libolm (#3982) * Fix `CryptoStore.countEndToEndSessions` This was apparently never tested, and was implemented incorrectly. * Add `CryptoStore.countEndToEndInboundGroupSessions` * Emit events to indicate migration progress --- spec/integ/crypto/rust-crypto.spec.ts | 25 ++++++++- spec/unit/crypto/store/CryptoStore.spec.ts | 8 +++ spec/unit/rust-crypto/rust-crypto.spec.ts | 5 ++ src/client.ts | 6 ++- src/crypto/index.ts | 11 ++++ src/crypto/store/base.ts | 7 +++ .../store/indexeddb-crypto-store-backend.ts | 19 +++++++ src/crypto/store/indexeddb-crypto-store.ts | 11 ++++ src/crypto/store/localStorage-crypto-store.ts | 24 ++++++++- src/crypto/store/memory-crypto-store.ts | 17 +++++- src/rust-crypto/index.ts | 7 +++ src/rust-crypto/libolm_migration.ts | 54 +++++++++++++++++-- 12 files changed, 185 insertions(+), 9 deletions(-) diff --git a/spec/integ/crypto/rust-crypto.spec.ts b/spec/integ/crypto/rust-crypto.spec.ts index 478702f81fd..506f66e45c5 100644 --- a/spec/integ/crypto/rust-crypto.spec.ts +++ b/spec/integ/crypto/rust-crypto.spec.ts @@ -18,7 +18,7 @@ import "fake-indexeddb/auto"; import { IDBFactory } from "fake-indexeddb"; import fetchMock from "fetch-mock-jest"; -import { createClient, IndexedDBCryptoStore } from "../../../src"; +import { createClient, CryptoEvent, IndexedDBCryptoStore } from "../../../src"; import { populateStore } from "../../test-utils/test_indexeddb_cryptostore_dump"; jest.setTimeout(15000); @@ -124,6 +124,9 @@ describe("MatrixClient.initRustCrypto", () => { pickleKey: "+1k2Ppd7HIisUY824v7JtV3/oEE4yX0TqtmNPyhaD7o", }); + const progressListener = jest.fn(); + matrixClient.addListener(CryptoEvent.LegacyCryptoStoreMigrationProgress, progressListener); + await matrixClient.initRustCrypto(); // Do some basic checks on the imported data @@ -132,7 +135,25 @@ describe("MatrixClient.initRustCrypto", () => { expect(deviceKeys.ed25519).toEqual("qK70DEqIXq7T+UU3v/al47Ab4JkMEBLpNrTBMbS5rrw"); expect(await matrixClient.getCrypto()!.getActiveSessionBackupVersion()).toEqual("7"); - }); + + // check the progress callback + expect(progressListener.mock.calls.length).toBeGreaterThan(50); + + // The first call should have progress == 0 + const [firstProgress, totalSteps] = progressListener.mock.calls[0]; + expect(totalSteps).toBeGreaterThan(3000); + expect(firstProgress).toEqual(0); + + for (let i = 1; i < progressListener.mock.calls.length - 1; i++) { + const [progress, total] = progressListener.mock.calls[i]; + expect(total).toEqual(totalSteps); + expect(progress).toBeGreaterThan(progressListener.mock.calls[i - 1][0]); + expect(progress).toBeLessThanOrEqual(totalSteps); + } + + // The final call should have progress == total == -1 + expect(progressListener).toHaveBeenLastCalledWith(-1, -1); + }, 60000); }); describe("MatrixClient.clearStores", () => { diff --git a/spec/unit/crypto/store/CryptoStore.spec.ts b/spec/unit/crypto/store/CryptoStore.spec.ts index b6b6f652509..9070c2a5c44 100644 --- a/spec/unit/crypto/store/CryptoStore.spec.ts +++ b/spec/unit/crypto/store/CryptoStore.spec.ts @@ -74,6 +74,12 @@ describe.each([ const N_SESSIONS_PER_DEVICE = 6; await createSessions(N_DEVICES, N_SESSIONS_PER_DEVICE); + let nSessions = 0; + await store.doTxn("readonly", [IndexedDBCryptoStore.STORE_SESSIONS], (txn) => + store.countEndToEndSessions(txn, (n) => (nSessions = n)), + ); + expect(nSessions).toEqual(N_DEVICES * N_SESSIONS_PER_DEVICE); + // Then, get a batch and check it looks right. const batch = await store.getEndToEndSessionsBatch(); expect(batch!.length).toEqual(N_DEVICES * N_SESSIONS_PER_DEVICE); @@ -150,6 +156,8 @@ describe.each([ await store.markSessionsNeedingBackup([{ senderKey: pad43("device5"), sessionId: "session5" }], txn); }); + expect(await store.countEndToEndInboundGroupSessions()).toEqual(N_DEVICES * N_SESSIONS_PER_DEVICE); + const batch = await store.getEndToEndInboundGroupSessionsBatch(); expect(batch!.length).toEqual(N_DEVICES * N_SESSIONS_PER_DEVICE); for (let i = 0; i < N_DEVICES; i++) { diff --git a/spec/unit/rust-crypto/rust-crypto.spec.ts b/spec/unit/rust-crypto/rust-crypto.spec.ts index c3d3d33bbbb..28f3c653e4f 100644 --- a/spec/unit/rust-crypto/rust-crypto.spec.ts +++ b/spec/unit/rust-crypto/rust-crypto.spec.ts @@ -193,6 +193,10 @@ describe("initRustCrypto", () => { fetchMock.get("path:/_matrix/client/v3/room_keys/version", { version: "45" }); + function legacyMigrationProgressListener(progress: number, total: number): void { + logger.log(`migrated ${progress} of ${total}`); + } + await initRustCrypto({ logger, http: makeMatrixHttpApi(), @@ -204,6 +208,7 @@ describe("initRustCrypto", () => { storePassphrase: "storePassphrase", legacyCryptoStore: legacyStore, legacyPickleKey: PICKLE_KEY, + legacyMigrationProgressListener, }); // Check that the migration functions were correctly called diff --git a/src/client.ts b/src/client.ts index 01bd7e370d5..f0de776d0a6 100644 --- a/src/client.ts +++ b/src/client.ts @@ -962,7 +962,8 @@ type CryptoEvents = | CryptoEvent.KeysChanged | CryptoEvent.Warning | CryptoEvent.DevicesUpdated - | CryptoEvent.WillUpdateDevices; + | CryptoEvent.WillUpdateDevices + | CryptoEvent.LegacyCryptoStoreMigrationProgress; type MatrixEventEvents = MatrixEventEvent.Decrypted | MatrixEventEvent.Replaced | MatrixEventEvent.VisibilityChange; @@ -2330,6 +2331,9 @@ export class MatrixClient extends TypedEventEmitter { + this.emit(CryptoEvent.LegacyCryptoStoreMigrationProgress, progress, total); + }, }); rustCrypto.setSupportedVerificationMethods(this.verificationMethods); diff --git a/src/crypto/index.ts b/src/crypto/index.ts index c1ef1ac5435..16e37605b94 100644 --- a/src/crypto/index.ts +++ b/src/crypto/index.ts @@ -259,6 +259,15 @@ export enum CryptoEvent { WillUpdateDevices = "crypto.willUpdateDevices", DevicesUpdated = "crypto.devicesUpdated", KeysChanged = "crossSigning.keysChanged", + + /** + * Fires when data is being migrated from legacy crypto to rust crypto. + * + * The payload is a pair `(progress, total)`, where `progress` is the number of steps completed so far, and + * `total` is the total number of steps. When migration is complete, a final instance of the event is emitted, with + * `progress === total === -1`. + */ + LegacyCryptoStoreMigrationProgress = "crypto.legacyCryptoStoreMigrationProgress", } export type CryptoEventHandlerMap = { @@ -368,6 +377,8 @@ export type CryptoEventHandlerMap = { */ [CryptoEvent.DevicesUpdated]: (users: string[], initialFetch: boolean) => void; [CryptoEvent.UserCrossSigningUpdated]: (userId: string) => void; + + [CryptoEvent.LegacyCryptoStoreMigrationProgress]: (progress: number, total: number) => void; }; export class Crypto extends TypedEventEmitter implements CryptoBackend { diff --git a/src/crypto/store/base.ts b/src/crypto/store/base.ts index 0a0a9010ac8..b7003853b3a 100644 --- a/src/crypto/store/base.ts +++ b/src/crypto/store/base.ts @@ -176,6 +176,13 @@ export interface CryptoStore { txn: unknown, ): void; + /** + * Count the number of Megolm sessions in the database. + * + * @internal + */ + countEndToEndInboundGroupSessions(): Promise; + /** * Get a batch of Megolm sessions from the database. * diff --git a/src/crypto/store/indexeddb-crypto-store-backend.ts b/src/crypto/store/indexeddb-crypto-store-backend.ts index 6d2548518dd..454f69e0e3d 100644 --- a/src/crypto/store/indexeddb-crypto-store-backend.ts +++ b/src/crypto/store/indexeddb-crypto-store-backend.ts @@ -811,6 +811,25 @@ export class Backend implements CryptoStore { }); } + /** + * Count the number of Megolm sessions in the database. + * + * Implementation of {@link CryptoStore.countEndToEndInboundGroupSessions}. + * + * @internal + */ + public async countEndToEndInboundGroupSessions(): Promise { + let result = 0; + await this.doTxn("readonly", [IndexedDBCryptoStore.STORE_INBOUND_GROUP_SESSIONS], (txn) => { + const sessionStore = txn.objectStore(IndexedDBCryptoStore.STORE_INBOUND_GROUP_SESSIONS); + const countReq = sessionStore.count(); + countReq.onsuccess = (): void => { + result = countReq.result; + }; + }); + return result; + } + /** * Fetch a batch of Megolm sessions from the database. * diff --git a/src/crypto/store/indexeddb-crypto-store.ts b/src/crypto/store/indexeddb-crypto-store.ts index 8598e8454f3..dc33d5f81d9 100644 --- a/src/crypto/store/indexeddb-crypto-store.ts +++ b/src/crypto/store/indexeddb-crypto-store.ts @@ -506,6 +506,17 @@ export class IndexedDBCryptoStore implements CryptoStore { return this.backend!.filterOutNotifiedErrorDevices(devices); } + /** + * Count the number of Megolm sessions in the database. + * + * Implementation of {@link CryptoStore.countEndToEndInboundGroupSessions}. + * + * @internal + */ + public countEndToEndInboundGroupSessions(): Promise { + return this.backend!.countEndToEndInboundGroupSessions(); + } + /** * Fetch a batch of Olm sessions from the database. * diff --git a/src/crypto/store/localStorage-crypto-store.ts b/src/crypto/store/localStorage-crypto-store.ts index 440c644e24e..9a8dbf38a80 100644 --- a/src/crypto/store/localStorage-crypto-store.ts +++ b/src/crypto/store/localStorage-crypto-store.ts @@ -127,7 +127,11 @@ export class LocalStorageCryptoStore extends MemoryCryptoStore implements Crypto public countEndToEndSessions(txn: unknown, func: (count: number) => void): void { let count = 0; for (let i = 0; i < this.store.length; ++i) { - if (this.store.key(i)?.startsWith(keyEndToEndSessions(""))) ++count; + const key = this.store.key(i); + if (key?.startsWith(keyEndToEndSessions(""))) { + const sessions = getJsonItem(this.store, key); + count += Object.keys(sessions ?? {}).length; + } } func(count); } @@ -351,6 +355,24 @@ export class LocalStorageCryptoStore extends MemoryCryptoStore implements Crypto setJsonItem(this.store, keyEndToEndInboundGroupSessionWithheld(senderCurve25519Key, sessionId), sessionData); } + /** + * Count the number of Megolm sessions in the database. + * + * Implementation of {@link CryptoStore.countEndToEndInboundGroupSessions}. + * + * @internal + */ + public async countEndToEndInboundGroupSessions(): Promise { + let count = 0; + for (let i = 0; i < this.store.length; ++i) { + const key = this.store.key(i); + if (key?.startsWith(KEY_INBOUND_SESSION_PREFIX)) { + count += 1; + } + } + return count; + } + /** * Fetch a batch of Megolm sessions from the database. * diff --git a/src/crypto/store/memory-crypto-store.ts b/src/crypto/store/memory-crypto-store.ts index 2082f955aa1..de3dc3adaf9 100644 --- a/src/crypto/store/memory-crypto-store.ts +++ b/src/crypto/store/memory-crypto-store.ts @@ -336,7 +336,11 @@ export class MemoryCryptoStore implements CryptoStore { // Olm Sessions public countEndToEndSessions(txn: unknown, func: (count: number) => void): void { - func(Object.keys(this.sessions).length); + let count = 0; + for (const deviceSessions of Object.values(this.sessions)) { + count += Object.keys(deviceSessions).length; + } + func(count); } public getEndToEndSession( @@ -528,6 +532,17 @@ export class MemoryCryptoStore implements CryptoStore { this.inboundGroupSessionsWithheld[k] = sessionData; } + /** + * Count the number of Megolm sessions in the database. + * + * Implementation of {@link CryptoStore.countEndToEndInboundGroupSessions}. + * + * @internal + */ + public async countEndToEndInboundGroupSessions(): Promise { + return Object.keys(this.inboundGroupSessions).length; + } + /** * Fetch a batch of Megolm sessions from the database. * diff --git a/src/rust-crypto/index.ts b/src/rust-crypto/index.ts index 93994bebcb9..9108761b7d8 100644 --- a/src/rust-crypto/index.ts +++ b/src/rust-crypto/index.ts @@ -72,6 +72,13 @@ export async function initRustCrypto(args: { /** The pickle key for `legacyCryptoStore` */ legacyPickleKey?: string; + + /** + * A callback which will receive progress updates on migration from `legacyCryptoStore`. + * + * Called with (-1, -1) to mark the end of migration. + */ + legacyMigrationProgressListener?: (progress: number, total: number) => void; }): Promise { const { logger } = args; diff --git a/src/rust-crypto/libolm_migration.ts b/src/rust-crypto/libolm_migration.ts index 9e7605fffdd..e03d401cc93 100644 --- a/src/rust-crypto/libolm_migration.ts +++ b/src/rust-crypto/libolm_migration.ts @@ -52,6 +52,13 @@ export async function migrateFromLegacyCrypto(args: { /** Rust crypto store to migrate data into. */ storeHandle: RustSdkCryptoJs.StoreHandle; + + /** + * A callback which will receive progress updates on migration from `legacyStore`. + * + * Called with (-1, -1) to mark the end of migration. + */ + legacyMigrationProgressListener?: (progress: number, total: number) => void; }): Promise { const { logger, legacyStore } = args; @@ -74,6 +81,20 @@ export async function migrateFromLegacyCrypto(args: { return; } + const nOlmSessions = await countOlmSessions(logger, legacyStore); + const nMegolmSessions = await countMegolmSessions(logger, legacyStore); + const totalSteps = 1 + nOlmSessions + nMegolmSessions; + logger.info( + `Migrating data from legacy crypto store. ${nOlmSessions} olm sessions and ${nMegolmSessions} megolm sessions to migrate.`, + ); + + let stepsDone = 0; + function onProgress(steps: number): void { + stepsDone += steps; + args.legacyMigrationProgressListener?.(stepsDone, totalSteps); + } + onProgress(0); + const pickleKey = new TextEncoder().encode(args.legacyPickleKey); if (migrationState === MigrationState.NOT_STARTED) { @@ -83,23 +104,30 @@ export async function migrateFromLegacyCrypto(args: { migrationState = MigrationState.INITIAL_DATA_MIGRATED; await legacyStore.setMigrationState(migrationState); } + onProgress(1); if (migrationState === MigrationState.INITIAL_DATA_MIGRATED) { - logger.info("Migrating data from legacy crypto store. Step 2: olm sessions"); - await migrateOlmSessions(logger, legacyStore, pickleKey, args.storeHandle); + logger.info( + `Migrating data from legacy crypto store. Step 2: olm sessions (${nOlmSessions} sessions to migrate).`, + ); + await migrateOlmSessions(logger, legacyStore, pickleKey, args.storeHandle, onProgress); migrationState = MigrationState.OLM_SESSIONS_MIGRATED; await legacyStore.setMigrationState(migrationState); } if (migrationState === MigrationState.OLM_SESSIONS_MIGRATED) { - logger.info("Migrating data from legacy crypto store. Step 3: megolm sessions"); - await migrateMegolmSessions(logger, legacyStore, pickleKey, args.storeHandle); + logger.info( + `Migrating data from legacy crypto store. Step 3: megolm sessions (${nMegolmSessions} sessions to migrate).`, + ); + await migrateMegolmSessions(logger, legacyStore, pickleKey, args.storeHandle, onProgress); migrationState = MigrationState.MEGOLM_SESSIONS_MIGRATED; await legacyStore.setMigrationState(migrationState); } + // Migration is done. + args.legacyMigrationProgressListener?.(-1, -1); logger.info("Migration from legacy crypto store complete"); } @@ -147,11 +175,26 @@ async function migrateBaseData( await RustSdkCryptoJs.Migration.migrateBaseData(migrationData, pickleKey, storeHandle); } +async function countOlmSessions(logger: Logger, legacyStore: CryptoStore): Promise { + logger.debug("Counting olm sessions to be migrated"); + let nSessions: number; + await legacyStore.doTxn("readonly", [IndexedDBCryptoStore.STORE_SESSIONS], (txn) => + legacyStore.countEndToEndSessions(txn, (n) => (nSessions = n)), + ); + return nSessions!; +} + +async function countMegolmSessions(logger: Logger, legacyStore: CryptoStore): Promise { + logger.debug("Counting megolm sessions to be migrated"); + return await legacyStore.countEndToEndInboundGroupSessions(); +} + async function migrateOlmSessions( logger: Logger, legacyStore: CryptoStore, pickleKey: Uint8Array, storeHandle: RustSdkCryptoJs.StoreHandle, + onBatchDone: (batchSize: number) => void, ): Promise { // eslint-disable-next-line no-constant-condition while (true) { @@ -170,6 +213,7 @@ async function migrateOlmSessions( await RustSdkCryptoJs.Migration.migrateOlmSessions(migrationData, pickleKey, storeHandle); await legacyStore.deleteEndToEndSessionsBatch(batch); + onBatchDone(batch.length); } } @@ -178,6 +222,7 @@ async function migrateMegolmSessions( legacyStore: CryptoStore, pickleKey: Uint8Array, storeHandle: RustSdkCryptoJs.StoreHandle, + onBatchDone: (batchSize: number) => void, ): Promise { // eslint-disable-next-line no-constant-condition while (true) { @@ -204,6 +249,7 @@ async function migrateMegolmSessions( await RustSdkCryptoJs.Migration.migrateMegolmSessions(migrationData, pickleKey, storeHandle); await legacyStore.deleteEndToEndInboundGroupSessionsBatch(batch); + onBatchDone(batch.length); } }