Add support for device dehydration v2 (Element R)

package.json

@@ -52,7 +52,7 @@
     ],
     "dependencies": {
         "@babel/runtime": "^7.12.5",
-        "@matrix-org/matrix-sdk-crypto-wasm": "^4.6.0",
+        "@matrix-org/matrix-sdk-crypto-wasm": "^4.8.0",
         "another-json": "^0.2.0",
         "bs58": "^5.0.0",
         "content-type": "^1.0.4",

spec/unit/rust-crypto/rust-crypto.spec.ts

@@ -762,8 +762,11 @@ describe("RustCrypto", () => {
                             },
                         },
                     };
-                } else if (request instanceof RustSdkCryptoJs.UploadSigningKeysRequest) {
-                    // SigningKeysUploadRequest does not implement OutgoingRequest and does not need to be marked as sent.
+                } else if (
+                    request instanceof RustSdkCryptoJs.UploadSigningKeysRequest ||
+                    request instanceof RustSdkCryptoJs.PutDehydratedDeviceRequest
+                ) {
+                    // These request types do not implement OutgoingRequest and do not need to be marked as sent.
                     return;
                 }
                 if (request.id) {
@@ -1395,6 +1398,261 @@ describe("RustCrypto", () => {
             });
         });
     });
+
+    describe("dehydration", () => {
+        it("should rehydrate and dehydrate a device", async () => {
+            const secretStorageCallbacks = {
+                getSecretStorageKey: async (keys: any, name: string) => {
+                    return [[...Object.keys(keys.keys)][0], new Uint8Array(32)];
+                },
+            } as SecretStorageCallbacks;
+            const secretStorage = new ServerSideSecretStorageImpl(new DummyAccountDataClient(), secretStorageCallbacks);
+
+            const rustCrypto = await makeTestRustCrypto(
+                makeMatrixHttpApi(),
+                testData.TEST_USER_ID,
+                undefined,
+                secretStorage,
+            );
+            const olmMachine: OlmMachine = rustCrypto["olmMachine"];
+            rustCrypto["checkKeyBackupAndEnable"] = async () => {
+                return null;
+            };
+
+            const outgoingRequestProcessor = {
+                makeOutgoingRequest: jest.fn(),
+            } as unknown as OutgoingRequestProcessor;
+            rustCrypto["outgoingRequestProcessor"] = outgoingRequestProcessor;
+            (rustCrypto["crossSigningIdentity"] as any)["outgoingRequestProcessor"] = outgoingRequestProcessor;
+            const outgoingRequestsManager = new OutgoingRequestsManager(logger, olmMachine, outgoingRequestProcessor);
+            rustCrypto["outgoingRequestsManager"] = outgoingRequestsManager;
+
+            async function createSecretStorageKey() {
+                return {
+                    keyInfo: {} as AddSecretStorageKeyOpts,
+                    privateKey: new Uint8Array(32),
+                };
+            }
+
+            // create initial secret storage
+            await rustCrypto.bootstrapCrossSigning({ setupNewCrossSigning: true });
+            await rustCrypto.bootstrapSecretStorage({
+                createSecretStorageKey,
+                setupNewSecretStorage: true,
+                setupNewKeyBackup: false,
+            });
+
+            // there isn't any dehydrated device yet
+            fetchMock.config.overwriteRoutes = true;
+            fetchMock.get("path:/_matrix/client/unstable/org.matrix.msc3814.v1/dehydrated_device", {
+                status: 404,
+                body: {
+                    errcode: "M_NOT_FOUND",
+                    error: "Not found",
+                },
+            });
+            expect(await rustCrypto.rehydrateDeviceIfAvailable()).toBe(false);
+
+            // create a dehydrated device
+            let dehydratedDeviceBody: any;
+            fetchMock.put("path:/_matrix/client/unstable/org.matrix.msc3814.v1/dehydrated_device", (_, opts) => {
+                dehydratedDeviceBody = JSON.parse(opts.body as string);
+                return {};
+            });
+            await rustCrypto.createAndUploadDehydratedDevice();
+
+            // rehydrate the device that we just created
+            fetchMock.get("path:/_matrix/client/unstable/org.matrix.msc3814.v1/dehydrated_device", {
+                device_id: dehydratedDeviceBody.device_id,
+                device_data: dehydratedDeviceBody.device_data,
+            });
+            const eventsResponse = jest.fn((url, opts) => {
+                // rehydrating should make two calls to the /events endpoint.
+                // The first time will return a single event, and the second
+                // time will return no events (which will signal to the
+                // rehydration function that it can stop)
+                const body = JSON.parse(opts.body as string);
+                const nextBatch = body.next_batch ?? "0";
+                const events =
+                    nextBatch === "0" ? [{ sender: testData.TEST_USER_ID, type: "m.dummy", content: {} }] : [];
+                return {
+                    events,
+                    next_batch: nextBatch + "1",
+                };
+            });
+            fetchMock.post(
+                `path:/_matrix/client/unstable/org.matrix.msc3814.v1/dehydrated_device/${encodeURIComponent(dehydratedDeviceBody.device_id)}/events`,
+                eventsResponse,
+            );
+
+            expect(await rustCrypto.rehydrateDeviceIfAvailable()).toBe(true);
+            expect(eventsResponse.mock.calls).toHaveLength(2);
+        });
+
+        it("should schedule regular creation of dehydrated devices", async () => {
+            jest.useFakeTimers({ doNotFake: ["queueMicrotask"] });
+
+            const secretStorageCallbacks = {
+                getSecretStorageKey: async (keys: any, name: string) => {
+                    return [[...Object.keys(keys.keys)][0], new Uint8Array(32)];
+                },
+            } as SecretStorageCallbacks;
+            const secretStorage = new ServerSideSecretStorageImpl(new DummyAccountDataClient(), secretStorageCallbacks);
+
+            const rustCrypto = await makeTestRustCrypto(
+                makeMatrixHttpApi(),
+                testData.TEST_USER_ID,
+                undefined,
+                secretStorage,
+            );
+            const olmMachine: OlmMachine = rustCrypto["olmMachine"];
+            rustCrypto["checkKeyBackupAndEnable"] = async () => {
+                return null;
+            };
+
+            const outgoingRequestProcessor = {
+                makeOutgoingRequest: jest.fn(),
+            } as unknown as OutgoingRequestProcessor;
+            rustCrypto["outgoingRequestProcessor"] = outgoingRequestProcessor;
+            (rustCrypto["crossSigningIdentity"] as any)["outgoingRequestProcessor"] = outgoingRequestProcessor;
+            (rustCrypto["dehydrationManager"] as any)["outgoingRequestProcessor"] = outgoingRequestProcessor;
+            const outgoingRequestsManager = new OutgoingRequestsManager(logger, olmMachine, outgoingRequestProcessor);
+            rustCrypto["outgoingRequestsManager"] = outgoingRequestsManager;
+
+            async function createSecretStorageKey() {
+                return {
+                    keyInfo: {} as AddSecretStorageKeyOpts,
+                    privateKey: new Uint8Array(32),
+                };
+            }
+
+            // create initial secret storage
+            await rustCrypto.bootstrapCrossSigning({ setupNewCrossSigning: true });
+            await rustCrypto.bootstrapSecretStorage({
+                createSecretStorageKey,
+                setupNewSecretStorage: true,
+                setupNewKeyBackup: false,
+            });
+
+            // when we schedule dehydration with no delay, it should create a
+            // dehydrated device immediately
+            outgoingRequestProcessor.makeOutgoingRequest = jest.fn(async (req, uiaCallback) => {
+                expect(req).toBeInstanceOf(RustSdkCryptoJs.PutDehydratedDeviceRequest);
+            });
+            await rustCrypto.scheduleDeviceDehydration(30000);
+
+            expect(outgoingRequestProcessor.makeOutgoingRequest).toHaveBeenCalled();
+
+            // after we advance the timer, it should create another dehydrated device
+            // we make this a promise so that we can await it to make sure it gets
+            // called
+            const secondDehydrationPromise = new Promise<void>((resolve, reject) => {
+                outgoingRequestProcessor.makeOutgoingRequest = jest.fn(async (req, uiaCallback) => {
+                    expect(req).toBeInstanceOf(RustSdkCryptoJs.PutDehydratedDeviceRequest);
+                    resolve();
+                });
+            });
+            jest.advanceTimersByTime(35000);
+
+            await secondDehydrationPromise;
+
+            // when we schedule dehydration with a delay, it should not create
+            // a dehydrated device immediately
+            const thirdDehydrationPromise = new Promise<void>((resolve, reject) => {
+                outgoingRequestProcessor.makeOutgoingRequest = jest.fn(async (req, uiaCallback) => {
+                    expect(req).toBeInstanceOf(RustSdkCryptoJs.PutDehydratedDeviceRequest);
+                    resolve();
+                });
+            });
+            await rustCrypto.scheduleDeviceDehydration(30000, 10000);
+            expect(outgoingRequestProcessor.makeOutgoingRequest).not.toHaveBeenCalled();
+            jest.advanceTimersByTime(15000);
+            await thirdDehydrationPromise;
+
+            // when we stop rustCrypto, any pending device dehydration tasks
+            // should be cancelled
+            outgoingRequestProcessor.makeOutgoingRequest = jest.fn(async (req, uiaCallback) => {
+                expect(req).toBeInstanceOf(RustSdkCryptoJs.PutDehydratedDeviceRequest);
+            });
+            await rustCrypto.scheduleDeviceDehydration(30000, 10000);
+            rustCrypto.stop();
+            jest.advanceTimersByTime(15000);
+            expect(outgoingRequestProcessor.makeOutgoingRequest).not.toHaveBeenCalled();
+        });
+
+        it("should detect if dehydration is supported", async () => {
+            const rustCrypto = await makeTestRustCrypto(makeMatrixHttpApi());
+            fetchMock.config.overwriteRoutes = true;
+            fetchMock.get("path:/_matrix/client/unstable/org.matrix.msc3814.v1/dehydrated_device", {
+                status: 404,
+                body: {
+                    errcode: "M_UNRECOGNIZED",
+                    error: "Unknown endpoint",
+                },
+            });
+            expect(await rustCrypto.isDehydrationSupported()).toBe(false);
+            fetchMock.get("path:/_matrix/client/unstable/org.matrix.msc3814.v1/dehydrated_device", {
+                status: 404,
+                body: {
+                    errcode: "M_NOT_FOUND",
+                    error: "Unknown endpoint",
+                },
+            });
+            expect(await rustCrypto.isDehydrationSupported()).toBe(true);
+            fetchMock.get("path:/_matrix/client/unstable/org.matrix.msc3814.v1/dehydrated_device", {
+                device_id: "DEVICE_ID",
+                device_data: "data",
+            });
+            expect(await rustCrypto.isDehydrationSupported()).toBe(true);
+        });
+
+        it("should detect if dehydration key is stored", async () => {
+            const secretStorageCallbacks = {
+                getSecretStorageKey: async (keys: any, name: string) => {
+                    return [[...Object.keys(keys.keys)][0], new Uint8Array(32)];
+                },
+            } as SecretStorageCallbacks;
+            const secretStorage = new ServerSideSecretStorageImpl(new DummyAccountDataClient(), secretStorageCallbacks);
+
+            const rustCrypto = await makeTestRustCrypto(
+                makeMatrixHttpApi(),
+                testData.TEST_USER_ID,
+                undefined,
+                secretStorage,
+            );
+            async function createSecretStorageKey() {
+                return {
+                    keyInfo: {} as AddSecretStorageKeyOpts,
+                    privateKey: new Uint8Array(32),
+                };
+            }
+            await rustCrypto.bootstrapSecretStorage({
+                createSecretStorageKey,
+                setupNewSecretStorage: true,
+                setupNewKeyBackup: false,
+            });
+
+            expect(await rustCrypto.isDehydrationKeyStored()).toBe(false);
+
+            await rustCrypto.resetDehydrationKey();
+            expect(await rustCrypto.isDehydrationKeyStored()).toBe(true);
+
+            // rehydrateDeviceIfAvailable should check the server to see if a
+            // dehydrated device is present, because the dehydration key is set
+            const response = jest.fn(() => {
+                return {
+                    status: 404,
+                    body: {
+                        errcode: "M_NOT_FOUND",
+                        error: "Not found",
+                    },
+                };
+            });
+            fetchMock.get("path:/_matrix/client/unstable/org.matrix.msc3814.v1/dehydrated_device", response);
+            expect(await rustCrypto.rehydrateDeviceIfAvailable()).toBe(false);
+            expect(response).toHaveBeenCalled();
+        });
+    });
 });
 
 /** Build a MatrixHttpApi instance */

