From a9c6f4738591d2eb136223d9930edd7cf5e83b06 Mon Sep 17 00:00:00 2001
From: Richard van der Hoff <richard@matrix.org>
Date: Thu, 28 Dec 2023 15:05:54 +0000
Subject: [PATCH] Use a `StoreHandle` to init OlmMachine

This will be faster if we need to prepare the store.
---
 spec/unit/rust-crypto/rust-crypto.spec.ts | 30 ++++++++++++++---------
 src/client.ts                             |  1 +
 src/rust-crypto/index.ts                  | 21 ++++++++++------
 3 files changed, 33 insertions(+), 19 deletions(-)

diff --git a/spec/unit/rust-crypto/rust-crypto.spec.ts b/spec/unit/rust-crypto/rust-crypto.spec.ts
index c9146c24677..a2519171717 100644
--- a/spec/unit/rust-crypto/rust-crypto.spec.ts
+++ b/spec/unit/rust-crypto/rust-crypto.spec.ts
@@ -15,7 +15,7 @@ limitations under the License.
 */
 
 import * as RustSdkCryptoJs from "@matrix-org/matrix-sdk-crypto-wasm";
-import { KeysQueryRequest, OlmMachine } from "@matrix-org/matrix-sdk-crypto-wasm";
+import { KeysQueryRequest, OlmMachine, StoreHandle } from "@matrix-org/matrix-sdk-crypto-wasm";
 import { mocked, Mocked } from "jest-mock";
 import fetchMock from "fetch-mock-jest";
 
@@ -81,8 +81,11 @@ describe("initRustCrypto", () => {
     }
 
     it("passes through the store params", async () => {
+        const mockStore = { free: jest.fn() } as unknown as StoreHandle;
+        jest.spyOn(StoreHandle, "open").mockResolvedValue(mockStore);
+
         const testOlmMachine = makeTestOlmMachine();
-        jest.spyOn(OlmMachine, "initialize").mockResolvedValue(testOlmMachine);
+        jest.spyOn(OlmMachine, "init_from_store").mockResolvedValue(testOlmMachine);
 
         await initRustCrypto({
             logger,
@@ -95,17 +98,16 @@ describe("initRustCrypto", () => {
             storePassphrase: "storePassphrase",
         });
 
-        expect(OlmMachine.initialize).toHaveBeenCalledWith(
-            expect.anything(),
-            expect.anything(),
-            "storePrefix",
-            "storePassphrase",
-        );
+        expect(StoreHandle.open).toHaveBeenCalledWith("storePrefix", "storePassphrase");
+        expect(OlmMachine.init_from_store).toHaveBeenCalledWith(expect.anything(), expect.anything(), mockStore);
     });
 
     it("suppresses the storePassphrase if storePrefix is unset", async () => {
+        const mockStore = { free: jest.fn() } as unknown as StoreHandle;
+        jest.spyOn(StoreHandle, "open").mockResolvedValue(mockStore);
+
         const testOlmMachine = makeTestOlmMachine();
-        jest.spyOn(OlmMachine, "initialize").mockResolvedValue(testOlmMachine);
+        jest.spyOn(OlmMachine, "init_from_store").mockResolvedValue(testOlmMachine);
 
         await initRustCrypto({
             logger,
@@ -118,12 +120,16 @@ describe("initRustCrypto", () => {
             storePassphrase: "storePassphrase",
         });
 
-        expect(OlmMachine.initialize).toHaveBeenCalledWith(expect.anything(), expect.anything(), undefined, undefined);
+        expect(StoreHandle.open).toHaveBeenCalledWith(undefined, undefined);
+        expect(OlmMachine.init_from_store).toHaveBeenCalledWith(expect.anything(), expect.anything(), mockStore);
     });
 
     it("Should get secrets from inbox on start", async () => {
-        const testOlmMachine = makeTestOlmMachine() as OlmMachine;
-        jest.spyOn(OlmMachine, "initialize").mockResolvedValue(testOlmMachine);
+        const mockStore = { free: jest.fn() } as unknown as StoreHandle;
+        jest.spyOn(StoreHandle, "open").mockResolvedValue(mockStore);
+
+        const testOlmMachine = makeTestOlmMachine();
+        jest.spyOn(OlmMachine, "init_from_store").mockResolvedValue(testOlmMachine);
 
         await initRustCrypto({
             logger,
diff --git a/src/client.ts b/src/client.ts
index 27e70c43520..8b431c81d80 100644
--- a/src/client.ts
+++ b/src/client.ts
@@ -2329,6 +2329,7 @@ export class MatrixClient extends TypedEventEmitter<EmittedEvents, ClientEventHa
             storePrefix: useIndexedDB ? RUST_SDK_STORE_PREFIX : null,
             storePassphrase: this.pickleKey,
         });
+
         rustCrypto.setSupportedVerificationMethods(this.verificationMethods);
 
         this.cryptoBackend = rustCrypto;
diff --git a/src/rust-crypto/index.ts b/src/rust-crypto/index.ts
index 15778092b51..bc54387e1d6 100644
--- a/src/rust-crypto/index.ts
+++ b/src/rust-crypto/index.ts
@@ -15,6 +15,7 @@ limitations under the License.
 */
 
 import * as RustSdkCryptoJs from "@matrix-org/matrix-sdk-crypto-wasm";
+import { StoreHandle } from "@matrix-org/matrix-sdk-crypto-wasm";
 
 import { RustCrypto } from "./rust-crypto";
 import { IHttpOpts, MatrixHttpApi } from "../http-api";
@@ -73,6 +74,12 @@ export async function initRustCrypto(args: {
     // enable tracing in the rust-sdk
     new RustSdkCryptoJs.Tracing(RustSdkCryptoJs.LoggerLevel.Debug).turnOn();
 
+    logger.debug("Opening Rust CryptoStore");
+    const storeHandle: StoreHandle = await StoreHandle.open(
+        args.storePrefix ?? undefined,
+        (args.storePrefix && args.storePassphrase) ?? undefined,
+    );
+
     const rustCrypto = await initOlmMachine(
         logger,
         args.http,
@@ -80,10 +87,11 @@ export async function initRustCrypto(args: {
         args.deviceId,
         args.secretStorage,
         args.cryptoCallbacks,
-        args.storePrefix,
-        args.storePassphrase,
+        storeHandle,
     );
 
+    storeHandle.free();
+
     logger.debug("Completed rust crypto-sdk setup");
     return rustCrypto;
 }
@@ -95,15 +103,14 @@ async function initOlmMachine(
     deviceId: string,
     secretStorage: ServerSideSecretStorage,
     cryptoCallbacks: ICryptoCallbacks,
-    storePrefix: string | null,
-    storePassphrase: string | undefined,
+    storeHandle: StoreHandle,
 ): Promise<RustCrypto> {
     logger.debug("Init OlmMachine");
-    const olmMachine = await RustSdkCryptoJs.OlmMachine.initialize(
+
+    const olmMachine = await RustSdkCryptoJs.OlmMachine.init_from_store(
         new RustSdkCryptoJs.UserId(userId),
         new RustSdkCryptoJs.DeviceId(deviceId),
-        storePrefix ?? undefined,
-        (storePrefix && storePassphrase) ?? undefined,
+        storeHandle,
     );
 
     // Disable room key requests, per https://github.com/vector-im/element-web/issues/26524.