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;

src/crypto/store/indexeddb-crypto-store-backend.ts

@@ -23,23 +23,31 @@ import {
     ISession,
     ISessionInfo,
     IWithheld,
+    MigrationState,
     Mode,
     OutgoingRoomKeyRequest,
     ParkedSharedHistory,
     SecretStorePrivateKeys,
+    SESSION_BATCH_SIZE,
 } from "./base";
 import { IRoomKeyRequestBody, IRoomKeyRequestRecipient } from "../index";
 import { ICrossSigningKey } from "../../client";
 import { IOlmDevice } from "../algorithms/megolm";
 import { IRoomEncryption } from "../RoomList";
 import { InboundGroupSessionData } from "../OlmDevice";
+import { IndexedDBCryptoStore } from "./indexeddb-crypto-store";
 
 const PROFILE_TRANSACTIONS = false;
 
+/* Keys for the `account` object store */
+const ACCOUNT_OBJECT_KEY_MIGRATION_STATE = "migrationState";
+
 /**
  * Implementation of a CryptoStore which is backed by an existing
  * IndexedDB connection. Generally you want IndexedDBCryptoStore
  * which connects to the database and defers to one of these.
+ *
+ * @internal
  */
 export class Backend implements CryptoStore {
     private nextTxnId = 0;
@@ -56,15 +64,49 @@ export class Backend implements CryptoStore {
         };
     }
 
+    public async containsData(): Promise<boolean> {
+        throw Error("Not implemented for Backend");
+    }
+
     public async startup(): Promise<CryptoStore> {
         // No work to do, as the startup is done by the caller (e.g IndexedDBCryptoStore)
         // by passing us a ready IDBDatabase instance
         return this;
     }
+
     public async deleteAllData(): Promise<void> {
         throw Error("This is not implemented, call IDBFactory::deleteDatabase(dbName) instead.");
     }
 
+    /**
+     * Get data on how much of the libolm to Rust Crypto migration has been done.
+     *
+     * Implementation of {@link CryptoStore.getMigrationState}.
+     */
+    public async getMigrationState(): Promise<MigrationState> {
+        let migrationState = MigrationState.NOT_STARTED;
+        await this.doTxn("readonly", [IndexedDBCryptoStore.STORE_ACCOUNT], (txn) => {
+            const objectStore = txn.objectStore(IndexedDBCryptoStore.STORE_ACCOUNT);
+            const getReq = objectStore.get(ACCOUNT_OBJECT_KEY_MIGRATION_STATE);
+            getReq.onsuccess = (): void => {
+                migrationState = getReq.result ?? MigrationState.NOT_STARTED;
+            };
+        });
+        return migrationState;
+    }
+
+    /**
+     * Set data on how much of the libolm to Rust Crypto migration has been done.
+     *
+     * Implementation of {@link CryptoStore.setMigrationState}.
+     */
+    public async setMigrationState(migrationState: MigrationState): Promise<void> {
+        await this.doTxn("readwrite", [IndexedDBCryptoStore.STORE_ACCOUNT], (txn) => {
+            const objectStore = txn.objectStore(IndexedDBCryptoStore.STORE_ACCOUNT);
+            objectStore.put(migrationState, ACCOUNT_OBJECT_KEY_MIGRATION_STATE);
+        });
+    }
+
     /**
      * Look for an existing outgoing room key request, and if none is found,
      * add a new one
@@ -588,6 +630,62 @@ export class Backend implements CryptoStore {
         return ret;
     }
 
+    /**
+     * Fetch a batch of Olm sessions from the database.
+     *
+     * Implementation of {@link CryptoStore.getEndToEndSessionsBatch}.
+     */
+    public async getEndToEndSessionsBatch(): Promise<null | ISessionInfo[]> {
+        const result: ISessionInfo[] = [];
+        await this.doTxn("readonly", [IndexedDBCryptoStore.STORE_SESSIONS], (txn) => {
+            const objectStore = txn.objectStore(IndexedDBCryptoStore.STORE_SESSIONS);
+            const getReq = objectStore.openCursor();
+            getReq.onsuccess = function (): void {
+                try {
+                    const cursor = getReq.result;
+                    if (cursor) {
+                        result.push(cursor.value);
+                        if (result.length < SESSION_BATCH_SIZE) {
+                            cursor.continue();
+                        }
+                    }
+                } catch (e) {
+                    abortWithException(txn, <Error>e);
+                }
+            };
+        });
+
+        if (result.length === 0) {
+            // No sessions left.
+            return null;
+        }
+
+        return result;
+    }
+
+    /**
+     * Delete a batch of Olm sessions from the database.
+     *
+     * Implementation of {@link CryptoStore.deleteEndToEndSessionsBatch}.
+     *
+     * @internal
+     */
+    public async deleteEndToEndSessionsBatch(sessions: { deviceKey: string; sessionId: string }[]): Promise<void> {
+        await this.doTxn("readwrite", [IndexedDBCryptoStore.STORE_SESSIONS], async (txn) => {
+            try {
+                const objectStore = txn.objectStore(IndexedDBCryptoStore.STORE_SESSIONS);
+                for (const { deviceKey, sessionId } of sessions) {
+                    const req = objectStore.delete([deviceKey, sessionId]);
+                    await new Promise((resolve) => {
+                        req.onsuccess = resolve;
+                    });
+                }
+            } catch (e) {
+                abortWithException(txn, <Error>e);
+            }
+        });
+    }
+
     // Inbound group sessions
 
     public getEndToEndInboundGroupSession(
@@ -712,6 +810,68 @@ export class Backend implements CryptoStore {
         });
     }
 
