Add new methods to CryptoStore for Rust Crypto migration (#3969)

* Add `CryptoStore.containsData`

* add `CryptoStore.{get,set}MigrationState`

* Implement `CryptoStore.getEndToEnd{,InboundGroup}SessionsBatch`

* Implement `CryptoStore.deleteEndToEnd{,InboundGroup}SessionsBatch`

* fix typedoc errors

spec/unit/crypto/store/CryptoStore.spec.ts

@@ -0,0 +1,210 @@
+/*
+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 "fake-indexeddb/auto";
+import "jest-localstorage-mock";
+import { IndexedDBCryptoStore, LocalStorageCryptoStore, MemoryCryptoStore } from "../../../../src";
+import { CryptoStore, MigrationState, SESSION_BATCH_SIZE } from "../../../../src/crypto/store/base";
+
+describe.each([
+    ["IndexedDBCryptoStore", () => new IndexedDBCryptoStore(global.indexedDB, "tests")],
+    ["LocalStorageCryptoStore", () => new LocalStorageCryptoStore(localStorage)],
+    ["MemoryCryptoStore", () => new MemoryCryptoStore()],
+])("CryptoStore tests for %s", function (name, dbFactory) {
+    let store: CryptoStore;
+
+    beforeEach(async () => {
+        store = dbFactory();
+    });
+
+    describe("containsData", () => {
+        it("returns false at first", async () => {
+            expect(await store.containsData()).toBe(false);
+        });
+
+        it("returns true after startup and account setup", async () => {
+            await store.startup();
+            await store.doTxn("readwrite", [IndexedDBCryptoStore.STORE_ACCOUNT], (txn) => {
+                store.storeAccount(txn, "not a real account");
+            });
+            expect(await store.containsData()).toBe(true);
+        });
+    });
+
+    describe("migrationState", () => {
+        beforeEach(async () => {
+            await store.startup();
+        });
+
+        it("returns 0 at first", async () => {
+            expect(await store.getMigrationState()).toEqual(MigrationState.NOT_STARTED);
+        });
+
+        it("stores updates", async () => {
+            await store.setMigrationState(MigrationState.INITIAL_DATA_MIGRATED);
+            expect(await store.getMigrationState()).toEqual(MigrationState.INITIAL_DATA_MIGRATED);
+        });
+    });
+
+    describe("get/delete EndToEndSessionsBatch", () => {
+        beforeEach(async () => {
+            await store.startup();
+        });
+
+        it("returns null at first", async () => {
+            expect(await store.getEndToEndSessionsBatch()).toBe(null);
+        });
+
+        it("returns a batch of sessions", async () => {
+            // First store some sessions in the db
+            const N_DEVICES = 6;
+            const N_SESSIONS_PER_DEVICE = 6;
+            await createSessions(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);
+            for (let i = 0; i < N_DEVICES; i++) {
+                for (let j = 0; j < N_SESSIONS_PER_DEVICE; j++) {
+                    const r = batch![i * N_DEVICES + j];
+
+                    expect(r.deviceKey).toEqual(`device${i}`);
+                    expect(r.sessionId).toEqual(`session${j}`);
+                }
+            }
+        });
+
+        it("returns another batch of sessions after the first batch is deleted", async () => {
+            // First store some sessions in the db
+            const N_DEVICES = 8;
+            const N_SESSIONS_PER_DEVICE = 8;
+            await createSessions(N_DEVICES, N_SESSIONS_PER_DEVICE);
+
+            // Get the first batch
+            const batch = (await store.getEndToEndSessionsBatch())!;
+            expect(batch.length).toEqual(SESSION_BATCH_SIZE);
+
+            // ... and delete.
+            await store.deleteEndToEndSessionsBatch(batch);
+
+            // Fetch a second batch
+            const batch2 = (await store.getEndToEndSessionsBatch())!;
+            expect(batch2.length).toEqual(N_DEVICES * N_SESSIONS_PER_DEVICE - SESSION_BATCH_SIZE);
+
+            // ... and delete.
+            await store.deleteEndToEndSessionsBatch(batch2);
+
+            // the batch should now be null.
+            expect(await store.getEndToEndSessionsBatch()).toBe(null);
+        });
+
+        /** Create a bunch of fake Olm sessions and stash them in the DB. */
+        async function createSessions(nDevices: number, nSessionsPerDevice: number) {
+            await store.doTxn("readwrite", IndexedDBCryptoStore.STORE_SESSIONS, (txn) => {
+                for (let i = 0; i < nDevices; i++) {
+                    for (let j = 0; j < nSessionsPerDevice; j++) {
+                        store.storeEndToEndSession(
+                            `device${i}`,
+                            `session${j}`,
+                            {
+                                deviceKey: `device${i}`,
+                                sessionId: `session${j}`,
+                            },
+                            txn,
+                        );
+                    }
+                }
+            });
+        }
+    });
+
+    describe("get/delete EndToEndInboundGroupSessionsBatch", () => {
+        beforeEach(async () => {
+            await store.startup();
+        });
+
+        it("returns null at first", async () => {
+            expect(await store.getEndToEndInboundGroupSessionsBatch()).toBe(null);
+        });
+
+        it("returns a batch of sessions", async () => {
+            const N_DEVICES = 6;
+            const N_SESSIONS_PER_DEVICE = 6;
+            await createSessions(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++) {
+                for (let j = 0; j < N_SESSIONS_PER_DEVICE; j++) {
+                    const r = batch![i * N_DEVICES + j];
+
+                    expect(r.senderKey).toEqual(pad43(`device${i}`));
+                    expect(r.sessionId).toEqual(`session${j}`);
+                }
+            }
+        });
+
+        it("returns another batch of sessions after the first batch is deleted", async () => {
+            // First store some sessions in the db
+            const N_DEVICES = 8;
+            const N_SESSIONS_PER_DEVICE = 8;
+            await createSessions(N_DEVICES, N_SESSIONS_PER_DEVICE);
+
+            // Get the first batch
+            const batch = (await store.getEndToEndInboundGroupSessionsBatch())!;
+            expect(batch.length).toEqual(SESSION_BATCH_SIZE);
+
+            // ... and delete.
+            await store.deleteEndToEndInboundGroupSessionsBatch(batch);
+
+            // Fetch a second batch
+            const batch2 = (await store.getEndToEndInboundGroupSessionsBatch())!;
+            expect(batch2.length).toEqual(N_DEVICES * N_SESSIONS_PER_DEVICE - SESSION_BATCH_SIZE);
+
+            // ... and delete.
+            await store.deleteEndToEndInboundGroupSessionsBatch(batch2);
+
+            // the batch should now be null.
+            expect(await store.getEndToEndInboundGroupSessionsBatch()).toBe(null);
+        });
+
+        /** Create a bunch of fake megolm sessions and stash them in the DB. */
+        async function createSessions(nDevices: number, nSessionsPerDevice: number) {
+            await store.doTxn("readwrite", IndexedDBCryptoStore.STORE_INBOUND_GROUP_SESSIONS, (txn) => {
+                for (let i = 0; i < nDevices; i++) {
+                    for (let j = 0; j < nSessionsPerDevice; j++) {
+                        store.storeEndToEndInboundGroupSession(
+                            pad43(`device${i}`),
+                            `session${j}`,
+                            {
+                                forwardingCurve25519KeyChain: [],
+                                keysClaimed: {},
+                                room_id: "",
+                                session: "",
+                            },
+                            txn,
+                        );
+                    }
+                }
+            });
+        }
+    });
+});
+
+/** Pad a string to 43 characters long */
+function pad43(x: string): string {
+    return x + ".".repeat(43 - x.length);
+}

