feat(dehydrated): Add signalling to device dehydration manager

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 { encodeUri } from "../utils.ts";
 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,7 +68,7 @@ const DEHYDRATION_INTERVAL = 7 * 24 * 60 * 60 * 1000;
  *
  * @internal
  */
-export class DehydratedDeviceManager {
+export class DehydratedDeviceManager extends DehydratedDevicesAPI {
     /** the ID of the interval for periodically replacing the dehydrated device */
     private intervalId?: ReturnType<typeof setInterval>;
 
@@ -77,7 +78,9 @@ export class DehydratedDeviceManager {
         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();
@@ -228,6 +231,7 @@ export class DehydratedDeviceManager {
         }
 
         this.logger.info("dehydration: dehydrated device found");
+        this.emit(DehydratedDevicesEvents.RehydrationStarted);
 
         const rehydratedDevice = await this.olmMachine
             .dehydratedDevices()
@@ -264,8 +268,11 @@ export class DehydratedDeviceManager {
             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;
     }
@@ -279,9 +286,11 @@ export class DehydratedDeviceManager {
         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");
     }
@@ -296,6 +305,7 @@ export class DehydratedDeviceManager {
         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}.
      */