src/crypto-api.ts

@@ -496,6 +496,57 @@ export interface CryptoApi {
      * @param version - The backup version to delete.
      */
     deleteKeyBackupVersion(version: string): Promise<void>;
+
+    ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+    //
+    // Dehydrated devices
+    //
+    ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+
+    /**
+     * Returns whether MSC3814 dehydrated devices are supported by the crypto
+     * backend and by the server.
+     */
+    isDehydrationSupported(): Promise<boolean>;
+
+    /**
+     * Rehydrate the dehydrated device stored on the server
+     *
+     * Checks if there is a dehydrated device on the server.  If so, rehydrates
+     * the device and processes the to-device events.
+     *
+     * Returns whether or not a dehydrated device was found.
+     */
+    rehydrateDeviceIfAvailable(): Promise<boolean>;
+
+    /**
+     * Creates and uploads a new dehydrated device
+     *
+     * Creates and stores a new key in secret storage if none is available.
+     */
+    createAndUploadDehydratedDevice(): Promise<void>;
+
+    /**
+     * Schedule periodic creation of dehydrated devices
+     *
+     * @param interval - the time to wait between creating dehydrated devices
+     * @param delay - how long to wait before creating the first dehydrated device.
+     *     Defaults to creating the device immediately.
+     */
+    scheduleDeviceDehydration(interval: number, delay?: number): Promise<void>;
+
+    /** Return whether the dehydration key is stored in Secret Storage
+     */
+    isDehydrationKeyStored(): Promise<boolean>;
+
+    /**
+     * Reset the dehydration key
+     *
+     * Note: this does not create a new dehydrated device.  This will need to
+     * be done either by calling `createAndUploadDehydratedDevice` or
+     * `scheduleDeviceDehydration`.
+     */
+    resetDehydrationKey(): Promise<void>;
 }
 
 /** A reason code for a failure to decrypt an event. */