+    /**
+     * Fetch a batch of Megolm sessions from the database.
+     *
+     * Implementation of {@link CryptoStore.getEndToEndInboundGroupSessionsBatch}.
+     */
+    public async getEndToEndInboundGroupSessionsBatch(): Promise<null | ISession[]> {
+        const result: ISession[] = [];
+        await this.doTxn("readonly", [IndexedDBCryptoStore.STORE_INBOUND_GROUP_SESSIONS], (txn) => {
+            const objectStore = txn.objectStore(IndexedDBCryptoStore.STORE_INBOUND_GROUP_SESSIONS);
+            const getReq = objectStore.openCursor();
+            getReq.onsuccess = function (): void {
+                try {
+                    const cursor = getReq.result;
+                    if (cursor) {
+                        result.push({
+                            senderKey: cursor.value.senderCurve25519Key,
+                            sessionId: cursor.value.sessionId,
+                            sessionData: cursor.value.session,
+                        });
+                        if (result.length < SESSION_BATCH_SIZE) {
+                            cursor.continue();
+                        }
+                    }
+                } catch (e) {
+                    abortWithException(txn, <Error>e);
+                }
+            };
+        });
+
+        if (result.length === 0) {
+            // No sessions left.
+            return null;
+        }
+
+        return result;
+    }
+
+    /**
+     * Delete a batch of Megolm sessions from the database.
+     *
+     * Implementation of {@link CryptoStore.deleteEndToEndInboundGroupSessionsBatch}.
+     *
+     * @internal
+     */
+    public async deleteEndToEndInboundGroupSessionsBatch(
+        sessions: { senderKey: string; sessionId: string }[],
+    ): Promise<void> {
+        await this.doTxn("readwrite", [IndexedDBCryptoStore.STORE_INBOUND_GROUP_SESSIONS], async (txn) => {
+            try {
+                const objectStore = txn.objectStore(IndexedDBCryptoStore.STORE_INBOUND_GROUP_SESSIONS);
+                for (const { senderKey, sessionId } of sessions) {
+                    const req = objectStore.delete([senderKey, sessionId]);
+                    await new Promise((resolve) => {
+                        req.onsuccess = resolve;
+                    });
+                }
+            } catch (e) {
+                abortWithException(txn, <Error>e);
+            }
+        });
+    }
+
     public getEndToEndDeviceData(txn: IDBTransaction, func: (deviceData: IDeviceData | null) => void): void {
         const objectStore = txn.objectStore("device_data");
         const getReq = objectStore.get("-");

src/crypto/store/indexeddb-crypto-store.ts

@@ -27,6 +27,7 @@ import {
     ISession,
     ISessionInfo,
     IWithheld,
+    MigrationState,
     Mode,
     OutgoingRoomKeyRequest,
     ParkedSharedHistory,
@@ -38,7 +39,7 @@ import { IOlmDevice } from "../algorithms/megolm";
 import { IRoomEncryption } from "../RoomList";
 import { InboundGroupSessionData } from "../OlmDevice";
 
-/**
+/*
  * Internal module. indexeddb storage for e2e.
  */
 
@@ -72,6 +73,17 @@ export class IndexedDBCryptoStore implements CryptoStore {
      */
     public constructor(private readonly indexedDB: IDBFactory, private readonly dbName: string) {}
 
+    /**
+     * Returns true if this CryptoStore has ever been initialised (ie, it might contain data).
+     *
+     * Implementation of {@link CryptoStore.containsData}.
+     *
+     * @internal
+     */
+    public async containsData(): Promise<boolean> {
+        return IndexedDBCryptoStore.exists(this.indexedDB, this.dbName);
+    }
+
     /**
      * Ensure the database exists and is up-to-date, or fall back to
      * a local storage or in-memory store.
@@ -197,6 +209,28 @@ export class IndexedDBCryptoStore implements CryptoStore {
         });
     }
 
+    /**
+     * Get data on how much of the libolm to Rust Crypto migration has been done.
+     *
+     * Implementation of {@link CryptoStore.getMigrationState}.
+     *
+     * @internal
+     */
+    public getMigrationState(): Promise<MigrationState> {
+        return this.backend!.getMigrationState();
+    }
+
+    /**
+     * Set data on how much of the libolm to Rust Crypto migration has been done.
+     *
+     * Implementation of {@link CryptoStore.setMigrationState}.
+     *
+     * @internal
+     */
+    public setMigrationState(migrationState: MigrationState): Promise<void> {
+        return this.backend!.setMigrationState(migrationState);
+    }
+
     /**
      * Look for an existing outgoing room key request, and if none is found,
      * add a new one
@@ -468,6 +502,28 @@ export class IndexedDBCryptoStore implements CryptoStore {
         return this.backend!.filterOutNotifiedErrorDevices(devices);
     }
 
+    /**
+     * Fetch a batch of Olm sessions from the database.
+     *
+     * Implementation of {@link CryptoStore.getEndToEndSessionsBatch}.
+     *
+     * @internal
+     */
+    public getEndToEndSessionsBatch(): Promise<null | ISessionInfo[]> {
+        return this.backend!.getEndToEndSessionsBatch();
+    }
+
+    /**
+     * Delete a batch of Olm sessions from the database.
+     *
+     * Implementation of {@link CryptoStore.deleteEndToEndSessionsBatch}.
+     *
+     * @internal
+     */
+    public deleteEndToEndSessionsBatch(sessions: { deviceKey: string; sessionId: string }[]): Promise<void> {
+        return this.backend!.deleteEndToEndSessionsBatch(sessions);
+    }
+
     // Inbound group sessions
 
     /**
@@ -544,6 +600,30 @@ export class IndexedDBCryptoStore implements CryptoStore {
         this.backend!.storeEndToEndInboundGroupSessionWithheld(senderCurve25519Key, sessionId, sessionData, txn);
     }
 
+    /**
+     * Fetch a batch of Megolm sessions from the database.
+     *
+     * Implementation of {@link CryptoStore.getEndToEndInboundGroupSessionsBatch}.
+     *
+     * @internal
+     */
+    public getEndToEndInboundGroupSessionsBatch(): Promise<ISession[] | null> {
+        return this.backend!.getEndToEndInboundGroupSessionsBatch();
+    }
+
+    /**
+     * Delete a batch of Megolm sessions from the database.
+     *
+     * Implementation of {@link CryptoStore.deleteEndToEndInboundGroupSessionsBatch}.
+     *
+     * @internal
+     */
+    public deleteEndToEndInboundGroupSessionsBatch(
+        sessions: { senderKey: string; sessionId: string }[],
+    ): Promise<void> {
+        return this.backend!.deleteEndToEndInboundGroupSessionsBatch(sessions);
+    }
+
     // End-to-end device tracking
 
     /**

src/crypto/store/localStorage-crypto-store.ts

@@ -16,7 +16,18 @@ limitations under the License.
 
 import { logger } from "../../logger";
 import { MemoryCryptoStore } from "./memory-crypto-store";
-import { IDeviceData, IProblem, ISession, ISessionInfo, IWithheld, Mode, SecretStorePrivateKeys } from "./base";
+import {
+    CryptoStore,
+    IDeviceData,
+    IProblem,
+    ISession,
+    ISessionInfo,
+    IWithheld,
+    MigrationState,
+    Mode,
+    SecretStorePrivateKeys,
+    SESSION_BATCH_SIZE,
+} from "./base";
 import { IOlmDevice } from "../algorithms/megolm";
 import { IRoomEncryption } from "../RoomList";
 import { ICrossSigningKey } from "../../client";
@@ -32,6 +43,7 @@ import { safeSet } from "../../utils";
  */
 
 const E2E_PREFIX = "crypto.";
+const KEY_END_TO_END_MIGRATION_STATE = E2E_PREFIX + "migration";
 const KEY_END_TO_END_ACCOUNT = E2E_PREFIX + "account";
 const KEY_CROSS_SIGNING_KEYS = E2E_PREFIX + "cross_signing_keys";
 const KEY_NOTIFIED_ERROR_DEVICES = E2E_PREFIX + "notified_error_devices";
@@ -61,7 +73,7 @@ function keyEndToEndRoomsPrefix(roomId: string): string {
     return KEY_ROOMS_PREFIX + roomId;
 }
 
-export class LocalStorageCryptoStore extends MemoryCryptoStore {
+export class LocalStorageCryptoStore extends MemoryCryptoStore implements CryptoStore {
     public static exists(store: Storage): boolean {
         const length = store.length;
         for (let i = 0; i < length; i++) {
@@ -76,6 +88,39 @@ export class LocalStorageCryptoStore extends MemoryCryptoStore {
         super();
     }
 
+    /**
+     * Returns true if this CryptoStore has ever been initialised (ie, it might contain data).
+     *
+     * Implementation of {@link CryptoStore.containsData}.
+     *
+     * @internal
+     */
+    public async containsData(): Promise<boolean> {
+        return LocalStorageCryptoStore.exists(this.store);
+    }
+
+    /**
+     * Get data on how much of the libolm to Rust Crypto migration has been done.
+     *
+     * Implementation of {@link CryptoStore.getMigrationState}.
+     *
+     * @internal
+     */
+    public async getMigrationState(): Promise<MigrationState> {
+        return getJsonItem(this.store, KEY_END_TO_END_MIGRATION_STATE) ?? MigrationState.NOT_STARTED;
+    }
+
+    /**
+     * Set data on how much of the libolm to Rust Crypto migration has been done.
+     *
+     * Implementation of {@link CryptoStore.setMigrationState}.
+     *
+     * @internal
+     */
+    public async setMigrationState(migrationState: MigrationState): Promise<void> {
+        setJsonItem(this.store, KEY_END_TO_END_MIGRATION_STATE, migrationState);
+    }
+
     // Olm Sessions
 
     public countEndToEndSessions(txn: unknown, func: (count: number) => void): void {
@@ -192,6 +237,56 @@ export class LocalStorageCryptoStore extends MemoryCryptoStore {
         return ret;
     }
 
+    /**
+     * Fetch a batch of Olm sessions from the database.
+     *
+     * Implementation of {@link CryptoStore.getEndToEndSessionsBatch}.
+     *
+     * @internal
+     */
+    public async getEndToEndSessionsBatch(): Promise<null | ISessionInfo[]> {
+        const result: ISessionInfo[] = [];
+        for (let i = 0; i < this.store.length; ++i) {
+            if (this.store.key(i)?.startsWith(keyEndToEndSessions(""))) {
+                const deviceKey = this.store.key(i)!.split("/")[1];
+                for (const session of Object.values(this._getEndToEndSessions(deviceKey))) {
+                    result.push(session);
+                    if (result.length >= SESSION_BATCH_SIZE) {
+                        return result;
+                    }
+                }
+            }
+        }
+
+        if (result.length === 0) {
+            // No sessions left.
+            return null;
+        }
+
+        // There are fewer sessions than the batch size; return the final batch of sessions.
+        return result;
+    }
+
+    /**
+     * Delete a batch of Olm sessions from the database.
+     *
+     * Implementation of {@link CryptoStore.deleteEndToEndSessionsBatch}.
+     *
+     * @internal
+     */
+    public async deleteEndToEndSessionsBatch(sessions: { deviceKey: string; sessionId: string }[]): Promise<void> {
+        for (const { deviceKey, sessionId } of sessions) {
+            const deviceSessions = this._getEndToEndSessions(deviceKey) || {};
+            delete deviceSessions[sessionId];
+            if (Object.keys(deviceSessions).length === 0) {
+                // No more sessions for this device.
+                this.store.removeItem(keyEndToEndSessions(deviceKey));
+            } else {
+                setJsonItem(this.store, keyEndToEndSessions(deviceKey), deviceSessions);
+            }
+        }
+    }
+
     // Inbound Group Sessions
 
     public getEndToEndInboundGroupSession(
@@ -255,6 +350,60 @@ export class LocalStorageCryptoStore extends MemoryCryptoStore {
         setJsonItem(this.store, keyEndToEndInboundGroupSessionWithheld(senderCurve25519Key, sessionId), sessionData);
     }
 
+    /**
+     * Fetch a batch of Megolm sessions from the database.
+     *
+     * Implementation of {@link CryptoStore.getEndToEndInboundGroupSessionsBatch}.
+     *
+     * @internal
+     */
+    public async getEndToEndInboundGroupSessionsBatch(): Promise<ISession[] | null> {
+        const result: ISession[] = [];
+        for (let i = 0; i < this.store.length; ++i) {
+            const key = this.store.key(i);
+            if (key?.startsWith(KEY_INBOUND_SESSION_PREFIX)) {
+                // we can't use split, as the components we are trying to split out
+                // might themselves contain '/' characters. We rely on the
+                // senderKey being a (32-byte) curve25519 key, base64-encoded
+                // (hence 43 characters long).
+
+                result.push({
+                    senderKey: key.slice(KEY_INBOUND_SESSION_PREFIX.length, KEY_INBOUND_SESSION_PREFIX.length + 43),
+                    sessionId: key.slice(KEY_INBOUND_SESSION_PREFIX.length + 44),
+                    sessionData: getJsonItem(this.store, key)!,
+                });
+
+                if (result.length >= SESSION_BATCH_SIZE) {
+                    return result;
+                }
+            }
+        }
+
+        if (result.length === 0) {
+            // No sessions left.
+            return null;
+        }
+
+        // There are fewer sessions than the batch size; return the final batch of sessions.
+        return result;
+    }
+
+    /**
+     * Delete a batch of Megolm sessions from the database.
+     *
+     * Implementation of {@link CryptoStore.deleteEndToEndInboundGroupSessionsBatch}.
+     *
+     * @internal
+     */
+    public async deleteEndToEndInboundGroupSessionsBatch(
+        sessions: { senderKey: string; sessionId: string }[],
+    ): Promise<void> {
+        for (const { senderKey, sessionId } of sessions) {
+            const k = keyEndToEndInboundGroupSession(senderKey, sessionId);
+            this.store.removeItem(k);
+        }
+    }
+
     public getEndToEndDeviceData(txn: unknown, func: (deviceData: IDeviceData | null) => void): void {
         func(getJsonItem(this.store, KEY_DEVICE_DATA));
     }

src/crypto/store/memory-crypto-store.ts

@@ -15,7 +15,7 @@ limitations under the License.
 */
 
 import { logger } from "../../logger";
-import { safeSet, deepCompare, promiseTry } from "../../utils";
+import { deepCompare, promiseTry, safeSet } from "../../utils";
 import {
     CryptoStore,
     IDeviceData,
@@ -23,10 +23,12 @@ import {
     ISession,
     ISessionInfo,
     IWithheld,
+    MigrationState,
     Mode,
     OutgoingRoomKeyRequest,
     ParkedSharedHistory,
     SecretStorePrivateKeys,
+    SESSION_BATCH_SIZE,
 } from "./base";
 import { IRoomKeyRequestBody } from "../index";
 import { ICrossSigningKey } from "../../client";
@@ -39,6 +41,7 @@ import { InboundGroupSessionData } from "../OlmDevice";
  */
 
 export class MemoryCryptoStore implements CryptoStore {
+    private migrationState: MigrationState = MigrationState.NOT_STARTED;
     private outgoingRoomKeyRequests: OutgoingRoomKeyRequest[] = [];
     private account: string | null = null;
     private crossSigningKeys: Record<string, ICrossSigningKey> | null = null;
@@ -56,6 +59,18 @@ export class MemoryCryptoStore implements CryptoStore {
     private sharedHistoryInboundGroupSessions: { [roomId: string]: [senderKey: string, sessionId: string][] } = {};
     private parkedSharedHistory = new Map<string, ParkedSharedHistory[]>(); // keyed by room ID
 
+    /**
+     * Returns true if this CryptoStore has ever been initialised (ie, it might contain data).
+     *
+     * Implementation of {@link CryptoStore.containsData}.
+     *
+     * @internal
+     */
+    public async containsData(): Promise<boolean> {
+        // If it contains anything, it should contain an account.
+        return this.account !== null;
+    }
+
     /**
      * Ensure the database exists and is up-to-date.
      *
@@ -77,6 +92,28 @@ export class MemoryCryptoStore implements CryptoStore {
         return Promise.resolve();
     }
 
+    /**
+     * Get data on how much of the libolm to Rust Crypto migration has been done.
+     *
+     * Implementation of {@link CryptoStore.getMigrationState}.
+     *
+     * @internal
+     */
+    public async getMigrationState(): Promise<MigrationState> {
+        return this.migrationState;
+    }
+
+    /**
+     * Set data on how much of the libolm to Rust Crypto migration has been done.
+     *
+     * Implementation of {@link CryptoStore.setMigrationState}.
+     *
+     * @internal
+     */
+    public async setMigrationState(migrationState: MigrationState): Promise<void> {
+        this.migrationState = migrationState;
+    }
+
     /**
      * Look for an existing outgoing room key request, and if none is found,
      * add a new one
@@ -386,6 +423,51 @@ export class MemoryCryptoStore implements CryptoStore {
         return ret;
     }
 
+    /**
+     * Fetch a batch of Olm sessions from the database.
+     *
+     * Implementation of {@link CryptoStore.getEndToEndSessionsBatch}.
+     *
+     * @internal
+     */
+    public async getEndToEndSessionsBatch(): Promise<null | ISessionInfo[]> {
+        const result: ISessionInfo[] = [];
+        for (const deviceSessions of Object.values(this.sessions)) {
+            for (const session of Object.values(deviceSessions)) {
+                result.push(session);
+                if (result.length >= SESSION_BATCH_SIZE) {
+                    return result;
+                }
+            }
+        }
+
+        if (result.length === 0) {
+            // No sessions left.
+            return null;
+        }
+
+        // There are fewer sessions than the batch size; return the final batch of sessions.
+        return result;
+    }
+
+    /**
+     * Delete a batch of Olm sessions from the database.
+     *
+     * Implementation of {@link CryptoStore.deleteEndToEndSessionsBatch}.
+     *
+     * @internal
+     */
+    public async deleteEndToEndSessionsBatch(sessions: { deviceKey: string; sessionId: string }[]): Promise<void> {
+        for (const { deviceKey, sessionId } of sessions) {
+            const deviceSessions = this.sessions[deviceKey] || {};
+            delete deviceSessions[sessionId];
+            if (Object.keys(deviceSessions).length === 0) {
+                // No more sessions for this device.
+                delete this.sessions[deviceKey];
+            }
+        }
+    }
+
     // Inbound Group Sessions
 
     public getEndToEndInboundGroupSession(
@@ -445,6 +527,51 @@ export class MemoryCryptoStore implements CryptoStore {
         this.inboundGroupSessionsWithheld[k] = sessionData;
     }
 
+    /**
+     * Fetch a batch of Megolm sessions from the database.
+     *
+     * Implementation of {@link CryptoStore.getEndToEndInboundGroupSessionsBatch}.
+     *
+     * @internal
+     */
+    public async getEndToEndInboundGroupSessionsBatch(): Promise<null | ISession[]> {
+        const result: ISession[] = [];
+        for (const [key, session] of Object.entries(this.inboundGroupSessions)) {
+            result.push({
+                senderKey: key.slice(0, 43),
+                sessionId: key.slice(44),
+                sessionData: session,
+            });
+            if (result.length >= SESSION_BATCH_SIZE) {
+                return result;
+            }
+        }
+
+        if (result.length === 0) {
+            // No sessions left.
+            return null;
+        }
+
+        // There are fewer sessions than the batch size; return the final batch of sessions.
+        return result;
+    }
+
+    /**
+     * Delete a batch of Megolm sessions from the database.
+     *
+     * Implementation of {@link CryptoStore.deleteEndToEndInboundGroupSessionsBatch}.
+     *
+     * @internal
+     */
+    public async deleteEndToEndInboundGroupSessionsBatch(
+        sessions: { senderKey: string; sessionId: string }[],
+    ): Promise<void> {
+        for (const { senderKey, sessionId } of sessions) {
+            const k = senderKey + "/" + sessionId;
+            delete this.inboundGroupSessions[k];
+        }
+    }
+
     // Device Data
 
     public getEndToEndDeviceData(txn: unknown, func: (deviceData: IDeviceData | null) => void): void {