src/crypto/store/base.ts

@@ -46,8 +46,41 @@ export interface SecretStorePrivateKeys {
  * Abstraction of things that can store data required for end-to-end encryption
  */
 export interface CryptoStore {
+    /**
+     * Returns true if this CryptoStore has ever been initialised (ie, it might contain data).
+     *
+     * Unlike the rest of the methods in this interface, can be called before {@link CryptoStore#startup}.
+     *
+     * @internal
+     */
+    containsData(): Promise<boolean>;
+
+    /**
+     * Initialise this crypto store.
+     *
+     * Typically, this involves provisioning storage, and migrating any existing data to the current version of the
+     * storage schema where appropriate.
+     *
+     * Must be called before any of the rest of the methods in this interface.
+     */
     startup(): Promise<CryptoStore>;
+
     deleteAllData(): Promise<void>;
+
+    /**
+     * Get data on how much of the libolm to Rust Crypto migration has been done.
+     *
+     * @internal
+     */
+    getMigrationState(): Promise<MigrationState>;
+
+    /**
+     * Set data on how much of the libolm to Rust Crypto migration has been done.
+     *
+     * @internal
+     */
+    setMigrationState(migrationState: MigrationState): Promise<void>;
+
     getOrAddOutgoingRoomKeyRequest(request: OutgoingRoomKeyRequest): Promise<OutgoingRoomKeyRequest>;
     getOutgoingRoomKeyRequest(requestBody: IRoomKeyRequestBody): Promise<OutgoingRoomKeyRequest | null>;
     getOutgoingRoomKeyRequestByState(wantedStates: number[]): Promise<OutgoingRoomKeyRequest | null>;
@@ -99,6 +132,23 @@ export interface CryptoStore {
     getEndToEndSessionProblem(deviceKey: string, timestamp: number): Promise<IProblem | null>;
     filterOutNotifiedErrorDevices(devices: IOlmDevice[]): Promise<IOlmDevice[]>;
 
+    /**
+     * Get a batch of end-to-end sessions from the database.
+     *
+     * @returns A batch of Olm Sessions, or `null` if no sessions are left.
+     * @internal
+     */
+    getEndToEndSessionsBatch(): Promise<ISessionInfo[] | null>;
+
+    /**
+     * Delete a batch of end-to-end sessions from the database.
+     *
+     * Any sessions in the list which are not found are silently ignored.
+     *
+     * @internal
+     */
+    deleteEndToEndSessionsBatch(sessions: { deviceKey?: string; sessionId?: string }[]): Promise<void>;
+
     // Inbound Group Sessions
     getEndToEndInboundGroupSession(
         senderCurve25519Key: string,
@@ -126,6 +176,23 @@ export interface CryptoStore {
         txn: unknown,
     ): void;
 
+    /**
+     * Get a batch of Megolm sessions from the database.
+     *
+     * @returns A batch of Megolm Sessions, or `null` if no sessions are left.
+     * @internal
+     */
+    getEndToEndInboundGroupSessionsBatch(): Promise<ISession[] | null>;
+
+    /**
+     * Delete a batch of Megolm sessions from the database.
+     *
+     * Any sessions in the list which are not found are silently ignored.
+     *
+     * @internal
+     */
+    deleteEndToEndInboundGroupSessionsBatch(sessions: { senderKey: string; sessionId: string }[]): Promise<void>;
+
     // Device Data
     getEndToEndDeviceData(txn: unknown, func: (deviceData: IDeviceData | null) => void): void;
     storeEndToEndDeviceData(deviceData: IDeviceData, txn: unknown): void;
@@ -149,12 +216,14 @@ export interface CryptoStore {
 
 export type Mode = "readonly" | "readwrite";
 
+/** Data on a Megolm session */
 export interface ISession {
     senderKey: string;
     sessionId: string;
     sessionData?: InboundGroupSessionData;
 }
 
+/** Data on an Olm session */
 export interface ISessionInfo {
     deviceKey?: string;
     sessionId?: string;
@@ -224,3 +293,30 @@ export interface ParkedSharedHistory {
     keysClaimed: ReturnType<MatrixEvent["getKeysClaimed"]>; // XXX: Less type dependence on MatrixEvent
     forwardingCurve25519KeyChain: string[];
 }
+
+/**
+ * A record of which steps have been completed in the libolm to Rust Crypto migration.
+ *
+ * Used by {@link CryptoStore#getMigrationState} and {@link CryptoStore#setMigrationState}.
+ *
+ * @internal
+ */
+export enum MigrationState {
+    /** No migration steps have yet been completed. */
+    NOT_STARTED,
+
+    /** We have migrated the account data, cross-signing keys, etc. */
+    INITIAL_DATA_MIGRATED,
+
+    /** INITIAL_DATA_MIGRATED, and in addition, we have migrated all the Olm sessions. */
+    OLM_SESSIONS_MIGRATED,
+
+    /** OLM_SESSIONS_MIGRATED, and in addition, we have migrated all the Megolm sessions. */
+    MEGOLM_SESSIONS_MIGRATED,
+}
+
+/**
+ * The size of batches to be returned by {@link CryptoStore#getEndToEndSessionsBatch} and
+ * {@link CryptoStore#getEndToEndInboundGroupSessionsBatch}.
+ */
+export const SESSION_BATCH_SIZE = 50;