src/crypto/index.ts

@@ -4287,6 +4287,49 @@ export class Crypto extends TypedEventEmitter<CryptoEvent, CryptoEventHandlerMap
     public getRoomEncryption(roomId: string): IRoomEncryption | null {
         return this.roomList.getRoomEncryption(roomId);
     }
+
+    /**
+     * Returns whether dehydrated devices are supported by the crypto backend
+     * and by the server.
+     */
+    public async isDehydrationSupported(): Promise<boolean> {
+        return false;
+    }
+
+    /**
+     * Stub function -- dehydration is not implemented here, so always return false
+     */
+    public async rehydrateDeviceIfAvailable(): Promise<boolean> {
+        return false;
+    }
+
+    /**
+     * Stub function -- dehydration is not implemented here, so throw error
+     */
+    public async createAndUploadDehydratedDevice(): Promise<void> {
+        throw new Error("Not implemented");
+    }
+
+    /**
+     * Stub function -- dehydration is not implemented here, so throw error
+     */
+    public async scheduleDeviceDehydration(interval: number, delay?: number): Promise<void> {
+        throw new Error("Not implemented");
+    }
+
+    /**
+     * Stub function -- dehydration is not implemented here, so throw error
+     */
+    public async isDehydrationKeyStored(): Promise<boolean> {
+        return false;
+    }
+
+    /**
+     * Stub function -- dehydration is not implemented here, so throw error
+     */
+    public async resetDehydrationKey(): Promise<void> {
+        throw new Error("Not implemented");
+    }
 }
 
 /**