Device Dehydration | js-sdk: store/load dehydration key

spec/integ/crypto/device-dehydration.spec.ts

@@ -17,11 +17,12 @@ limitations under the License.
 import "fake-indexeddb/auto";
 import fetchMock from "fetch-mock-jest";
 
-import { createClient, ClientEvent, MatrixClient, MatrixEvent } from "../../../src";
+import { ClientEvent, createClient, MatrixClient, MatrixEvent } from "../../../src";
 import { RustCrypto } from "../../../src/rust-crypto/rust-crypto";
 import { AddSecretStorageKeyOpts } from "../../../src/secret-storage";
 import { E2EKeyReceiver } from "../../test-utils/E2EKeyReceiver";
 import { E2EKeyResponder } from "../../test-utils/E2EKeyResponder";
+import { DehydratedDevicesEvents } from "../../../src/crypto-api";
 
 describe("Device dehydration", () => {
     it("should rehydrate and dehydrate a device", async () => {
@@ -40,6 +41,29 @@ describe("Device dehydration", () => {
 
         await initializeSecretStorage(matrixClient, "@alice:localhost", "http://test.server");
 
+        const dehydratedDevices = matrixClient.getCrypto()!.dehydratedDevices();
+        let creationEventCount = 0;
+        let pickleKeyCachedEventCount = 0;
+        let rehydrationStartedCount = 0;
+        let rehydrationEndedCount = 0;
+        let rehydrationProgressEvent = 0;
+
+        dehydratedDevices.on(DehydratedDevicesEvents.DeviceCreated, () => {
+            creationEventCount++;
+        });
+        dehydratedDevices.on(DehydratedDevicesEvents.PickleKeyCached, () => {
+            pickleKeyCachedEventCount++;
+        });
+        dehydratedDevices.on(DehydratedDevicesEvents.RehydrationStarted, () => {
+            rehydrationStartedCount++;
+        });
+        dehydratedDevices.on(DehydratedDevicesEvents.RehydrationEnded, () => {
+            rehydrationEndedCount++;
+        });
+        dehydratedDevices.on(DehydratedDevicesEvents.RehydrationProgress, (roomKeyCount, toDeviceCount) => {
+            rehydrationProgressEvent++;
+        });
+
         // count the number of times the dehydration key gets set
         let setDehydrationCount = 0;
         matrixClient.on(ClientEvent.AccountData, (event: MatrixEvent) => {
@@ -74,14 +98,19 @@ describe("Device dehydration", () => {
         await crypto.startDehydration();
 
         expect(dehydrationCount).toEqual(1);
+        expect(creationEventCount).toEqual(1);
+        expect(pickleKeyCachedEventCount).toEqual(1);
 
         // a week later, we should have created another dehydrated device
         const dehydrationPromise = new Promise<void>((resolve, reject) => {
             resolveDehydrationPromise = resolve;
         });
         jest.advanceTimersByTime(7 * 24 * 60 * 60 * 1000);
         await dehydrationPromise;
+
+        expect(pickleKeyCachedEventCount).toEqual(1);
         expect(dehydrationCount).toEqual(2);
+        expect(creationEventCount).toEqual(2);
 
         // restart dehydration -- rehydrate the device that we created above,
         // and create a new dehydrated device.  We also set `createNewKey`, so
@@ -113,6 +142,12 @@ describe("Device dehydration", () => {
         expect(setDehydrationCount).toEqual(2);
         expect(eventsResponse.mock.calls).toHaveLength(2);
 
+        expect(rehydrationStartedCount).toEqual(1);
+        expect(rehydrationEndedCount).toEqual(1);
+        expect(creationEventCount).toEqual(3);
+        expect(rehydrationProgressEvent).toEqual(1);
+        expect(pickleKeyCachedEventCount).toEqual(2);
+
         matrixClient.stopClient();
     });
 });

src/crypto-api/index.ts

@@ -31,6 +31,7 @@ import {
 } from "./keybackup.ts";
 import { ISignatures } from "../@types/signed.ts";
 import { MatrixEvent } from "../models/event.ts";
+import { TypedEventEmitter } from "../models/typed-event-emitter.ts";
 
 /**
  * `matrix-js-sdk/lib/crypto-api`: End-to-end encryption support.
@@ -615,12 +616,18 @@ export interface CryptoApi {
     //
     ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
+    /**
+     * Get the API object for interacting with dehydrated devices.
+     */
+    dehydratedDevices(): DehydratedDevicesAPI;
+
     /**
      * Returns whether MSC3814 dehydrated devices are supported by the crypto
      * backend and by the server.
      *
      * This should be called before calling `startDehydration`, and if this
      * returns `false`, `startDehydration` should not be called.
+     *
      */
     isDehydrationSupported(): Promise<boolean>;
 
@@ -1243,6 +1250,45 @@ export interface OwnDeviceKeys {
     curve25519: string;
 }
 
+export enum DehydratedDevicesEvents {
+    /** Emitted when a new dehydrated device is created locally */
+    DeviceCreated = "DeviceCreated",
+    /** Emitted when a new dehydrated device is successfully uploaded to the server */
+    DeviceUploaded = "DeviceUploaded",
+    /** Emitted when rehydration has started */
+    RehydrationStarted = "RehydrationStarted",
+    /** Emitted when rehydration has finished */
+    RehydrationEnded = "RehydrationEnded",
+    /** Emitted during rehydration, signalling the current `roomKeyCount` and `toDeviceCount` */
+    RehydrationProgress = "RehydrationProgress",
+    /** Emitted when a dehydrated device key has been cached */
+    PickleKeyCached = "PickleKeyCached",
+    /** Emitted when an error occurred during rotation of the dehydrated device */
+    SchedulingError = "SchedulingError",
+}
+
+export type DehydratedDevicesEventsMap = {
+    [DehydratedDevicesEvents.DeviceCreated]: () => void;
+    [DehydratedDevicesEvents.DeviceUploaded]: () => void;
+    [DehydratedDevicesEvents.RehydrationStarted]: () => void;
+    [DehydratedDevicesEvents.RehydrationEnded]: () => void;
+    [DehydratedDevicesEvents.RehydrationProgress]: (roomKeyCount: number, toDeviceCount: number) => void;
+    [DehydratedDevicesEvents.PickleKeyCached]: () => void;
+    [DehydratedDevicesEvents.SchedulingError]: (msg: string) => void;
+};
+
+export abstract class DehydratedDevicesAPI extends TypedEventEmitter<
+    DehydratedDevicesEvents,
+    DehydratedDevicesEventsMap
+> {
+    protected constructor() {
+        super();
+    }
+
+    public abstract isSupported(): Promise<boolean>;
+    public abstract start(createNewKey?: boolean): Promise<void>;
+}
+
 export * from "./verification.ts";
 export type * from "./keybackup.ts";
 export * from "./recovery-key.ts";

src/crypto/index.ts

@@ -105,6 +105,7 @@ import {
     CryptoEventHandlerMap as CryptoApiCryptoEventHandlerMap,
     KeyBackupRestoreResult,
     KeyBackupRestoreOpts,
+    DehydratedDevicesAPI,
 } from "../crypto-api/index.ts";
 import { Device, DeviceMap } from "../models/device.ts";
 import { deviceInfoToDevice } from "./device-converter.ts";
@@ -4324,6 +4325,10 @@ export class Crypto extends TypedEventEmitter<CryptoEvent, CryptoEventHandlerMap
         throw new Error("Not implemented");
     }
 
+    public dehydratedDevices(): DehydratedDevicesAPI {
+        throw new Error("Not implemented");
+    }
+
     /**
      * Stub function -- restoreKeyBackup is not implemented here, so throw error
      */

src/rust-crypto/DehydratedDeviceManager.ts

@@ -21,8 +21,9 @@
 import { IHttpOpts, MatrixError, MatrixHttpApi, Method } from "../http-api/index.ts";
 import { IToDeviceEvent } from "../sync-accumulator.ts";
 import { ServerSideSecretStorage } from "../secret-storage.ts";
-import { decodeBase64, encodeUnpaddedBase64 } from "../base64.ts";
+import { decodeBase64 } from "../base64.ts";
 import { Logger } from "../logger.ts";
+import { DehydratedDevicesEvents, DehydratedDevicesAPI } from "../crypto-api/index.ts";
 
 /**
  * The response body of `GET /_matrix/client/unstable/org.matrix.msc3814.v1/dehydrated_device`.
@@ -67,9 +68,7 @@
  *
  * @internal
  */
-export class DehydratedDeviceManager {
-    /** the secret key used for dehydrating and rehydrating */
-    private key?: Uint8Array;
+export class DehydratedDeviceManager extends DehydratedDevicesAPI {
     /** the ID of the interval for periodically replacing the dehydrated device */
     private intervalId?: ReturnType<typeof setInterval>;
 
@@ -79,8 +78,18 @@
         private readonly http: MatrixHttpApi<IHttpOpts & { onlyData: true }>,
         private readonly outgoingRequestProcessor: OutgoingRequestProcessor,
         private readonly secretStorage: ServerSideSecretStorage,
-    ) {}
+    ) {
+        super();
+    }
+
+    private async getCachedKey(): Promise<RustSdkCryptoJs.DehydratedDeviceKey | undefined> {
+        return await this.olmMachine.dehydratedDevices().getDehydratedDeviceKey();
+    }
 
+    private async cacheKey(key: RustSdkCryptoJs.DehydratedDeviceKey): Promise<void> {
+        await this.olmMachine.dehydratedDevices().saveDehydratedDeviceKey(key);
+        this.emit(DehydratedDevicesEvents.PickleKeyCached);
+    }
     /**
      * Return whether the server supports dehydrated devices.
      */
@@ -153,10 +162,10 @@
      * Creates a new key and stores it in secret storage.
      */
     public async resetKey(): Promise<void> {
-        const key = new Uint8Array(32);
-        globalThis.crypto.getRandomValues(key);
-        await this.secretStorage.store(SECRET_STORAGE_NAME, encodeUnpaddedBase64(key));
-        this.key = key;
+        const key = RustSdkCryptoJs.DehydratedDeviceKey.createRandomKey();
+        await this.secretStorage.store(SECRET_STORAGE_NAME, key.toBase64());
+        // also cache it
+        await this.cacheKey(key);
     }
 
     /**
@@ -166,19 +175,22 @@
      *
      * @returns the key, if available, or `null` if no key is available
      */
-    private async getKey(create: boolean): Promise<Uint8Array | null> {
-        if (this.key === undefined) {
+    private async getKey(create: boolean): Promise<RustSdkCryptoJs.DehydratedDeviceKey | null> {
+        const cachedKey = await this.getCachedKey();
+        if (!cachedKey) {
             const keyB64 = await this.secretStorage.get(SECRET_STORAGE_NAME);
             if (keyB64 === undefined) {
                 if (!create) {
                     return null;
                 }
                 await this.resetKey();
             } else {
-                this.key = decodeBase64(keyB64);
+                const bytes = decodeBase64(keyB64);
+                const key = RustSdkCryptoJs.DehydratedDeviceKey.createKeyFromArray(bytes);
+                await this.cacheKey(key);
             }
         }
-        return this.key!;
+        return (await this.getCachedKey())!;
     }
 
     /**
@@ -190,7 +202,7 @@
      * Returns whether or not a dehydrated device was found.
      */
     public async rehydrateDeviceIfAvailable(): Promise<boolean> {
-        const key = await this.getKey(false);
+        const key = (await this.getCachedKey()) || (await this.getKey(false));
         if (!key) {
             return false;
         }
@@ -219,6 +231,7 @@
         }
 
         this.logger.info("dehydration: dehydrated device found");
+        this.emit(DehydratedDevicesEvents.RehydrationStarted);
 
         const rehydratedDevice = await this.olmMachine
             .dehydratedDevices()
@@ -255,8 +268,11 @@
             nextBatch = eventResp.next_batch;
             const roomKeyInfos = await rehydratedDevice.receiveEvents(JSON.stringify(eventResp.events));
             roomKeyCount += roomKeyInfos.length;
+
+            this.emit(DehydratedDevicesEvents.RehydrationProgress, roomKeyCount, toDeviceCount);
         }
         this.logger.info(`dehydration: received ${roomKeyCount} room keys from ${toDeviceCount} to-device events`);
+        this.emit(DehydratedDevicesEvents.RehydrationEnded);
 
         return true;
     }
@@ -267,12 +283,14 @@
      * Creates and stores a new key in secret storage if none is available.
      */
     public async createAndUploadDehydratedDevice(): Promise<void> {
-        const key = (await this.getKey(true))!;
+        const key = ((await this.getCachedKey()) || (await this.getKey(true)))!;
 
         const dehydratedDevice = await this.olmMachine.dehydratedDevices().create();
+        this.emit(DehydratedDevicesEvents.DeviceCreated);
         const request = await dehydratedDevice.keysForUpload("Dehydrated device", key);
 
         await this.outgoingRequestProcessor.makeOutgoingRequest(request);
+        this.emit(DehydratedDevicesEvents.DeviceUploaded);
 
         this.logger.info("dehydration: uploaded device");
     }
@@ -287,6 +305,7 @@
         await this.createAndUploadDehydratedDevice();
         this.intervalId = setInterval(() => {
             this.createAndUploadDehydratedDevice().catch((error) => {
+                this.emit(DehydratedDevicesEvents.SchedulingError, error.message);
                 this.logger.error("Error creating dehydrated device:", error);
             });
         }, DEHYDRATION_INTERVAL);

src/rust-crypto/rust-crypto.ts

@@ -67,6 +67,7 @@ import {
     CryptoEventHandlerMap,
     KeyBackupRestoreOpts,
     KeyBackupRestoreResult,
+    DehydratedDevicesAPI,
 } from "../crypto-api/index.ts";
 import { deviceKeysToDeviceMap, rustDeviceToJsDevice } from "./device-converter.ts";
 import { IDownloadKeyResult, IQueryKeysRequest } from "../client.ts";
@@ -1406,6 +1407,13 @@ export class RustCrypto extends TypedEventEmitter<RustCryptoEvents, CryptoEventH
         return await this.dehydratedDeviceManager.start(createNewKey);
     }
 
+    /**
+     * Implementation of {@link CryptoApi#dehydratedDevices}.
+     */
+    public dehydratedDevices(): DehydratedDevicesAPI {
+        return this.dehydratedDeviceManager;
+    }
+
     /**
      * Implementation of {@link CryptoApi#importSecretsBundle}.
      */