Use a StoreHandle to init OlmMachine

This will be faster if we need to prepare the store.
matrix-org · Jan 8, 2024 · a9c6f47 · a9c6f47
1 parent 585cba8
commit a9c6f47
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 19 deletions.
diff --git 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
@@ -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
@@ -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,17 +74,24 @@ 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,
         args.userId,
         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.