Element-R: pass pickleKey in as raw key for indexeddb encryption

src/Lifecycle.ts

@@ -18,12 +18,12 @@ limitations under the License.
 */
 
 import { ReactNode } from "react";
-import { createClient, MatrixClient, SSOAction, OidcTokenRefresher } from "matrix-js-sdk/src/matrix";
+import { createClient, MatrixClient, SSOAction, OidcTokenRefresher, decodeBase64 } from "matrix-js-sdk/src/matrix";
 import { IEncryptedPayload } from "matrix-js-sdk/src/crypto/aes";
 import { QueryDict } from "matrix-js-sdk/src/utils";
 import { logger } from "matrix-js-sdk/src/logger";
 
-import { IMatrixClientCreds, MatrixClientPeg } from "./MatrixClientPeg";
+import { IMatrixClientCreds, MatrixClientPeg, MatrixClientPegAssignOpts } from "./MatrixClientPeg";
 import { ModuleRunner } from "./modules/ModuleRunner";
 import EventIndexPeg from "./indexing/EventIndexPeg";
 import createMatrixClient from "./utils/createMatrixClient";
@@ -422,6 +422,7 @@ async function onSuccessfulDelegatedAuthLogin(credentials: IMatrixClientCreds):
 }
 
 type TryAgainFunction = () => void;
+
 /**
  * Display a friendly error to the user when token login or OIDC authorization fails
  * @param description error description
@@ -821,7 +822,23 @@ async function doSetLoggedIn(credentials: IMatrixClientCreds, clearStorageEnable
     checkSessionLock();
 
     dis.fire(Action.OnLoggedIn);
-    await startMatrixClient(client, /*startSyncing=*/ !softLogout);
+
+    const clientPegOpts: MatrixClientPegAssignOpts = {};
+    if (credentials.pickleKey) {
+        // The pickleKey, if provided, is probably a base64-encoded 256-bit key, so can be used for the crypto store.
+        if (credentials.pickleKey.length == 43) {
+            clientPegOpts.rustCryptoStoreKey = decodeBase64(credentials.pickleKey);
+        } else {
+            // We have some legacy pickle key. Continue using it as a password.
+            clientPegOpts.rustCryptoStorePassword = credentials.pickleKey;
+        }
+    }
+
+    try {
+        await startMatrixClient(client, /*startSyncing=*/ !softLogout, clientPegOpts);
+    } finally {
+        clientPegOpts.rustCryptoStoreKey?.fill(0);
+    }
 
     return client;
 }
