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

src/BasePlatform.ts

@@ -38,7 +38,7 @@ import { idbLoad, idbSave, idbDelete } from "./utils/StorageAccess";
 import { ViewRoomPayload } from "./dispatcher/payloads/ViewRoomPayload";
 import { IConfigOptions } from "./IConfigOptions";
 import SdkConfig from "./SdkConfig";
-import { buildAndEncodePickleKey, getPickleAdditionalData } from "./utils/tokens/pickling";
+import { buildAndEncodePickleKey, encryptPickleKey } from "./utils/tokens/pickling";
 
 export const SSO_HOMESERVER_URL_KEY = "mx_sso_hs_url";
 export const SSO_ID_SERVER_URL_KEY = "mx_sso_is_url";
@@ -378,24 +378,16 @@ export default abstract class BasePlatform {
      *     support storing pickle keys.
      */
     public async createPickleKey(userId: string, deviceId: string): Promise<string | null> {
-        if (!window.crypto || !window.crypto.subtle) {
-            return null;
-        }
-        const crypto = window.crypto;
         const randomArray = new Uint8Array(32);
         crypto.getRandomValues(randomArray);
-        const cryptoKey = await crypto.subtle.generateKey({ name: "AES-GCM", length: 256 }, false, [
-            "encrypt",
-            "decrypt",
-        ]);
-        const iv = new Uint8Array(32);
-        crypto.getRandomValues(iv);
-
-        const additionalData = getPickleAdditionalData(userId, deviceId);
-        const encrypted = await crypto.subtle.encrypt({ name: "AES-GCM", iv, additionalData }, cryptoKey, randomArray);
+        const data = await encryptPickleKey(randomArray, userId, deviceId);
+        if (data === undefined) {
+            // no crypto support
+            return null;
+        }
 
         try {
-            await idbSave("pickleKey", [userId, deviceId], { encrypted, iv, cryptoKey });
+            await idbSave("pickleKey", [userId, deviceId], data);
         } catch (e) {
             return null;
         }

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
@@ -94,14 +115,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
@@ -248,7 +269,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");
         }
@@ -275,7 +299,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);
@@ -301,8 +325,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");
         }
@@ -338,7 +370,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
@@ -367,8 +405,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);

src/utils/tokens/pickling.ts

@@ -20,6 +20,20 @@ limitations under the License.
 import { encodeUnpaddedBase64 } from "matrix-js-sdk/src/matrix";
 import { logger } from "matrix-js-sdk/src/logger";
 
+/**
+ * Encrypted format of a pickle key, as stored in IndexedDB.
+ */
+export interface EncryptedPickleKey {
+    /** The encrypted payload. */
+    encrypted?: BufferSource;
+
+    /** Initialisation vector for the encryption. */
+    iv?: BufferSource;
+
+    /** The encryption key which was used to encrypt the payload. */
+    cryptoKey?: CryptoKey;
+}
+
 /**
  * Calculates the `additionalData` for the AES-GCM key used by the pickling processes. This
  * additional data is *not* encrypted, but *is* authenticated. The additional data is constructed
@@ -46,6 +60,32 @@ export function getPickleAdditionalData(userId: string, deviceId: string): Uint8
     return additionalData;
 }
 
+/**
+ * Encrypt the given pickle key, ready for storage in the database.
+ *
+ * @param pickleKey - The key to be encrypted.
+ * @param userId - The user ID the pickle key belongs to.
+ * @param deviceId - The device ID the pickle key belongs to.
+ *
+ * @returns Data object ready for storing in indexeddb.
+ */
+export async function encryptPickleKey(
+    pickleKey: Uint8Array,
+    userId: string,
+    deviceId: string,
+): Promise<EncryptedPickleKey | undefined> {
+    if (!crypto?.subtle) {
+        return undefined;
+    }
+    const cryptoKey = await crypto.subtle.generateKey({ name: "AES-GCM", length: 256 }, false, ["encrypt", "decrypt"]);
+    const iv = new Uint8Array(32);
+    crypto.getRandomValues(iv);
+
+    const additionalData = getPickleAdditionalData(userId, deviceId);
+    const encrypted = await crypto.subtle.encrypt({ name: "AES-GCM", iv, additionalData }, cryptoKey, pickleKey);
+    return { encrypted, iv, cryptoKey };
+}
+
 /**
  * Decrypts the provided data into a pickle key and base64-encodes it ready for use elsewhere.
  *
@@ -59,7 +99,7 @@ export function getPickleAdditionalData(userId: string, deviceId: string): Uint8
  * @returns A promise that resolves to the encoded pickle key, or undefined if the key cannot be built and encoded.
  */
 export async function buildAndEncodePickleKey(
-    data: { encrypted?: BufferSource; iv?: BufferSource; cryptoKey?: CryptoKey } | undefined,
+    data: EncryptedPickleKey | undefined,
     userId: string,
     deviceId: string,
 ): Promise<string | undefined> {