@@ -955,11 +972,16 @@ export function isLoggingOut(): boolean {
 /**
  * Starts the matrix client and all other react-sdk services that
  * listen for events while a session is logged in.
+ *
  * @param client the matrix client to start
- * @param {boolean} startSyncing True (default) to actually start
- * syncing the client.
+ * @param startSyncing - `true` to actually start syncing the client.
+ * @param clientPegOpts - Options to pass through to {@link MatrixClientPeg.start}.
  */
-async function startMatrixClient(client: MatrixClient, startSyncing = true): Promise<void> {
+async function startMatrixClient(
+    client: MatrixClient,
+    startSyncing: boolean,
+    clientPegOpts: MatrixClientPegAssignOpts,
+): Promise<void> {
     logger.log(`Lifecycle: Starting MatrixClient`);
 
     // dispatch this before starting the matrix client: it's used
@@ -990,10 +1012,10 @@ async function startMatrixClient(client: MatrixClient, startSyncing = true): Pro
         // index (e.g. the FilePanel), therefore initialize the event index
         // before the client.
         await EventIndexPeg.init();
-        await MatrixClientPeg.start();
+        await MatrixClientPeg.start(clientPegOpts);
     } else {
         logger.warn("Caller requested only auxiliary services be started");
-        await MatrixClientPeg.assign();
+        await MatrixClientPeg.assign(clientPegOpts);
     }
 
     checkSessionLock();

src/MatrixClientPeg.ts

@@ -66,6 +66,27 @@ export interface IMatrixClientCreds {
     freshLogin?: boolean;
 }
 
+export interface MatrixClientPegAssignOpts {
+    /**
+     * If we are using Rust crypto, a key with which to encrypt the indexeddb.
+     *
+     * If provided, it must be exactly 32 bytes of data. If both this and
+     * {@link MatrixClientPegAssignOpts.rustCryptoStorePassword} are undefined,
+     * the store will be unencrypted.
+     */
+    rustCryptoStoreKey?: Uint8Array;
+
+    /**
+     * If we are using Rust crypto, a password which will be used to derive a key to encrypt the store with.
+     *
+     * An alternative to {@link MatrixClientPegAssignOpts.rustCryptoStoreKey}. Ignored if `rustCryptoStoreKey` is set.
+     *
+     * Deriving a key from a password is (deliberately) a slow operation, so prefer to pass a `rustCryptoStoreKey`
+     * directly where possible.
+     */
+    rustCryptoStorePassword?: string;
+}
+
 /**
  * Holds the current instance of the `MatrixClient` to use across the codebase.
  * Looking for an `MatrixClient`? Just look for the `MatrixClientPeg` on the peg
@@ -103,14 +124,14 @@ export interface IMatrixClientPeg {
     unset(): void;
 
     /**
-     * Prepare the MatrixClient for use, including initialising the store and crypto, but do not start it
+     * Prepare the MatrixClient for use, including initialising the store and crypto, but do not start it.
      */
-    assign(): Promise<IStartClientOpts>;
+    assign(opts?: MatrixClientPegAssignOpts): Promise<IStartClientOpts>;
 
     /**
-     * Prepare the MatrixClient for use, including initialising the store and crypto, and start it
+     * Prepare the MatrixClient for use, including initialising the store and crypto, and start it.
      */
-    start(): Promise<void>;
+    start(opts?: MatrixClientPegAssignOpts): Promise<void>;
 
     /**
      * If we've registered a user ID we set this to the ID of the
@@ -257,7 +278,10 @@ class MatrixClientPegClass implements IMatrixClientPeg {
         PlatformPeg.get()?.reload();
     };
 
-    public async assign(): Promise<IStartClientOpts> {
+    /**
+     * Implementation of {@link IMatrixClientPeg.assign}.
+     */
+    public async assign(assignOpts: MatrixClientPegAssignOpts = {}): Promise<IStartClientOpts> {
         if (!this.matrixClient) {
             throw new Error("createClient must be called first");
         }
@@ -284,7 +308,7 @@ class MatrixClientPegClass implements IMatrixClientPeg {
 
         // try to initialise e2e on the new client
         if (!SettingsStore.getValue("lowBandwidth")) {
-            await this.initClientCrypto();
+            await this.initClientCrypto(assignOpts.rustCryptoStoreKey, assignOpts.rustCryptoStorePassword);
         }
 
         const opts = utils.deepCopy(this.opts);
@@ -310,8 +334,16 @@ class MatrixClientPegClass implements IMatrixClientPeg {
 
     /**
      * Attempt to initialize the crypto layer on a newly-created MatrixClient
+     *
+     * @param rustCryptoStoreKey - If we are using Rust crypto, a key with which to encrypt the indexeddb.
+     *   If provided, it must be exactly 32 bytes of data. If both this and `rustCryptoStorePassword` are
+     *   undefined, the store will be unencrypted.
+     *
+     * @param rustCryptoStorePassword - An alternative to `rustCryptoStoreKey`. Ignored if `rustCryptoStoreKey` is set.
+     *    A password which will be used to derive a key to encrypt the store with. Deriving a key from a password is
+     *    (deliberately) a slow operation, so prefer to pass a `rustCryptoStoreKey` directly where possible.
      */
-    private async initClientCrypto(): Promise<void> {
+    private async initClientCrypto(rustCryptoStoreKey?: Uint8Array, rustCryptoStorePassword?: string): Promise<void> {
         if (!this.matrixClient) {
             throw new Error("createClient must be called first");
         }
@@ -347,7 +379,13 @@ class MatrixClientPegClass implements IMatrixClientPeg {
 
         // Now we can initialise the right crypto impl.
         if (useRustCrypto) {
-            await this.matrixClient.initRustCrypto();
+            if (!rustCryptoStoreKey && !rustCryptoStorePassword) {
+                logger.error("Warning! Not using an encryption key for rust crypto store.");
+            }
+            await this.matrixClient.initRustCrypto({
+                storageKey: rustCryptoStoreKey,
+                storagePassword: rustCryptoStorePassword,
+            });
 
             StorageManager.setCryptoInitialised(true);
             // TODO: device dehydration and whathaveyou
@@ -376,8 +414,11 @@ class MatrixClientPegClass implements IMatrixClientPeg {
         }
     }
 
-    public async start(): Promise<void> {
-        const opts = await this.assign();
+    /**
+     * Implementation of {@link IMatrixClientPeg.start}.
+     */
+    public async start(assignOpts?: MatrixClientPegAssignOpts): Promise<void> {
+        const opts = await this.assign(assignOpts);
 
         logger.log(`MatrixClientPeg: really starting MatrixClient`);
         await this.matrixClient!.startClient(opts);

test/MatrixClientPeg-test.ts

@@ -217,7 +217,7 @@ describe("MatrixClientPeg", () => {
                 testPeg.safeGet().store.on = emitter.on.bind(emitter);
                 const platform: any = { reload: jest.fn() };
                 PlatformPeg.set(platform);
-                await testPeg.assign();
+                await testPeg.assign({});
                 emitter.emit("closed" as any);
                 expect(platform.reload).toHaveBeenCalled();
             });
@@ -229,7 +229,7 @@ describe("MatrixClientPeg", () => {
                 PlatformPeg.set(platform);
                 testPeg.safeGet().store.on = emitter.on.bind(emitter);
                 const spy = jest.spyOn(Modal, "createDialog");
-                await testPeg.assign();
+                await testPeg.assign({});
                 emitter.emit("closed" as any);
                 expect(spy).toHaveBeenCalled();
             });
@@ -243,9 +243,10 @@ describe("MatrixClientPeg", () => {
             const mockInitCrypto = jest.spyOn(testPeg.safeGet(), "initCrypto").mockResolvedValue(undefined);
             const mockInitRustCrypto = jest.spyOn(testPeg.safeGet(), "initRustCrypto").mockResolvedValue(undefined);
 
-            await testPeg.start();
+            const cryptoStoreKey = new Uint8Array([1, 2, 3, 4]);
+            await testPeg.start({ rustCryptoStoreKey: cryptoStoreKey });
             expect(mockInitCrypto).not.toHaveBeenCalled();
-            expect(mockInitRustCrypto).toHaveBeenCalledTimes(1);
+            expect(mockInitRustCrypto).toHaveBeenCalledWith({ storageKey: cryptoStoreKey });
 
             // we should have stashed the setting in the settings store
             expect(mockSetValue).toHaveBeenCalledWith("feature_rust_crypto", null, SettingLevel.DEVICE, true);
@@ -271,7 +272,7 @@ describe("MatrixClientPeg", () => {
 
             await testPeg.start();
             expect(mockInitCrypto).toHaveBeenCalled();
-            expect(mockInitRustCrypto).not.toHaveBeenCalledTimes(1);
+            expect(mockInitRustCrypto).not.toHaveBeenCalled();
 
             // we should have stashed the setting in the settings store
             expect(mockSetValue).toHaveBeenCalledWith("feature_rust_crypto", null, SettingLevel.DEVICE, false);